feat: added crossencoder (#181)

* feat: added crossencoder

* refactor

* feat: added arg similarity

* Update optimizer_config.schema.json

* feat: added tests

* feat: added errors

* fix: scoring test

* fix: description vectors error

* fix: description vectors error

* fix: lint

* fix: test

* add node validators (#177)

* add node validators

* add comments

* Update optimizer_config.schema.json

* rename bert model

* lint

* fixes

* fix test

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: voorhs <[email protected]>

* fix: unit tests

* feat: added test for description

* feat: delete encoder_type from the class args

* feat: update assets

* feat: update assets

* fix: fixed test

* Update optimizer_config.schema.json

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Roman Solomatin <[email protected]>
Co-authored-by: voorhs <[email protected]>

Author: 4 people

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:
@@ -76,6 +78,7 @@ def __init__(self, embedder_config: EmbedderConfig) -> None:
             self.config.model_name,
             device=self.config.device,
             prompts=embedder_config.get_prompt_config(),
+            similarity_fn_name=self.config.similarity_fn_name,
             trust_remote_code=self.config.trust_remote_code,
         )
 
@@ -119,6 +122,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:
@@ -189,3 +193,18 @@ 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.
+        """
+        result = self.embedding_model.similarity(embeddings1, embeddings2)
+        return result.detach().cpu().numpy().astype(np.float32)

autointent/configs/_transformers.py

@@ -61,6 +61,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

@@ -1,15 +1,14 @@
 """DescriptionScorer class for scoring utterances based on intent descriptions."""
 
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import scipy
 from numpy.typing import NDArray
 from pydantic import PositiveFloat
-from sklearn.metrics.pairwise import cosine_similarity
 
-from autointent import Context, Embedder
-from autointent.configs import EmbedderConfig, TaskTypeEnum
+from autointent import Context, Embedder, Ranker
+from autointent.configs import CrossEncoderConfig, EmbedderConfig, TaskTypeEnum
 from autointent.context.optimization_info import ScorerArtifact
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL
@@ -19,29 +18,38 @@
 class DescriptionScorer(BaseScorer):
     """Scoring module that scores utterances based on similarity to intent descriptions.
 
-    DescriptionScorer embeds both the utterances and the intent descriptions, then computes
-    a similarity score between the two, using either cosine similarity and softmax.
+    DescriptionScorer can use either a bi-encoder or cross-encoder architecture:
+    - Bi-encoder: Embeds both utterances and descriptions separately, then computes cosine similarity
+    - Cross-encoder: Directly computes similarity between each utterance-description pair
 
     Args:
-        embedder_config: Config of the embedder model
+        embedder_config: Config of the embedder model (for bi-encoder mode)
+        cross_encoder_config: Config of the cross-encoder model (for cross-encoder mode)
+        encoder_type: Type of encoder to use, either "bi" or "cross"
         temperature: Temperature parameter for scaling logits, defaults to 1.0
     """
 
-    _embedder: Embedder
+    _embedder: Embedder | None = None
+    _cross_encoder: Ranker | None = None
     name = "description"
     _n_classes: int
     _multilabel: bool
-    _description_vectors: NDArray[Any]
+    _description_vectors: NDArray[Any] | None = None
+    _description_texts: list[str] | None = None
     supports_multiclass = True
     supports_multilabel = True
 
     def __init__(
         self,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
+        cross_encoder_config: CrossEncoderConfig | str | dict[str, Any] | None = None,
+        encoder_type: Literal["bi", "cross"] = "bi",
         temperature: PositiveFloat = 1.0,
     ) -> None:
         self.temperature = temperature
         self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
+        self.cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
+        self._encoder_type = encoder_type
 
         if self.temperature < 0 or not isinstance(self.temperature, float | int):
             msg = "`temperature` argument of `DescriptionScorer` must be a positive float"
@@ -51,35 +59,51 @@ def __init__(
     def from_context(
         cls,
         context: Context,
-        temperature: PositiveFloat,
+        temperature: PositiveFloat = 1.0,
         embedder_config: EmbedderConfig | str | None = None,
+        cross_encoder_config: CrossEncoderConfig | str | None = None,
+        encoder_type: Literal["bi", "cross"] = "bi",
     ) -> "DescriptionScorer":
         """Create a DescriptionScorer instance using a Context object.
 
         Args:
             context: Context containing configurations and utilities
             temperature: Temperature parameter for scaling logits
             embedder_config: Config of the embedder model. If None, the best embedder is used
+            cross_encoder_config: Config of the cross-encoder model. If None, the default config is used
+            encoder_type: Type of encoder to use, either "bi" or "cross"
 
         Returns:
             Initialized DescriptionScorer instance
         """
         if embedder_config is None:
             embedder_config = context.resolve_embedder()
+        if cross_encoder_config is None:
+            cross_encoder_config = context.resolve_ranker()
 
         return cls(
             temperature=temperature,
             embedder_config=embedder_config,
+            cross_encoder_config=cross_encoder_config,
+            encoder_type=encoder_type,
         )
 
     def get_embedder_config(self) -> dict[str, Any]:
-        """Get the name of the embedder.
+        """Get the configuration of the embedder.
 
         Returns:
-            Embedder name
+            Embedder configuration
         """
         return self.embedder_config.model_dump()
 
+    def get_cross_encoder_config(self) -> dict[str, Any]:
+        """Get the configuration of the cross-encoder.
+
+        Returns:
+            Cross-encoder configuration
+        """
+        return self.cross_encoder_config.model_dump()
+
     def fit(
         self,
         utterances: list[str],
@@ -96,8 +120,10 @@ def fit(
         Raises:
             ValueError: If descriptions contain None values or embeddings mismatch utterances
         """
-        if hasattr(self, "_embedder"):
+        if hasattr(self, "_embedder") and self._embedder is not None:
             self._embedder.clear_ram()
+        if hasattr(self, "_cross_encoder") and self._cross_encoder is not None:
+            self._cross_encoder.clear_ram()
 
         self._validate_task(labels)
 
@@ -108,10 +134,17 @@ def fit(
             )
             raise ValueError(error_text)
 
-        embedder = Embedder(self.embedder_config)
-
-        self._description_vectors = embedder.embed(descriptions, TaskTypeEnum.sts)
-        self._embedder = embedder
+        if self._encoder_type == "bi":
+            embedder = Embedder(self.embedder_config)
+            self._description_vectors = embedder.embed(descriptions, TaskTypeEnum.sts)
+            self._embedder = embedder
+            self._cross_encoder = None
+            self._description_texts = None
+        else:
+            self._cross_encoder = Ranker(self.cross_encoder_config)
+            self._description_texts = descriptions
+            self._embedder = None
+            self._description_vectors = None
 
     def predict(self, utterances: list[str]) -> NDArray[np.float64]:
         """Predict scores for utterances based on similarity to intent descriptions.
@@ -122,8 +155,32 @@ def predict(self, utterances: list[str]) -> NDArray[np.float64]:
         Returns:
             Array of probabilities for each utterance
         """
-        utterance_vectors = self._embedder.embed(utterances, TaskTypeEnum.sts)
-        similarities: NDArray[np.float64] = cosine_similarity(utterance_vectors, self._description_vectors)
+        if self._encoder_type == "bi":
+            if self._description_vectors is None:
+                error_text = "Description vectors are not initialized. Call fit() before predict()."
+                raise RuntimeError(error_text)
+
+            if self._embedder is None:
+                error_text = "Embedder is not initialized. Call fit() before predict()."
+                raise RuntimeError(error_text)
+
+            utterance_vectors = self._embedder.embed(utterances, TaskTypeEnum.sts)
+            similarities: NDArray[np.float64] = np.array(
+                self._embedder.similarity(utterance_vectors, self._description_vectors), dtype=np.float64
+            )
+        else:
+            if self._cross_encoder is None:
+                error_text = "Cross encoder is not initialized. Call fit() before predict()."
+                raise RuntimeError(error_text)
+
+            if self._description_texts is None:
+                error_text = "Description texts are not initialized. Call fit() before predict()."
+                raise RuntimeError(error_text)
+
+            pairs = [(utterance, description) for utterance in utterances for description in self._description_texts]
+
+            scores = self._cross_encoder.predict(pairs)
+            similarities = np.array(scores, dtype=np.float64).reshape(len(utterances), len(self._description_texts))
 
         if self._multilabel:
             probabilities = scipy.special.expit(similarities / self.temperature)
@@ -132,8 +189,11 @@ def predict(self, utterances: list[str]) -> NDArray[np.float64]:
         return probabilities  # type: ignore[no-any-return]
 
     def clear_cache(self) -> None:
-        """Clear cached data in memory used by the embedder."""
-        self._embedder.clear_ram()
+        """Clear cached data in memory used by the embedder or cross-encoder."""
+        if self._embedder is not None:
+            self._embedder.clear_ram()
+        if self._cross_encoder is not None:
+            self._cross_encoder.clear_ram()
 
     def get_train_data(self, context: Context) -> tuple[list[str], ListOfLabels, list[str]]:
         """Get training data from context.

docs/optimizer_config.schema.json

@@ -212,6 +212,19 @@
                     "description": "Prompt for passage.",
                     "title": "Passage Prompt"
                 },
+                "similarity_fn_name": {
+                    "anyOf": [
+                        {
+                            "type": "string"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": "cosine",
+                    "description": "Name of the similarity function to use (cosine, dot, euclidean, manhattan).",
+                    "title": "Similarity Fn Name"
+                },
                 "use_cache": {
                     "default": false,
                     "description": "Whether to use embeddings caching.",
@@ -389,6 +402,7 @@
                 "sts_prompt": null,
                 "query_prompt": null,
                 "passage_prompt": null,
+                "similarity_fn_name": "cosine",
                 "use_cache": false
             }
         },

tests/assets/configs/description.yaml

@@ -11,6 +11,11 @@
   search_space:
     - module_name: description
       temperature: [1.0, 0.5, 0.1, 0.05]
+      embedder_config:
+        - model_name: sentence-transformers/all-MiniLM-L6-v2
+      cross_encoder_config:
+        - model_name: cross-encoder/ms-marco-MiniLM-L-6-v2
+      encoder_type: [cross, bi]
 - node_type: decision
   target_metric: decision_accuracy
   search_space:

tests/callback/test_callback.py

@@ -146,6 +146,7 @@ def test_pipeline_callbacks(dataset):
                         "query_prompt": None,
                         "sts_prompt": None,
                         "use_cache": False,
+                        "similarity_fn_name": "cosine",
                         "trust_remote_code": False,
                     },
                     "k": 1,
@@ -181,6 +182,7 @@ def test_pipeline_callbacks(dataset):
                         "query_prompt": None,
                         "sts_prompt": None,
                         "use_cache": False,
+                        "similarity_fn_name": "cosine",
                         "trust_remote_code": False,
                     },
                     "k": 1,
@@ -216,6 +218,7 @@ def test_pipeline_callbacks(dataset):
                         "query_prompt": None,
                         "sts_prompt": None,
                         "use_cache": False,
+                        "similarity_fn_name": "cosine",
                         "trust_remote_code": False,
                     },
                 },

