fix typing problems (except DataHandler._split_cv)

Author: voorhs

autointent/_pipeline/_pipeline.py

@@ -3,7 +3,7 @@
 import json
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import yaml
@@ -122,7 +122,7 @@ def _is_inference(self) -> bool:
         """
         return isinstance(self.nodes[NodeType.scoring], InferenceNode)
 
-    def fit(self, dataset: Dataset) -> Context:
+    def fit(self, dataset: Dataset, scheme: Literal["ho", "cv"] = "ho") -> Context:
         """
         Optimize the pipeline from dataset.
 
@@ -134,7 +134,7 @@ def fit(self, dataset: Dataset) -> Context:
             raise RuntimeError(msg)
 
         context = Context()
-        context.set_dataset(dataset)
+        context.set_dataset(dataset, scheme)
         context.configure_logging(self.logging_config)
         context.configure_vector_index(self.vector_index_config, self.embedder_config)
         context.configure_cross_encoder(self.cross_encoder_config)

autointent/configs/_optimization.py

@@ -1,6 +1,7 @@
 """Configuration for the optimization process."""
 
 from pathlib import Path
+from typing import Literal
 
 from pydantic import BaseModel, Field
 
@@ -12,6 +13,8 @@ class DataConfig(BaseModel):
 
     train_path: str | Path
     """Path to the training data. Can be local path or HF repo."""
+    scheme: Literal["ho", "cv"]
+    """Hold-out or cross-validation."""
 
 
 class TaskConfig(BaseModel):

autointent/context/_context.py

@@ -3,7 +3,7 @@
 import json
 import logging
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import yaml
 
@@ -83,9 +83,10 @@ def configure_data(self, config: DataConfig) -> None:
         self.data_handler = DataHandler(
             dataset=load_dataset(config.train_path),
             random_seed=self.seed,
+            scheme=config.scheme
         )
 
-    def set_dataset(self, dataset: Dataset) -> None:
+    def set_dataset(self, dataset: Dataset, scheme: Literal["ho", "cv"]) -> None:
         """
         Set the datasets for training, validation and testing.
 
@@ -94,6 +95,7 @@ def set_dataset(self, dataset: Dataset) -> None:
         self.data_handler = DataHandler(
             dataset=dataset,
             random_seed=self.seed,
+            scheme=scheme,
         )
 
     def get_inference_config(self) -> dict[str, Any]:

autointent/context/data_handler/_data_handler.py

@@ -9,7 +9,7 @@
 from transformers import set_seed
 
 from autointent import Dataset
-from autointent.custom_types import ListOfGenericLabels, Split
+from autointent.custom_types import ListOfGenericLabels, ListOfLabels, Split
 
 from ._stratification import split_dataset
 
@@ -169,7 +169,7 @@ def test_labels(self, idx: int | None = None) -> ListOfGenericLabels:
         split = f"{Split.TEST}_{idx}" if idx is not None else Split.TEST
         return cast(ListOfGenericLabels, self.dataset[split][self.dataset.label_feature])
 
-    def validation_iterator(self) -> Generator[tuple[list, list, list, list]]:
+    def validation_iterator(self) -> Generator[tuple[list[str], ListOfLabels, list[str], ListOfLabels]]:
         if self.scheme == "ho":
             msg = "Cannot call cross-validation on hold-out DataHandler"
             raise RuntimeError(msg)
@@ -180,7 +180,7 @@ def validation_iterator(self) -> Generator[tuple[list, list, list, list]]:
             train_folds = [i for i in range(self.n_folds) if i != j]
             train_utterances = [ut for i_fold in train_folds for ut in self.train_utterances(i_fold)]
             train_labels = [ut for i_fold in train_folds for ut in self.train_labels(i_fold)]
-            yield train_utterances, train_labels, val_utterances, val_labels
+            yield train_utterances, train_labels, val_utterances, val_labels  # type: ignore[misc]
 
         msg = "something's wrong"
         raise RuntimeError(msg)

autointent/modules/abc/_base.py

@@ -12,7 +12,7 @@
 from autointent._dump_tools import Dumper
 from autointent.context import Context
 from autointent.context.optimization_info import Artifact
-from autointent.custom_types import ListOfGenericLabels
+from autointent.custom_types import ListOfGenericLabels, ListOfLabels
 from autointent.exceptions import WrongClassificationError
 
 logger = logging.getLogger(__name__)
@@ -133,20 +133,22 @@ def score_metrics_ho(params: tuple[Any, Any], metrics_dict: dict[str, Any]) -> d
         return metrics
 
     def score_metrics_cv(
-        self, metrics_dict: dict[str, Any], cv_iterator: Iterable[tuple[list, list, list, list]]
+        self,
+        metrics_dict: dict[str, Any],
+        cv_iterator: Iterable[tuple[list[str], ListOfLabels, list[str], ListOfLabels]],
     ) -> tuple[dict[str, float], list[ListOfGenericLabels] | list[npt.NDArray[Any]]]:
-        metrics_values = {name: [] for name in metrics_dict}
+        metrics_values: dict[str, list[float]] = {name: [] for name in metrics_dict}
         all_val_preds = []
 
         for train_utterances, train_labels, val_utterances, val_labels in cv_iterator:
-            self.fit(train_utterances, train_labels)
+            self.fit(train_utterances, train_labels)  # type: ignore[arg-type]
             val_preds = self.predict(val_utterances)
             for name, fn in metrics_dict.items():
                 metrics_values[name].append(fn(val_labels, val_preds))
             all_val_preds.append(val_preds)
 
-        metrics = {name: np.mean(values_list) for name, values_list in metrics_values.items()}
-        return metrics, all_val_preds
+        metrics = {name: float(np.mean(values_list)) for name, values_list in metrics_values.items()}
+        return metrics, all_val_preds  # type: ignore[return-value]
 
     def _validate_multilabel(self, data_is_multilabel: bool) -> None:
         if data_is_multilabel and not self.supports_multilabel:

autointent/modules/abc/_decision.py

@@ -48,8 +48,8 @@ def score_ho(self, context: Context, metrics: list[str]) -> dict[str, float]:
         :param split: Target split
         :return: Computed metrics value for the test set or error code of metrics
         """
