Skip to content

Commit

Permalink
feat!: Rename model_name_or_path to model in `TransformersSimilar…
Browse files Browse the repository at this point in the history
…ityRanker` (#6734)

* rename model parameter in transformers ranker

* fix tests for transformers ranker

* reno

* reno

* typo
  • Loading branch information
ZanSara authored Jan 12, 2024
1 parent e0afe27 commit 0616197
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 31 deletions.
28 changes: 14 additions & 14 deletions haystack/components/rankers/transformers_similarity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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]:
"""
Expand All @@ -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,
Expand Down Expand Up @@ -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()'."
)
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
---
upgrade:
- Rename `model_name_or_path` to `model` in `TransformersSimilarityRanker`.
34 changes: 17 additions & 17 deletions test/components/rankers/test_transformers_similarity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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": [],
Expand Down Expand Up @@ -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,
Expand All @@ -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": [],
Expand All @@ -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 == []
Expand All @@ -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)]
Expand All @@ -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()
Expand All @@ -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()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down

0 comments on commit 0616197

Please sign in to comment.