Rerank scorer: опция для выбора источника для расчета вектора вероятностей (#115)

* Enable rerank scorer to use crossencoder scores for the probability vector

* add cross encoder scores range options

* upd test

---------

Co-authored-by: voorhs <[email protected]>

Author: Dmitryv-2024

autointent/_ranker.py

@@ -10,7 +10,7 @@
 import logging
 from pathlib import Path
 from random import shuffle
-from typing import Any, TypedDict
+from typing import Any, Literal, TypedDict
 
 import joblib
 import numpy as np
@@ -101,12 +101,14 @@ def __init__(
         self,
         cross_encoder_config: CrossEncoderConfig | str | dict[str, Any],
         classifier_head: LogisticRegressionCV | None = None,
+        output_range: Literal["sigmoid", "tanh"] = "sigmoid",
     ) -> None:
         """Initialize the Ranker.
 
         Args:
             cross_encoder_config: Configuration for the cross-encoder model
             classifier_head: Optional pre-trained classifier head
+            output_range: Range of the output probabilities ([0, 1] for sigmoid, [-1, 1] for tanh)
         """
         self.config = CrossEncoderConfig.from_search_config(cross_encoder_config)
         self.cross_encoder = st.CrossEncoder(
@@ -117,6 +119,7 @@ def __init__(
         )
         self._train_head = False
         self._clf = classifier_head
+        self.output_range = output_range
 
         if classifier_head is not None or self.config.train_head:
             self._train_head = True
@@ -148,7 +151,7 @@ def _get_features_or_predictions(self, pairs: list[tuple[str, str]]) -> npt.NDAr
                 self.cross_encoder.predict(
                     pairs,
                     batch_size=self.config.batch_size,
-                    activation_fct=nn.Sigmoid(),
+                    activation_fct=nn.Sigmoid() if self.output_range == "sigmoid" else nn.Tanh(),
                 )
             )
 
@@ -210,7 +213,10 @@ def predict(self, pairs: list[tuple[str, str]]) -> npt.NDArray[Any]:
         features = self._get_features_or_predictions(pairs)
 
         if self._clf is not None:
-            return np.array(self._clf.predict_proba(features)[:, 1])
+            probs = np.array(self._clf.predict_proba(features)[:, 1])
+            if self.output_range == "tanh":
+                probs = probs * 2 - 1
+            return probs
         return features
 
     def rank(

autointent/modules/scoring/_dnnc/dnnc.py

@@ -119,7 +119,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         self._vector_index = VectorIndex(self.embedder_config)
         self._vector_index.add(utterances, labels)
 
-        self._cross_encoder = Ranker(self.cross_encoder_config)
+        self._cross_encoder = Ranker(self.cross_encoder_config, output_range="sigmoid")
         self._cross_encoder.fit(utterances, labels)
 
     def predict(self, utterances: list[str]) -> npt.NDArray[Any]:

autointent/modules/scoring/_knn/rerank_scorer.py

@@ -29,7 +29,6 @@ class RerankScorer(KNNScorer):
 
         cross_encoder_config: Config of the cross-encoder model used for re-ranking
         m: Number of top-ranked neighbors to consider, or None to use k
-        rank_threshold_cutoff: Rank threshold cutoff for re-ranking, or None
     """
 
     name = "rerank"
@@ -38,9 +37,9 @@ class RerankScorer(KNNScorer):
     def __init__(
         self,
         k: int,
-        weights: WeightType,
+        weights: WeightType = "distance",
+        use_crosencoder_scores: bool = False,
         m: int | None = None,
-        rank_threshold_cutoff: int | None = None,
         cross_encoder_config: CrossEncoderConfig | str | dict[str, Any] | None = None,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
     ) -> None:
@@ -53,18 +52,12 @@ def __init__(
         self.cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
 
         self.m = k if m is None else m
-        self.rank_threshold_cutoff = rank_threshold_cutoff
+        self.use_crosencoder_scores = use_crosencoder_scores
 
         if self.m < 0 or not isinstance(self.m, int):
             msg = "`m` argument of `RerankScorer` must be a positive int"
             raise ValueError(msg)
 
-        if self.rank_threshold_cutoff is not None and (
-            self.rank_threshold_cutoff < 0 or not isinstance(self.rank_threshold_cutoff, int)
-        ):
-            msg = "`rank_threshold_cutoff` argument of `RerankScorer` must be a positive int or None"
-            raise ValueError(msg)
-
     @classmethod
     def from_context(
         cls,
@@ -74,7 +67,7 @@ def from_context(
         m: PositiveInt | None = None,
         cross_encoder_config: CrossEncoderConfig | str | None = None,
         embedder_config: EmbedderConfig | str | None = None,
-        rank_threshold_cutoff: int | None = None,
+        use_crosencoder_scores: bool = False,
     ) -> "RerankScorer":
         """Create a RerankScorer instance from a given context.
 
@@ -86,7 +79,7 @@ def from_context(
             embedder_config: Config of the embedder used for vectorization,
                 or None to use the best existing embedder
             m: Number of top-ranked neighbors to consider, or None to use k
-            rank_threshold_cutoff: Rank threshold cutoff for re-ranking, or None
+            use_crosencoder_scores: use crosencoder scores for the output probability vector computation
         """
         if embedder_config is None:
             embedder_config = context.resolve_embedder()
@@ -98,7 +91,7 @@ def from_context(
             k=k,
             weights=weights,
             m=m,
-            rank_threshold_cutoff=rank_threshold_cutoff,
+            use_crosencoder_scores=use_crosencoder_scores,
             embedder_config=embedder_config,
             cross_encoder_config=cross_encoder_config,
         )
@@ -113,9 +106,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         if hasattr(self, "_scorer"):
             self.clear_cache()
 
-        self._scorer = Ranker(
-            self.cross_encoder_config,
-        )
+        self._scorer = Ranker(self.cross_encoder_config, output_range="tanh")
         self._scorer.fit(utterances, labels)
 
         super().fit(utterances, labels, clear_cache=False)
@@ -147,10 +138,14 @@ def _predict(self, utterances: list[str]) -> tuple[npt.NDArray[Any], list[list[s
         ):
             cur_ranks = self._scorer.rank(query, query_docs, top_k=self.m)
 
-            for dst, src in zip(
-                [labels, distances, neighbours], [query_labels, query_distances, query_docs], strict=True
-            ):
+            for dst, src in zip([labels, neighbours], [query_labels, query_docs], strict=True):
                 dst.append([src[rank["corpus_id"]] for rank in cur_ranks])  # type: ignore[attr-defined]
 
+            if self.use_crosencoder_scores:
+                distances.append([rank["score"] for rank in cur_ranks])
+            else:
+                distances.append([query_distances[rank["corpus_id"]] for rank in cur_ranks])
+
         scores = self._count_scores(np.array(labels), np.array(distances))
+
         return scores, neighbours

tests/assets/configs/multiclass.yaml

@@ -23,6 +23,7 @@
       k: [ 5, 10 ]
       weights: [uniform, distance, closest]
       m: [ 2, 3 ]
+      use_crosencoder_scores: [true, false]
       cross_encoder_config:
         - cross-encoder/ms-marco-MiniLM-L-6-v2
     - module_name: sklearn

tests/assets/configs/multilabel.yaml

@@ -18,6 +18,7 @@
     - module_name: rerank
       k: [ 5, 10 ]
       weights: [ uniform, distance, closest ]
+      use_crosencoder_scores: [true, false]
       m: [ 2, 3 ]
       cross_encoder_config:
         - model_name: cross-encoder/ms-marco-MiniLM-L-6-v2

tests/configs/test_scoring.py

@@ -44,7 +44,7 @@ def valid_scoring_config():
                 "embedder_config": ["sergeyzh/rubert-tiny-turbo"],
                 "k": [5],
                 "weights": ["distance"],
-                "rank_threshold_cutoff": [None, 3],
+                "use_crosencoder_scores": [True, False],
             },
             {
                 "module_name": "sklearn",