feat: added arg similarity

Author: Darinochka

autointent/_embedder.py

@@ -50,6 +50,8 @@ class EmbedderDumpMetadata(TypedDict):
     """Maximum sequence length for the embedding model."""
     use_cache: bool
     """Whether to use embeddings caching."""
+    similarity_fn_name: str | None
+    """Name of the similarity function to use."""
 
 
 class Embedder:
@@ -73,7 +75,10 @@ def __init__(self, embedder_config: EmbedderConfig) -> None:
         self.config = embedder_config
 
         self.embedding_model = SentenceTransformer(
-            self.config.model_name, device=self.config.device, prompts=embedder_config.get_prompt_config()
+            self.config.model_name,
+            device=self.config.device,
+            prompts=embedder_config.get_prompt_config(),
+            similarity_fn_name=self.config.similarity_fn_name,
         )
 
         self._logger = logging.getLogger(__name__)
@@ -116,6 +121,7 @@ def dump(self, path: Path) -> None:
             batch_size=self.config.batch_size,
             max_length=self.config.tokenizer_config.max_length,
             use_cache=self.config.use_cache,
+            similarity_fn_name=self.config.similarity_fn_name,
         )
         path.mkdir(parents=True, exist_ok=True)
         with (path / self._metadata_dict_name).open("w") as file:
@@ -186,3 +192,17 @@ def embed(self, utterances: list[str], task_type: TaskTypeEnum | None = None) ->
             np.save(embeddings_path, embeddings)
 
         return embeddings
+
+    def similarity(
+        self, embeddings1: npt.NDArray[np.float32], embeddings2: npt.NDArray[np.float32]
+    ) -> npt.NDArray[np.float32]:
+        """Calculate similarity between two sets of embeddings.
+
+        Args:
+            embeddings1: First set of embeddings.
+            embeddings2: Second set of embeddings.
+
+        Returns:
+            A numpy array of similarities.
+        """
+        return self.embedding_model.similarity(embeddings1, embeddings2)

autointent/configs/_transformers.py

@@ -60,6 +60,9 @@ class EmbedderConfig(HFModelConfig):
     sts_prompt: str | None = Field(None, description="Prompt for finding most similar sentences.")
     query_prompt: str | None = Field(None, description="Prompt for query.")
     passage_prompt: str | None = Field(None, description="Prompt for passage.")
+    similarity_fn_name: str | None = Field(
+        "cosine", description="Name of the similarity function to use (cosine, dot, euclidean, manhattan)."
+    )
 
     def get_prompt_config(self) -> dict[str, str] | None:
         """Get the prompt config for the given prompt type.

autointent/modules/scoring/_description/description.py

@@ -6,7 +6,6 @@
 import scipy
 from numpy.typing import NDArray
 from pydantic import PositiveFloat
-from sklearn.metrics.pairwise import cosine_similarity
 
 from autointent import Context, Embedder, Ranker
 from autointent.configs import CrossEncoderConfig, EmbedderConfig, TaskTypeEnum
@@ -159,7 +158,7 @@ def predict(self, utterances: list[str]) -> NDArray[np.float64]:
         """
         if self._encoder_type == "bi":
             utterance_vectors = self._embedder.embed(utterances, TaskTypeEnum.sts)
-            similarities: NDArray[np.float64] = cosine_similarity(utterance_vectors, self._description_vectors)
+            similarities: NDArray[np.float64] = self._embedder.similarity(utterance_vectors, self._description_vectors)
         else:
             pairs = [(utterance, description) for utterance in utterances for description in self._description_texts]