diff --git a/haystack/components/rankers/transformers_similarity.py b/haystack/components/rankers/transformers_similarity.py index 033c06aa57..05cdcea8f1 100644 --- a/haystack/components/rankers/transformers_similarity.py +++ b/haystack/components/rankers/transformers_similarity.py @@ -37,7 +37,7 @@ class TransformersSimilarityRanker: def __init__( self, - model_name_or_path: Union[str, Path] = "cross-encoder/ms-marco-MiniLM-L-6-v2", + model: Union[str, Path] = "cross-encoder/ms-marco-MiniLM-L-6-v2", device: Optional[str] = "cpu", token: Union[bool, str, None] = None, top_k: int = 10, @@ -51,7 +51,7 @@ def __init__( """ Creates an instance of TransformersSimilarityRanker. - :param model_name_or_path: The name or path of a pre-trained cross-encoder model + :param model: The name or path of a pre-trained cross-encoder model from the Hugging Face Hub. :param device: The torch device (for example, cuda:0, cpu, mps) to which you want to limit model inference. :param token: The API token used to download private models from Hugging Face. @@ -66,18 +66,18 @@ def __init__( `sigmoid(logits * calibration_factor)`. This is only used if `scale_score` is set to True. :param score_threshold: If provided only returns documents with a score above this threshold. :param model_kwargs: Additional keyword arguments passed to `AutoModelForSequenceClassification.from_pretrained` - when loading the model specified in `model_name_or_path`. For details on what kwargs you can pass, + when loading the model specified in `model`. For details on what kwargs you can pass, see the model's documentation. """ torch_and_transformers_import.check() - self.model_name_or_path = model_name_or_path + self.model = model if top_k <= 0: raise ValueError(f"top_k must be > 0, but got {top_k}") self.top_k = top_k self.device = device self.token = token - self.model = None + self._model = None self.tokenizer = None self.meta_fields_to_embed = meta_fields_to_embed or [] self.embedding_separator = embedding_separator @@ -94,20 +94,20 @@ def _get_telemetry_data(self) -> Dict[str, Any]: """ Data that is sent to Posthog for usage analytics. """ - return {"model": str(self.model_name_or_path)} + return {"model": str(self.model)} def warm_up(self): """ Warm up the model and tokenizer used for scoring the Documents. """ - if self.model is None: + if self._model is None: if self.device is None: self.device = get_device() - self.model = AutoModelForSequenceClassification.from_pretrained( - self.model_name_or_path, token=self.token, **self.model_kwargs + self._model = AutoModelForSequenceClassification.from_pretrained( + self.model, token=self.token, **self.model_kwargs ).to(self.device) - self.model.eval() - self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, token=self.token) + self._model.eval() + self.tokenizer = AutoTokenizer.from_pretrained(self.model, token=self.token) def to_dict(self) -> Dict[str, Any]: """ @@ -116,7 +116,7 @@ def to_dict(self) -> Dict[str, Any]: serialization_dict = default_to_dict( self, device=self.device, - model_name_or_path=self.model_name_or_path, + model=self.model, token=self.token if not isinstance(self.token, str) else None, # don't serialize valid tokens top_k=self.top_k, meta_fields_to_embed=self.meta_fields_to_embed, @@ -208,7 +208,7 @@ def run( score_threshold = self.score_threshold # If a model path is provided but the model isn't loaded - if self.model_name_or_path and not self.model: + if self.model and not self._model: raise ComponentError( f"The component {self.__class__.__name__} wasn't warmed up. Run 'warm_up()' before calling 'run()'." ) @@ -227,7 +227,7 @@ def run( self.device ) with torch.inference_mode(): - similarity_scores = self.model(**features).logits.squeeze(dim=1) # type: ignore + similarity_scores = self._model(**features).logits.squeeze(dim=1) # type: ignore if scale_score: similarity_scores = torch.sigmoid(similarity_scores * calibration_factor) diff --git a/releasenotes/notes/rename-model-param-rankers-be3d6163597e0a9b.yaml b/releasenotes/notes/rename-model-param-rankers-be3d6163597e0a9b.yaml new file mode 100644 index 0000000000..afc85b5fd9 --- /dev/null +++ b/releasenotes/notes/rename-model-param-rankers-be3d6163597e0a9b.yaml @@ -0,0 +1,3 @@ +--- +upgrade: + - Rename `model_name_or_path` to `model` in `TransformersSimilarityRanker`. diff --git a/test/components/rankers/test_transformers_similarity.py b/test/components/rankers/test_transformers_similarity.py index c81af7be53..1cc97176fd 100644 --- a/test/components/rankers/test_transformers_similarity.py +++ b/test/components/rankers/test_transformers_similarity.py @@ -17,7 +17,7 @@ def test_to_dict(self): "device": "cpu", "top_k": 10, "token": None, - "model_name_or_path": "cross-encoder/ms-marco-MiniLM-L-6-v2", + "model": "cross-encoder/ms-marco-MiniLM-L-6-v2", "meta_fields_to_embed": [], "embedding_separator": "\n", "scale_score": True, @@ -29,7 +29,7 @@ def test_to_dict(self): def test_to_dict_with_custom_init_parameters(self): component = TransformersSimilarityRanker( - model_name_or_path="my_model", + model="my_model", device="cuda", token="my_token", top_k=5, @@ -43,7 +43,7 @@ def test_to_dict_with_custom_init_parameters(self): "type": "haystack.components.rankers.transformers_similarity.TransformersSimilarityRanker", "init_parameters": { "device": "cuda", - "model_name_or_path": "my_model", + "model": "my_model", "token": None, # we don't serialize valid tokens, "top_k": 5, "meta_fields_to_embed": [], @@ -71,7 +71,7 @@ def test_to_dict_with_quantization_options(self): "device": "cpu", "top_k": 10, "token": None, - "model_name_or_path": "cross-encoder/ms-marco-MiniLM-L-6-v2", + "model": "cross-encoder/ms-marco-MiniLM-L-6-v2", "meta_fields_to_embed": [], "embedding_separator": "\n", "scale_score": True, @@ -91,7 +91,7 @@ def test_from_dict(self): "type": "haystack.components.rankers.transformers_similarity.TransformersSimilarityRanker", "init_parameters": { "device": "cuda", - "model_name_or_path": "my_model", + "model": "my_model", "token": None, "top_k": 5, "meta_fields_to_embed": [], @@ -105,7 +105,7 @@ def test_from_dict(self): component = TransformersSimilarityRanker.from_dict(data) assert component.device == "cuda" - assert component.model_name_or_path == "my_model" + assert component.model == "my_model" assert component.token is None assert component.top_k == 5 assert component.meta_fields_to_embed == [] @@ -122,9 +122,9 @@ def test_embed_meta(self, mocked_sort, mocked_sigmoid): mocked_sort.return_value = (None, torch.tensor([0])) mocked_sigmoid.return_value = torch.tensor([0]) embedder = TransformersSimilarityRanker( - model_name_or_path="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n" + model="model", meta_fields_to_embed=["meta_field"], embedding_separator="\n" ) - embedder.model = MagicMock() + embedder._model = MagicMock() embedder.tokenizer = MagicMock() documents = [Document(content=f"document number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)] @@ -147,9 +147,9 @@ def test_embed_meta(self, mocked_sort, mocked_sigmoid): @patch("torch.sort") def test_scale_score_false(self, mocked_sort): mocked_sort.return_value = (None, torch.tensor([0, 1])) - embedder = TransformersSimilarityRanker(model_name_or_path="model", scale_score=False) - embedder.model = MagicMock() - embedder.model.return_value = SequenceClassifierOutput( + embedder = TransformersSimilarityRanker(model="model", scale_score=False) + embedder._model = MagicMock() + embedder._model.return_value = SequenceClassifierOutput( loss=None, logits=torch.FloatTensor([[-10.6859], [-8.9874]]), hidden_states=None, attentions=None ) embedder.tokenizer = MagicMock() @@ -162,9 +162,9 @@ def test_scale_score_false(self, mocked_sort): @patch("torch.sort") def test_score_threshold(self, mocked_sort): mocked_sort.return_value = (None, torch.tensor([0, 1])) - embedder = TransformersSimilarityRanker(model_name_or_path="model", scale_score=False, score_threshold=0.1) - embedder.model = MagicMock() - embedder.model.return_value = SequenceClassifierOutput( + embedder = TransformersSimilarityRanker(model="model", scale_score=False, score_threshold=0.1) + embedder._model = MagicMock() + embedder._model.return_value = SequenceClassifierOutput( loss=None, logits=torch.FloatTensor([[0.955], [0.001]]), hidden_states=None, attentions=None ) embedder.tokenizer = MagicMock() @@ -201,7 +201,7 @@ def test_run(self, query, docs_before_texts, expected_first_text, scores): """ Test if the component ranks documents correctly. """ - ranker = TransformersSimilarityRanker(model_name_or_path="cross-encoder/ms-marco-MiniLM-L-6-v2") + ranker = TransformersSimilarityRanker(model="cross-encoder/ms-marco-MiniLM-L-6-v2") ranker.warm_up() docs_before = [Document(content=text) for text in docs_before_texts] output = ranker.run(query=query, documents=docs_before) @@ -244,7 +244,7 @@ def test_run_top_k(self, query, docs_before_texts, expected_first_text): """ Test if the component ranks documents correctly with a custom top_k. """ - ranker = TransformersSimilarityRanker(model_name_or_path="cross-encoder/ms-marco-MiniLM-L-6-v2", top_k=2) + ranker = TransformersSimilarityRanker(model="cross-encoder/ms-marco-MiniLM-L-6-v2", top_k=2) ranker.warm_up() docs_before = [Document(content=text) for text in docs_before_texts] output = ranker.run(query=query, documents=docs_before) @@ -261,7 +261,7 @@ def test_run_single_document(self): """ Test if the component runs with a single document. """ - ranker = TransformersSimilarityRanker(model_name_or_path="cross-encoder/ms-marco-MiniLM-L-6-v2", device=None) + ranker = TransformersSimilarityRanker(model="cross-encoder/ms-marco-MiniLM-L-6-v2", device=None) ranker.warm_up() docs_before = [Document(content="Berlin")] output = ranker.run(query="City in Germany", documents=docs_before)