fix: typing

Author: Darinochka

autointent/modules/abc/_base.py

@@ -34,9 +34,7 @@ def fit(self, *args: tuple[Any], **kwargs: dict[str, Any]) -> None:
         """
 
     @abstractmethod
-    def score(
-        self, context: Context, split: Literal["validation", "test"], metrics: list[str]
-    ) -> dict[str, float | str]:
+    def score(self, context: Context, split: Literal["validation", "test"], metrics: list[str]) -> dict[str, float]:
         """
         Calculate metric on test set and return metric value.
 
@@ -108,7 +106,7 @@ def get_embedder_name(self) -> str | None:
         return None
 
     @staticmethod
-    def score_metrics(params: tuple[Any, Any], metrics_dict: dict[str, Any]) -> dict[str, float | str]:
+    def score_metrics(params: tuple[Any, Any], metrics_dict: dict[str, Any]) -> dict[str, float]:
         """
         Score metrics on the test set.
 

autointent/modules/abc/_decision.py

@@ -40,11 +40,7 @@ def predict(self, scores: npt.NDArray[Any]) -> ListOfGenericLabels:
         :param scores: Scores to predict
         """
 
-    def score(
-        self,
-        context: Context,
-        split: Literal["validation", "test"],
-    ) -> dict[str, float | str]:
+    def score(self, context: Context, split: Literal["validation", "test"], metrics: list[str]) -> dict[str, float]:
         """
         Calculate metric on test set and return metric value.
 
@@ -54,7 +50,8 @@ def score(
         """
         labels, scores = get_decision_evaluation_data(context, split)
         self._decisions = self.predict(scores)
-        return self.score_metrics((labels, self._decisions), PREDICTION_METRICS_MULTICLASS)
+        chosen_metrics = {name: fn for name, fn in PREDICTION_METRICS_MULTICLASS.items() if name in metrics}
+        return self.score_metrics((labels, self._decisions), chosen_metrics)
 
     def get_assets(self) -> DecisionArtifact:
         """Return useful assets that represent intermediate data into context."""

autointent/modules/abc/_scoring.py

@@ -21,9 +21,7 @@ class ScoringModule(Module, ABC):
 
     supports_oos = False
 
-    def score(
-        self, context: Context, split: Literal["validation", "test"], metrics: list[str]
-    ) -> dict[str, float | str]:
+    def score(self, context: Context, split: Literal["validation", "test"], metrics: list[str]) -> dict[str, float]:
         """
         Evaluate the scorer on a test set and compute the specified metric.
 

autointent/modules/embedding/_logreg.py

@@ -129,9 +129,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
 
         self._classifier.fit(embeddings, labels)
 
-    def score(
-        self, context: Context, split: Literal["validation", "test"], metrics: list[str]
-    ) -> dict[str, float | str]:
+    def score(self, context: Context, split: Literal["validation", "test"], metrics: list[str]) -> dict[str, float]:
         """
         Evaluate the embedding model using a specified metric function.
 

autointent/modules/embedding/_retrieval.py

@@ -109,9 +109,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         )
         self._vector_index.add(utterances, labels)
 
-    def score(
-        self, context: Context, split: Literal["validation", "test"], metrics: list[str]
-    ) -> dict[str, float | str]:
+    def score(self, context: Context, split: Literal["validation", "test"], metrics: list[str]) -> dict[str, float]:
         """
         Evaluate the embedding model using a specified metric function.
 

autointent/modules/regexp/_regexp.py

@@ -108,9 +108,7 @@ def _predict_single(self, utterance: str) -> tuple[LabelType, dict[str, list[str
             matches["partial_matches"].extend(intent_matches["partial_matches"])
         return list(prediction), matches
 
-    def score(
-        self, context: Context, split: Literal["validation", "test"], metrics: list[str]
-    ) -> dict[str, float | str]:
+    def score(self, context: Context, split: Literal["validation", "test"], metrics: list[str]) -> dict[str, float]:
         """
         Calculate metric on test set and return metric value.