tests/modules/scoring/test_description.py

@@ -1,3 +1,5 @@
+import tempfile
+
 import numpy as np
 import pytest
 
@@ -45,3 +47,63 @@ def test_description_scorer(dataset, expected_prediction, multilabel):
     assert metadata is None
 
     scorer.clear_cache()
+
+
+@pytest.mark.parametrize(
+    ("expected_prediction", "multilabel"),
+    [
+        ([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]], True),
+        ([[0.2, 0.3, 0.2, 0.2], [0.2, 0.3, 0.2, 0.2]], False),
+    ],
+)
+def test_description_scorer_cross_encoder(dataset, expected_prediction, multilabel):
+    if multilabel:
+        dataset = dataset.to_multilabel()
+    data_handler = DataHandler(dataset)
+
+    scorer = DescriptionScorer(
+        cross_encoder_config="cross-encoder/ms-marco-MiniLM-L-6-v2", encoder_type="cross", temperature=0.3
+    )
+
+    scorer.fit(
+        data_handler.train_utterances(0),
+        data_handler.train_labels(0),
+        data_handler.intent_descriptions,
+    )
+    assert scorer._description_texts is not None
+    assert len(scorer._description_texts) == len(data_handler.intent_descriptions)
+    assert scorer._cross_encoder is not None
+
+    test_utterances = [
+        "What is the balance on my account?",
+        "How do I reset my online banking password?",
+    ]
+
+    predictions = scorer.predict(test_utterances)
+    if multilabel:
+        assert np.sum(predictions) <= len(test_utterances) * 4
+    else:
+        np.testing.assert_almost_equal(np.sum(predictions), len(test_utterances))
+
+    assert predictions.shape == (len(test_utterances), len(data_handler.intent_descriptions))
+    np.testing.assert_almost_equal(predictions, np.array(expected_prediction).reshape(predictions.shape), decimal=1)
+
+    predictions, metadata = scorer.predict_with_metadata(test_utterances)
+    assert len(predictions) == len(test_utterances)
+    assert metadata is None
+
+    scorer.clear_cache()
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        scorer.dump(temp_dir)
+
+        new_scorer = DescriptionScorer(
+            cross_encoder_config="cross-encoder/ms-marco-MiniLM-L-6-v2", encoder_type="cross", temperature=0.3
+        )
+        new_scorer.load(temp_dir)
+
+        loaded_predictions = new_scorer.predict(test_utterances)
+
+        np.testing.assert_almost_equal(predictions, loaded_predictions, decimal=5)
+
+        new_scorer.clear_cache()