fix names (#37)

* fix names

* lint

Author: Samoed

autointent/context/optimization_info/data_models.py

@@ -4,6 +4,8 @@
 from numpy.typing import NDArray
 from pydantic import BaseModel, ConfigDict, Field
 
+from autointent.custom_types import NodeType
+
 
 class Artifact(BaseModel): ...
 
@@ -40,7 +42,7 @@ class PredictorArtifact(Artifact):
 
 
 def validate_node_name(value: str) -> str:
-    if value in ["regexp", "retrieval", "scoring", "prediction"]:
+    if value in [NodeType.retrieval, NodeType.scoring, NodeType.prediction, NodeType.regexp]:
         return value
     msg = f"Unknown node_type: {value}. Expected one of ['regexp', 'retrieval', 'scoring', 'prediction']"
     raise ValueError(msg)

autointent/context/optimization_info/optimization_info.py

@@ -4,6 +4,7 @@
 from numpy.typing import NDArray
 
 from autointent.configs.node import InferenceNodeConfig
+from autointent.custom_types import NODE_TYPES, NodeType
 from autointent.logger import get_logger
 
 from .data_models import Artifact, Artifacts, RetrieverArtifact, ScorerArtifact, Trial, Trials, TrialsIds
@@ -71,32 +72,28 @@ def _get_best_artifact(self, node_type: str) -> RetrieverArtifact | ScorerArtifa
         return self.artifacts.get_best_artifact(node_type, i_best)
 
     def get_best_embedder(self) -> str:
-        best_retriever_artifact: RetrieverArtifact = self._get_best_artifact(node_type="retrieval")  # type: ignore[assignment]
+        best_retriever_artifact: RetrieverArtifact = self._get_best_artifact(node_type=NodeType.retrieval)  # type: ignore[assignment]
         return best_retriever_artifact.embedder_name
 
     def get_best_test_scores(self) -> NDArray[np.float64] | None:
-        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type="scoring")  # type: ignore[assignment]
+        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
         return best_scorer_artifact.test_scores
 
     def get_best_oos_scores(self) -> NDArray[np.float64] | None:
-        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type="scoring")  # type: ignore[assignment]
+        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
         return best_scorer_artifact.oos_scores
 
     def dump_evaluation_results(self) -> dict[str, dict[str, list[float]]]:
-        node_wise_metrics = {
-            node_type: self._get_metrics_values(node_type)
-            for node_type in ["regexp", "retrieval", "scoring", "prediction"]
-        }
+        node_wise_metrics = {node_type.value: self._get_metrics_values(node_type) for node_type in NODE_TYPES}
         return {
             "metrics": node_wise_metrics,
             "configs": self.trials.model_dump(),
         }
 
     def get_inference_nodes_config(self) -> list[InferenceNodeConfig]:
-        node_types = ["regexp", "retrieval", "scoring", "prediction"]
-        trial_ids = [self._get_best_trial_idx(node_type) for node_type in node_types]
+        trial_ids = [self._get_best_trial_idx(node_type) for node_type in NODE_TYPES]
         res = []
-        for idx, node_type in zip(trial_ids, node_types, strict=True):
+        for idx, node_type in zip(trial_ids, NODE_TYPES, strict=True):
             if idx is None:
                 continue
             trial = self.trials.get_trial(node_type, idx)

autointent/custom_types.py

@@ -19,3 +19,13 @@ class LogLevel(Enum):
 
 class BaseMetadataDict(TypedDict):
     pass
+
+
+class NodeType(str, Enum):
+    retrieval = "retrieval"
+    prediction = "prediction"
+    scoring = "scoring"
+    regexp = "regexp"
+
+
+NODE_TYPES = [NodeType.retrieval, NodeType.prediction, NodeType.scoring, NodeType.regexp]

autointent/generation/prompt_scheme.py

@@ -6,8 +6,8 @@
 class PromptDescription(BaseModel):
     text: str = PROMPT_DESCRIPTION
 
-    @field_validator("text")
     @classmethod
