finish fixing typing

Author: voorhs

autointent/_vector_index.py

@@ -86,7 +86,7 @@ def add(self, texts: list[str], labels: ListOfLabels) -> None:
         if not hasattr(self, "index"):
             self.index = faiss.IndexFlatIP(embeddings.shape[1])
         self.index.add(embeddings)
-        self.labels.extend(labels)
+        self.labels.extend(labels)  # type: ignore[arg-type]
         self.texts.extend(texts)
 
     def is_empty(self) -> bool:
@@ -186,9 +186,9 @@ def query(
         func = self._search_by_text if isinstance(queries[0], str) else self._search_by_embedding
         all_results = func(queries, k)  # type: ignore[arg-type]
 
-        all_labels = [[self.labels[result["id"]] for result in results] for results in all_results]
+        all_labels: list[ListOfLabels] = [[self.labels[result["id"]] for result in results] for results in all_results]
         all_distances = [[float(result["distance"]) for result in results] for results in all_results]
-        all_texts = [[self.texts[result["id"]] for result in results] for results in all_results]
+        all_texts: list[list[str]] = [[self.texts[result["id"]] for result in results] for results in all_results]
 
         return all_labels, all_distances, all_texts
 

autointent/metrics/custom_types.py

@@ -4,10 +4,10 @@
 
 import numpy.typing as npt
 
-from autointent.custom_types import LabelType
+from autointent.custom_types import ListOfLabels
 
-LABELS_VALUE_TYPE = list[LabelType] | npt.NDArray[Any]
+LABELS_VALUE_TYPE = ListOfLabels
 
-CANDIDATE_TYPE = list[list[LabelType]] | npt.NDArray[Any]
+CANDIDATE_TYPE = list[ListOfLabels] | npt.NDArray[Any]
 
 SCORES_VALUE_TYPE = list[list[float]] | npt.NDArray[Any]

autointent/metrics/retrieval.py

@@ -1,6 +1,5 @@
 """Retrieval metrics."""
 
-from collections.abc import Callable
 from typing import Any, Protocol
 
 import numpy as np
@@ -36,7 +35,7 @@ def __call__(
 
 
 def _macrofy(
-    metric_fn: Callable[[npt.NDArray[Any], npt.NDArray[Any], int | None], float],
+    metric_fn: RetrievalMetricFn,
     query_labels: LABELS_VALUE_TYPE,
     candidates_labels: CANDIDATE_TYPE,
     k: int | None = None,
@@ -72,7 +71,7 @@ def _macrofy(
     for i in range(n_classes):
         binarized_query_labels = query_labels_[..., i]
         binarized_candidates_labels = candidates_labels_[..., i]
-        classwise_values.append(metric_fn(binarized_query_labels, binarized_candidates_labels, k))
+        classwise_values.append(metric_fn(binarized_query_labels, binarized_candidates_labels, k))  # type: ignore[arg-type]
 
     return np.mean(classwise_values)  # type: ignore[return-value]
 
@@ -136,12 +135,12 @@ def retrieval_map(query_labels: LABELS_VALUE_TYPE, candidates_labels: CANDIDATE_
     :param k: Number of top items to consider for each query
     :return: Score of the retrieval metric
     """
-    ap_list = [_average_precision(q, c, k) for q, c in zip(query_labels, candidates_labels, strict=True)]
+    ap_list = [_average_precision(q, c, k) for q, c in zip(query_labels, candidates_labels, strict=True)]  # type: ignore[arg-type]
     return sum(ap_list) / len(ap_list)
 
 
 def _average_precision_intersecting(
-    query_label: LABELS_VALUE_TYPE, candidate_labels: CANDIDATE_TYPE, k: int | None = None
+    query_label: list[int], candidate_labels: CANDIDATE_TYPE, k: int | None = None
 ) -> float:
     r"""
     Calculate the average precision at position k for the intersecting labels.
@@ -212,7 +211,7 @@ def retrieval_map_intersecting(
     :param k: Number of top items to consider for each query
     :return: Score of the retrieval metric
     """
-    ap_list = [_average_precision_intersecting(q, c, k) for q, c in zip(query_labels, candidates_labels, strict=True)]
+    ap_list = [_average_precision_intersecting(q, c, k) for q, c in zip(query_labels, candidates_labels, strict=True)]  # type: ignore[arg-type]
     return sum(ap_list) / len(ap_list)
 
 

autointent/modules/abc/_decision.py

@@ -63,9 +63,7 @@ def get_assets(self) -> DecisionArtifact:
     def clear_cache(self) -> None:
         """Clear cache."""
 
-    def _validate_inputs(
-        self, scores: npt.NDArray[Any], labels: ListOfGenericLabels
-    ) -> tuple[int, bool]:
+    def _validate_inputs(self, scores: npt.NDArray[Any], labels: ListOfGenericLabels) -> tuple[int, bool, bool]:
         """
         Sanity check if labels and scores are valid to be a training data for decision module.
 

autointent/modules/decision/_adaptive.py

@@ -7,7 +7,7 @@
 import numpy.typing as npt
 
 from autointent import Context
-from autointent.custom_types import ListOfGenericLabels, ListOfLabelsWithOOS
+from autointent.custom_types import ListOfGenericLabels, ListOfLabelsWithOOS, MultiLabel
 from autointent.metrics import decision_f1
 from autointent.modules.abc import DecisionModule
 from autointent.schemas import Tag
@@ -146,10 +146,11 @@ def multilabel_predict(scores: npt.NDArray[Any], r: float, tags: list[Tag] | Non
     res = (scores >= thresh[:, None]).astype(int)
     if tags:
         res = apply_tags(res, scores, tags)
-    return [lab if sum(lab) > 0 else None for lab in res.tolist()]
+    y_pred: list[MultiLabel] = res.tolist()  # type: ignore[assignment]
+    return [lab if sum(lab) > 0 else None for lab in y_pred]
 
 
-def multilabel_score(y_true: ListOfLabelsWithOOS, y_pred: ListOfLabelsWithOOS) -> float:
+def multilabel_score(y_true: ListOfGenericLabels, y_pred: ListOfGenericLabels) -> float:
     """
     Calculate the weighted F1 score for multi-label classification.
 

autointent/modules/decision/_jinoos.py

@@ -114,7 +114,7 @@ def predict(self, scores: npt.NDArray[Any]) -> list[int | None]:
         if scores.shape[1] != self._n_classes:
             raise InvalidNumClassesError
         pred_classes, best_scores = _predict(scores)
-        y_pred = _detect_oos(pred_classes, best_scores, self._thresh).tolist()
+        y_pred: list[int] = _detect_oos(pred_classes, best_scores, self._thresh).tolist()  # type: ignore[assignment]
         return [lab if lab != -1 else None for lab in y_pred]
 
     @staticmethod

autointent/modules/decision/_threshold.py

@@ -7,7 +7,7 @@
 import numpy.typing as npt
 
 from autointent import Context
-from autointent.custom_types import ListOfGenericLabels
+from autointent.custom_types import ListOfGenericLabels, MultiLabel
 from autointent.modules.abc import DecisionModule
 from autointent.schemas import Tag
 
@@ -141,7 +141,7 @@ def predict(self, scores: npt.NDArray[Any]) -> ListOfGenericLabels:
         return multiclass_predict(scores, self.thresh)
 
 
-def multiclass_predict(scores: npt.NDArray[Any], thresh: float | npt.NDArray[Any]) -> npt.NDArray[Any]:
+def multiclass_predict(scores: npt.NDArray[Any], thresh: float | npt.NDArray[Any]) -> ListOfGenericLabels:
     """
     Make predictions for multiclass classification task.
 
@@ -158,14 +158,15 @@ def multiclass_predict(scores: npt.NDArray[Any], thresh: float | npt.NDArray[Any
         thresh_selected = thresh[pred_classes]
         pred_classes[best_scores < thresh_selected] = -1  # out of scope
 
-    return [lab if lab != -1 else None for lab in pred_classes.tolist()]
+    y_pred: list[int] = pred_classes.tolist()  # type: ignore[assignment]
+    return [lab if lab != -1 else None for lab in y_pred]
 
 
 def multilabel_predict(
     scores: npt.NDArray[Any],
     thresh: float | npt.NDArray[Any],
     tags: list[Tag] | None,
-) -> npt.NDArray[Any]:
+) -> ListOfGenericLabels:
     """
     Make predictions for multilabel classification task.
 
@@ -177,4 +178,5 @@ def multilabel_predict(
     res = (scores >= thresh).astype(int) if isinstance(thresh, float) else (scores >= thresh[None, :]).astype(int)
     if tags:
         res = apply_tags(res, scores, tags)
-    return [lab if sum(lab) > 0 else None for lab in res.tolist()]
+    y_pred: list[MultiLabel] = res.tolist()  # type: ignore[assignment]
+    return [lab if sum(lab) > 0 else None for lab in y_pred]

autointent/modules/decision/_tunable.py

@@ -176,7 +176,7 @@ def objective(self, trial: Trial) -> float:
             y_pred = multilabel_predict(self.probas, thresholds, self.tags)
         else:
             y_pred = multiclass_predict(self.probas, thresholds)
-        return decision_f1(self.labels, y_pred)  # type: ignore[no-any-return]
+        return decision_f1(self.labels, y_pred)
 
     def fit(
         self,

autointent/modules/embedding/_retrieval.py

@@ -167,7 +167,7 @@ def load(self, path: str) -> None:
         """
         self._vector_index = VectorIndex.load(Path(path))
 
-    def predict(self, utterances: list[str]) -> tuple[list[list[int | list[int]]], list[list[float]], list[list[str]]]:
+    def predict(self, utterances: list[str]) -> tuple[list[ListOfLabels], list[list[float]], list[list[str]]]:
         """
         Predict the nearest neighbors for a list of utterances.
 

autointent/modules/regexp/_regexp.py

@@ -1,8 +1,6 @@
 """Module for regular expressions based intent detection."""
 
-import json
 import re
-from pathlib import Path
 from typing import Any, Literal, TypedDict
 
 from autointent import Context
@@ -140,30 +138,6 @@ def get_assets(self) -> Artifact:
         """Get assets."""
         return Artifact()
 
-    def load(self, path: str) -> None:
-        """
-        Load data from dump.
-
-        :param path: Path to load
-        """
-        dump_dir = Path(path)
-
-        with (dump_dir / self.metadata_dict_name).open() as file:
-            self.regexp_patterns = json.load(file)
-
-        self._compile_regex_patterns()
-
-    def dump(self, path: str) -> None:
-        """
-        Dump all data needed for inference.
-
-        :param path: Path to dump
-        """
-        dump_dir = Path(path)
-
-        with (dump_dir / self.metadata_dict_name).open("w") as file:
-            json.dump(self.regexp_patterns, file, indent=4)
-
     def _compile_regex_patterns(self) -> None:
         """Compile regex patterns."""
         self.regexp_patterns_compiled = [