resolve part conversations

Samoed · Samoed · commit 2e4b5e246289 · 2025-02-12T01:49:10.000+03:00
diff --git a/autointent/_ranker.py b/autointent/_ranker.py
@@ -20,7 +20,7 @@
 from torch import nn
 
 from autointent.custom_types import ListOfLabels
-from autointent.schemas._schemas import CrossEncoderConfig
+from autointent.schemas import CrossEncoderConfig
 
 logger = logging.getLogger(__name__)
 
@@ -106,7 +106,7 @@ class Ranker:
 
     def __init__(
         self,
-        cross_encoder_config: CrossEncoderConfig,
+        cross_encoder_config: CrossEncoderConfig | str | dict[str, Any],
         classifier_head: LogisticRegressionCV | None = None,
     ) -> None:
         """
@@ -116,10 +116,7 @@ def __init__(
         :param max_length (int, optional): Max length for input sequences for the cross encoder.
         :param classifier_head (LogisticRegressionCV, optional): Classifier (to be used in restore procedure mainly).
         """
-        if isinstance(cross_encoder_config, dict):
-            cross_encoder_config = CrossEncoderConfig(**cross_encoder_config)
-        if isinstance(cross_encoder_config, str):
-            cross_encoder_config = CrossEncoderConfig(model_name=cross_encoder_config)
+        cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
         self.cross_encoder = st.CrossEncoder(
             cross_encoder_config.model_name,
             trust_remote_code=True,
diff --git a/autointent/_vector_index.py b/autointent/_vector_index.py
@@ -16,7 +16,7 @@
 
 from autointent import Embedder
 from autointent.custom_types import ListOfLabels
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class VectorIndexMetadata(TypedDict):
diff --git a/autointent/context/optimization_info/_data_models.py b/autointent/context/optimization_info/_data_models.py
@@ -11,7 +11,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 from autointent.custom_types import ListOfLabelsWithOOS, NodeType
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class Artifact(BaseModel):
diff --git a/autointent/context/optimization_info/_optimization_info.py b/autointent/context/optimization_info/_optimization_info.py
@@ -13,7 +13,7 @@
 
 from autointent.configs import InferenceNodeConfig
 from autointent.custom_types import NodeType
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 from ._data_models import Artifact, Artifacts, RetrieverArtifact, ScorerArtifact, Trial, Trials, TrialsIds
 
diff --git a/autointent/modules/abc/_base.py b/autointent/modules/abc/_base.py
@@ -14,6 +14,7 @@
 from autointent.context.optimization_info import Artifact
 from autointent.custom_types import ListOfGenericLabels, ListOfLabels
 from autointent.exceptions import WrongClassificationError
+from autointent.schemas import EmbedderConfig
 
 logger = logging.getLogger(__name__)
 
@@ -114,8 +115,12 @@ def from_context(cls, context: Context, **kwargs: dict[str, Any]) -> "Module":
         :param kwargs: Additional kwargs.
         """
 
-    def get_embedder_name(self) -> str | None:
-        """Experimental method."""
+    def get_embedder_config(self) -> EmbedderConfig | None:
+        """
+        Get the config of the embedder.
+
+        :return: Embedder config.
+        """
         return None
 
     @staticmethod
diff --git a/autointent/modules/embedding/_logreg.py b/autointent/modules/embedding/_logreg.py
@@ -13,7 +13,7 @@
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL
 from autointent.modules.abc import EmbeddingModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class LogregAimedEmbedding(EmbeddingModule):
diff --git a/autointent/modules/embedding/_retrieval.py b/autointent/modules/embedding/_retrieval.py
@@ -7,7 +7,7 @@
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import RETRIEVAL_METRICS_MULTICLASS, RETRIEVAL_METRICS_MULTILABEL
 from autointent.modules.abc import EmbeddingModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class RetrievalAimedEmbedding(EmbeddingModule):
diff --git a/autointent/modules/scoring/_description/description.py b/autointent/modules/scoring/_description/description.py
@@ -12,7 +12,7 @@
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL
 from autointent.modules.abc import ScoringModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class DescriptionScorer(ScoringModule):
@@ -47,12 +47,7 @@ def __init__(
         :param temperature: Temperature parameter for scaling logits, defaults to 1.0.
         """
         self.temperature = temperature
-        if isinstance(embedder_config, dict):
-            embedder_config = EmbedderConfig(**embedder_config)
-        if isinstance(embedder_config, str):
-            embedder_config = EmbedderConfig(model_name=embedder_config)
-
-        self.embedder_config = embedder_config
+        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
 
     @classmethod
     def from_context(
@@ -77,7 +72,7 @@ def from_context(
             embedder_config=embedder_config,
         )
 
-    def get_embedder_name(self) -> EmbedderConfig:
+    def get_embedder_config(self) -> EmbedderConfig:
         """
         Get the name of the embedder.
 
@@ -129,10 +124,10 @@ def predict(self, utterances: list[str]) -> NDArray[np.float64]:
         similarities: NDArray[np.float64] = cosine_similarity(utterance_vectors, self._description_vectors)
 
         if self._multilabel:
-            probabilites = scipy.special.expit(similarities / self.temperature)
+            probabilities = scipy.special.expit(similarities / self.temperature)
         else:
-            probabilites = scipy.special.softmax(similarities / self.temperature, axis=1)
-        return probabilites  # type: ignore[no-any-return]
+            probabilities = scipy.special.softmax(similarities / self.temperature, axis=1)
+        return probabilities  # type: ignore[no-any-return]
 
     def clear_cache(self) -> None:
         """Clear cached data in memory used by the embedder."""
@@ -150,7 +145,7 @@ def score_cv(self, context: Context, metrics: list[str]) -> dict[str, float]:
         Evaluate the scorer on a test set and compute the specified metric.
 
         :param context: Context containing test set and other data.
-        :param split: Target split
+        :param metrics: List of metric names to compute.
         :return: Computed metrics value for the test set or error code of metrics
         """
         metrics_dict = SCORING_METRICS_MULTILABEL if context.is_multilabel() else SCORING_METRICS_MULTICLASS
diff --git a/autointent/modules/scoring/_dnnc/dnnc.py b/autointent/modules/scoring/_dnnc/dnnc.py
@@ -10,7 +10,7 @@
 from autointent import Context, Ranker, VectorIndex
 from autointent.custom_types import ListOfLabels
 from autointent.modules.abc import ScoringModule
-from autointent.schemas._schemas import CrossEncoderConfig, EmbedderConfig
+from autointent.schemas import CrossEncoderConfig, EmbedderConfig
 
 logger = logging.getLogger(__name__)
 
@@ -87,17 +87,8 @@ def __init__(
         :param embedder_config: Config of the embedder model.
         :param k: Number of nearest neighbors to retrieve.
         """
-        if isinstance(cross_encoder_config, dict):
-            cross_encoder_config = CrossEncoderConfig(**cross_encoder_config)
-        if isinstance(cross_encoder_config, str):
-            cross_encoder_config = CrossEncoderConfig(model_name=cross_encoder_config)
-        self.cross_encoder_config = cross_encoder_config
-
-        if isinstance(embedder_config, dict):
-            embedder_config = EmbedderConfig(**embedder_config)
-        if isinstance(embedder_config, str):
-            embedder_config = EmbedderConfig(model_name=embedder_config)
-        self.embedder_config = embedder_config
+        self.cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
+        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.k = k
 
     @classmethod
diff --git a/autointent/modules/scoring/_knn/knn.py b/autointent/modules/scoring/_knn/knn.py
@@ -8,7 +8,7 @@
 from autointent import Context, VectorIndex
 from autointent.custom_types import WEIGHT_TYPES, ListOfLabels
 from autointent.modules.abc import ScoringModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 from .weighting import apply_weights
 
@@ -71,12 +71,7 @@ def __init__(
             - "distance": Weight inversely proportional to distance.
             - "closest": Only the closest neighbor of each class is weighted.
         """
-        if isinstance(embedder_config, dict):
-            embedder_config = EmbedderConfig(**embedder_config)
-        if isinstance(embedder_config, str):
-            embedder_config = EmbedderConfig(model_name=embedder_config)
-
-        self.embedder_config = embedder_config
+        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.k = k
         self.weights = weights
 
@@ -106,13 +101,13 @@ def from_context(
             weights=weights,
         )
 
-    def get_embedder_name(self) -> str:
+    def get_embedder_config(self) -> EmbedderConfig:
         """
         Get the name of the embedder.
 
         :return: Embedder name.
         """
-        return self.embedder_config.model_name
+        return self.embedder_config
 
     def fit(self, utterances: list[str], labels: ListOfLabels, clear_cache: bool = False) -> None:
         """
diff --git a/autointent/modules/scoring/_knn/rerank_scorer.py b/autointent/modules/scoring/_knn/rerank_scorer.py
@@ -7,7 +7,7 @@
 
 from autointent import Context, Ranker
 from autointent.custom_types import WEIGHT_TYPES, ListOfLabels
-from autointent.schemas._schemas import CrossEncoderConfig, EmbedderConfig
+from autointent.schemas import CrossEncoderConfig, EmbedderConfig
 
 from .knn import KNNScorer
 
diff --git a/autointent/modules/scoring/_linear.py b/autointent/modules/scoring/_linear.py
@@ -10,7 +10,7 @@
 from autointent import Context, Embedder
 from autointent.custom_types import ListOfLabels
 from autointent.modules.abc import ScoringModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class LinearScorer(ScoringModule):
@@ -69,12 +69,7 @@ def __init__(
         self.cv = cv
         self.n_jobs = n_jobs
         self.seed = seed
-        if isinstance(embedder_config, dict):
-            embedder_config = EmbedderConfig(**embedder_config)
-        if isinstance(embedder_config, str):
-            embedder_config = EmbedderConfig(model_name=embedder_config)
-
-        self.embedder_config = embedder_config
+        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
 
     @classmethod
     def from_context(
@@ -96,7 +91,7 @@ def from_context(
             embedder_config=embedder_config,
         )
 
-    def get_embedder_name(self) -> EmbedderConfig:
+    def get_embedder_config(self) -> EmbedderConfig:
         """
         Get the name of the embedder.
 
diff --git a/autointent/modules/scoring/_mlknn/mlknn.py b/autointent/modules/scoring/_mlknn/mlknn.py
@@ -8,7 +8,7 @@
 from autointent import Context, VectorIndex
 from autointent.custom_types import ListOfLabels
 from autointent.modules.abc import ScoringModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 
 class MLKnnScorer(ScoringModule):
@@ -72,11 +72,7 @@ def __init__(
         :param ignore_first_neighbours: Number of closest neighbors to ignore, defaults to 0.
         """
         self.k = k
-        if isinstance(embedder_config, dict):
-            embedder_config = EmbedderConfig(**embedder_config)
-        if isinstance(embedder_config, str):
-            embedder_config = EmbedderConfig(model_name=embedder_config)
-        self.embedder_config = embedder_config
+        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.s = s
         self.ignore_first_neighbours = ignore_first_neighbours
 
@@ -109,7 +105,7 @@ def from_context(
             ignore_first_neighbours=ignore_first_neighbours,
         )
 
-    def get_embedder_name(self) -> EmbedderConfig:
+    def get_embedder_config(self) -> EmbedderConfig:
         """
         Get the name of the embedder.
 
diff --git a/autointent/modules/scoring/_sklearn/sklearn_scorer.py b/autointent/modules/scoring/_sklearn/sklearn_scorer.py
@@ -11,7 +11,7 @@
 from autointent import Context, Embedder
 from autointent.custom_types import ListOfLabels
 from autointent.modules.abc import ScoringModule
-from autointent.schemas._schemas import EmbedderConfig
+from autointent.schemas import EmbedderConfig
 
 logger = logging.getLogger(__name__)
 AVAILABLE_CLASSIFIERS = {
@@ -55,12 +55,7 @@ def __init__(
         :param clf_name: Name of the sklearn classifier to use.
         :param clf_args: dictionary with the chosen sklearn classifier arguments, defaults to {}.
         """
-        if isinstance(embedder_config, dict):
-            embedder_config = EmbedderConfig(**embedder_config)
-        if isinstance(embedder_config, str):
-            embedder_config = EmbedderConfig(model_name=embedder_config)
-
-        self.embedder_config = embedder_config
+        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.clf_name = clf_name
         self.clf_args = clf_args or {}
 
diff --git a/autointent/nodes/_optimization/_node_optimizer.py b/autointent/nodes/_optimization/_node_optimizer.py
@@ -63,7 +63,7 @@ def fit(self, context: Context) -> None:
                 )
                 module = self.node_info.modules_available[module_name].from_context(context, **module_kwargs)
 
-                embedder_config = module.get_embedder_name()
+                embedder_config = module.get_embedder_config()
                 if embedder_config is not None:
                     module_kwargs["embedder_config"] = embedder_config
 
diff --git a/autointent/schemas/__init__.py b/autointent/schemas/__init__.py
@@ -1,5 +1,5 @@
 """Data models related to :class:`autointent.Dataset`."""
 
-from ._schemas import Intent, Sample, Tag, TagsList
+from ._schemas import CrossEncoderConfig, EmbedderConfig, Intent, LLMConfig, Sample, STModelConfig, Tag, TagsList
 
-__all__ = ["Intent", "Sample", "Tag", "TagsList"]
+__all__ = ["CrossEncoderConfig", "EmbedderConfig", "Intent", "LLMConfig", "STModelConfig", "Sample", "Tag", "TagsList"]
diff --git a/autointent/schemas/_schemas.py b/autointent/schemas/_schemas.py
@@ -7,7 +7,13 @@
 from pathlib import Path
 from typing import Any
 
-from pydantic import BaseModel, model_validator
+from pydantic import (
+    BaseModel,
+    NonNegativeFloat,
+    PositiveInt,
+    model_validator,
+)
+from typing_extensions import Self
 
 from autointent.custom_types import LabelWithOOS
 
@@ -119,14 +125,14 @@ class Intent(BaseModel):
 
 
 class ModelConfig(BaseModel):
-    batch_size: int = 32
+    batch_size: PositiveInt = 32
     """Batch size for model inference."""
-    max_length: int | None = None
+    max_length: PositiveInt | None = None
     """Maximum length of input sequences."""
 
 
 class LLMConfig(ModelConfig):
-    temperature: float | None = None
+    temperature: NonNegativeFloat | None = None
     """Temperature for sampling from the model."""
     base_url: str | None = None
     """Base URL for the model API."""
@@ -141,8 +147,18 @@ class STModelConfig(ModelConfig):
     """Name of the hugging face model."""
     device: str | None = None
     """Torch notation for CPU or CUDA."""
-    use_cache: bool = True
-    """Whether to use embeddings caching."""
+
+    @classmethod
+    def from_search_config(cls, values: dict[str, Any] | str | BaseModel) -> Self:
+        """Validate the model configuration.
+
+        :param values: Model configuration values. If a string is provided, it is converted to a dictionary.
+        """
+        if isinstance(values, BaseModel):
+            return values
+        if isinstance(values, str):
+            return cls(model_name=values)
+        return cls(**values)
 
 
 class EmbedderConfig(STModelConfig):
@@ -152,6 +168,9 @@ class EmbedderConfig(STModelConfig):
     passage_instruction: str | None = None
     """Instruction for passage."""
 
+    use_cache: bool = True
+    """Whether to use embeddings caching."""
+
 
 class CrossEncoderConfig(STModelConfig):
     train_head: bool = False
diff --git a/tests/assets/configs/multiclass.yaml b/tests/assets/configs/multiclass.yaml
diff --git a/tests/context/test_vector_index.py b/tests/context/test_vector_index.py