-        train_scores, train_labels = self.get_train_data(context)
-        self.fit(train_scores, train_labels, context.data_handler.tags)
+        train_scores, train_labels, tags = self.get_train_data(context)
+        self.fit(train_scores, train_labels, tags)
 
         val_labels, val_scores = get_decision_evaluation_data(context, "validation")
         decisions = self.predict(val_scores)
@@ -73,22 +73,22 @@ def score_cv(self, context: Context, metrics: list[str]) -> dict[str, float]:
             raise RuntimeError(msg)
 
         chosen_metrics = {name: fn for name, fn in PREDICTION_METRICS_MULTICLASS.items() if name in metrics}
-        metrics_values = {name: [] for name in chosen_metrics}
+        metrics_values: dict[str, list[float]] = {name: [] for name in chosen_metrics}
         all_val_decisions = []
         for j in range(context.data_handler.n_folds):
             val_labels = labels[j]
             val_scores = scores[j]
             train_folds = [i for i in range(context.data_handler.n_folds) if i != j]
             train_labels = [ut for i_fold in train_folds for ut in labels[i_fold]]
             train_scores = [ut for i_fold in train_folds for ut in scores[i_fold]]
-            self.fit(train_scores, train_labels, context.data_handler.tags)
+            self.fit(train_scores, train_labels, context.data_handler.tags)  # type: ignore[arg-type]
             val_decisions = self.predict(val_scores)
             for name, fn in chosen_metrics.items():
                 metrics_values[name].append(fn(val_labels, val_decisions))
             all_val_decisions.append(val_decisions)
 
         self._artifact = DecisionArtifact(labels=[pred for pred_list in all_val_decisions for pred in pred_list])
-        return {name: np.mean(values_list) for name, values_list in metrics_values.items()}
+        return {name: float(np.mean(values_list)) for name, values_list in metrics_values.items()}
 
     def get_assets(self) -> DecisionArtifact:
         """Return useful assets that represent intermediate data into context."""

autointent/modules/abc/_embedding.py

@@ -11,4 +11,4 @@ class EmbeddingModule(Module, ABC):
     """Base class for embedding modules."""
 
     def get_train_data(self, context: Context) -> tuple[list[str], ListOfLabels]:
-        return (context.data_handler.train_utterances(0), context.data_handler.train_labels(0))
+        return (context.data_handler.train_utterances(0), context.data_handler.train_labels(0))  # type: ignore[return-value]

autointent/modules/abc/_scoring.py

@@ -22,6 +22,14 @@ class ScoringModule(Module, ABC):
 
     supports_oos = False
 
+    @abstractmethod
+    def fit(
+        self,
+        utterances: list[str],
+        labels: ListOfLabels,
+    ) -> None:
+        ...
+
     def score_ho(self, context: Context, metrics: list[str]) -> dict[str, float]:
         train_utterances, train_labels = self.get_train_data(context)
         self.fit(train_utterances, train_labels)
@@ -68,7 +76,7 @@ def get_assets(self) -> ScorerArtifact:
         return self._artifact
 
     def get_train_data(self, context: Context) -> tuple[list[str], ListOfLabels]:
-        return (context.data_handler.train_utterances(0), context.data_handler.train_labels(0))
+        return (context.data_handler.train_utterances(0), context.data_handler.train_labels(0))  # type: ignore[return-value]
 
     @abstractmethod
     def predict(self, utterances: list[str]) -> npt.NDArray[Any]:

autointent/modules/scoring/_description/description.py

@@ -148,7 +148,7 @@ def clear_cache(self) -> None:
         self._embedder.clear_ram()
 
     def get_train_data(self, context: Context) -> tuple[list[str], ListOfLabels, list[str]]:
-        return (
+        return (  # type: ignore[return-value]
             context.data_handler.train_utterances(0),
             context.data_handler.train_labels(0),
             context.data_handler.intent_descriptions,

autointent/nodes/_optimization/_node_optimizer.py

@@ -67,7 +67,7 @@ def fit(self, context: Context) -> None:
                     module_kwargs["embedder_name"] = embedder_name
 
                 self._logger.debug("scoring %s module...", module_name)
-                metrics_score = module.score(context, test=False, metrics=self.metrics)
+                metrics_score = module.score(context, metrics=self.metrics)
                 metric_value = metrics_score[self.target_metric]
 
                 context.callback_handler.log_metrics(metrics_score)