+    @field_validator("text")
     def check_valid_prompt(cls, value: str) -> str:
         if value.find("{intent_name}") == -1 or value.find("{user_utterances}") == -1:
             text_error = (

autointent/modules/__init__.py

@@ -1,3 +1,5 @@
+from typing import TypeVar
+
 from .base import Module
 from .prediction import (
     ArgmaxPredictor,
@@ -10,36 +12,31 @@
 from .retrieval import RetrievalModule, VectorDBModule
 from .scoring import DescriptionScorer, DNNCScorer, KNNScorer, LinearScorer, MLKnnScorer, ScoringModule
 
-RETRIEVAL_MODULES_MULTICLASS: dict[str, type[Module]] = {
-    "vector_db": VectorDBModule,
-}
+T = TypeVar("T", bound=Module)
+
+
+def create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]:
+    return {module.name: module for module in modules}
+
+
+RETRIEVAL_MODULES_MULTICLASS: dict[str, type[Module]] = create_modules_dict([VectorDBModule])
 
 RETRIEVAL_MODULES_MULTILABEL = RETRIEVAL_MODULES_MULTICLASS
 
-SCORING_MODULES_MULTICLASS: dict[str, type[ScoringModule]] = {
-    "dnnc": DNNCScorer,
-    "knn": KNNScorer,
-    "linear": LinearScorer,
-    "description": DescriptionScorer,
-}
+SCORING_MODULES_MULTICLASS: dict[str, type[ScoringModule]] = create_modules_dict(
+    [DNNCScorer, KNNScorer, LinearScorer, DescriptionScorer]
+)
+
+SCORING_MODULES_MULTILABEL: dict[str, type[ScoringModule]] = create_modules_dict(
+    [MLKnnScorer, LinearScorer, DescriptionScorer]
+)
 
-SCORING_MODULES_MULTILABEL: dict[str, type[ScoringModule]] = {
-    "knn": KNNScorer,
-    "linear": LinearScorer,
-    "mlknn": MLKnnScorer,
-}
+PREDICTION_MODULES_MULTICLASS: dict[str, type[Module]] = create_modules_dict(
+    [ArgmaxPredictor, JinoosPredictor, ThresholdPredictor, TunablePredictor]
+)
 
-PREDICTION_MODULES_MULTICLASS: dict[str, type[Module]] = {
-    "argmax": ArgmaxPredictor,
-    "jinoos": JinoosPredictor,
-    "threshold": ThresholdPredictor,
-    "tunable": TunablePredictor,
-}
+PREDICTION_MODULES_MULTILABEL: dict[str, type[Module]] = create_modules_dict([ThresholdPredictor, TunablePredictor])
 
-PREDICTION_MODULES_MULTILABEL: dict[str, type[Module]] = {
-    "threshold": ThresholdPredictor,
-    "tunable": TunablePredictor,
-}
 __all__ = [
     "Module",
     "ArgmaxPredictor",

autointent/modules/base.py

@@ -11,6 +11,8 @@
 
 
 class Module(ABC):
+    name: str
+
     metadata_dict_name: str = "metadata.json"
     metadata: BaseMetadataDict
 

autointent/modules/prediction/argmax.py

@@ -13,10 +13,12 @@
 
 class ArgmaxPredictor(PredictionModule):
     metadata = {}  # noqa: RUF012
+    name = "argmax"
 
     def __init__(self) -> None:
         pass
 
+
     @classmethod
     def from_context(cls, context: Context) -> Self:
         return cls()

autointent/modules/prediction/jinoos.py

@@ -22,6 +22,7 @@ class JinoosPredictorDumpMetadata(BaseMetadataDict):
 
 class JinoosPredictor(PredictionModule):
     thresh: float
+    name = "jinoos"
 
     def __init__(
         self,

autointent/modules/prediction/threshold.py

@@ -27,6 +27,7 @@ class ThresholdPredictor(PredictionModule):
     metadata: ThresholdPredictorDumpMetadata
     multilabel: bool
     tags: list[Tag] | None
+    name = "threshold"
 
     def __init__(
         self,

autointent/modules/prediction/tunable.py

@@ -24,6 +24,8 @@ class TunablePredictorDumpMetadata(BaseMetadataDict):
 
 
 class TunablePredictor(PredictionModule):
+    name = "tunable"
+
     def __init__(
         self,
         n_trials: int = 320,