From 7da1ba98e58ea4d41b06ded17212dcb0375da857 Mon Sep 17 00:00:00 2001
From: Roman Solomatin <36135455+Samoed@users.noreply.github.com>
Date: Tue, 5 Nov 2024 20:46:39 +0300
Subject: [PATCH 1/2] fix names

---
 .../context/optimization_info/data_models.py  |  4 +-
 .../optimization_info/optimization_info.py    | 17 +++----
 autointent/custom_types.py                    | 10 +++++
 autointent/generation/prompt_scheme.py        |  2 +-
 autointent/modules/__init__.py                | 45 +++++++++----------
 autointent/modules/base.py                    |  2 +
 autointent/modules/prediction/argmax.py       |  2 +
 autointent/modules/prediction/jinoos.py       |  1 +
 autointent/modules/prediction/threshold.py    |  1 +
 autointent/modules/prediction/tunable.py      |  1 +
 autointent/modules/regexp.py                  |  2 +
 autointent/modules/retrieval/vectordb.py      |  1 +
 .../scoring/description/description.py        |  1 +
 autointent/modules/scoring/dnnc/dnnc.py       |  2 +
 autointent/modules/scoring/knn/knn.py         |  1 +
 autointent/modules/scoring/linear.py          |  1 +
 autointent/modules/scoring/mlknn/mlknn.py     |  1 +
 autointent/nodes/nodes_info/__init__.py       |  8 ++--
 autointent/nodes/nodes_info/base.py           |  3 +-
 autointent/nodes/nodes_info/prediction.py     |  3 +-
 autointent/nodes/nodes_info/regexp.py         |  5 ++-
 autointent/nodes/nodes_info/retrieval.py      |  3 +-
 autointent/nodes/nodes_info/scoring.py        |  3 +-
 .../pipeline/inference/inference_pipeline.py  |  6 +--
 .../optimization/pipeline_optimizer.py        |  3 +-
 pyproject.toml                                |  1 +
 26 files changed, 80 insertions(+), 49 deletions(-)

diff --git a/autointent/context/optimization_info/data_models.py b/autointent/context/optimization_info/data_models.py
index c74dc089..30c4d583 100644
--- a/autointent/context/optimization_info/data_models.py
+++ b/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)
diff --git a/autointent/context/optimization_info/optimization_info.py b/autointent/context/optimization_info/optimization_info.py
index 36795151..949f196d 100644
--- a/autointent/context/optimization_info/optimization_info.py
+++ b/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)
diff --git a/autointent/custom_types.py b/autointent/custom_types.py
index 08083f3c..bc7652c7 100644
--- a/autointent/custom_types.py
+++ b/autointent/custom_types.py
@@ -15,3 +15,13 @@ class LogLevel(Enum):
 WEIGHT_TYPES = Literal["uniform", "distance", "closest"]
 
 LABEL_TYPE = int | list[int]
+
+
+class NodeType(str, Enum):
+    retrieval = "retrieval"
+    prediction = "prediction"
+    scoring = "scoring"
+    regexp = "regexp"
+
+
+NODE_TYPES = [NodeType.retrieval, NodeType.prediction, NodeType.scoring, NodeType.regexp]
diff --git a/autointent/generation/prompt_scheme.py b/autointent/generation/prompt_scheme.py
index 68de2c56..bcea3180 100644
--- a/autointent/generation/prompt_scheme.py
+++ b/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 = (
diff --git a/autointent/modules/__init__.py b/autointent/modules/__init__.py
index 10fd889d..a764f7d4 100644
--- a/autointent/modules/__init__.py
+++ b/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",
diff --git a/autointent/modules/base.py b/autointent/modules/base.py
index a97a0af7..77735ab3 100644
--- a/autointent/modules/base.py
+++ b/autointent/modules/base.py
@@ -9,6 +9,8 @@
 
 
 class Module(ABC):
+    name: str
+
     @abstractmethod
     def fit(self, context: Context) -> None:
         pass
diff --git a/autointent/modules/prediction/argmax.py b/autointent/modules/prediction/argmax.py
index 5b81b78e..5d803bcf 100644
--- a/autointent/modules/prediction/argmax.py
+++ b/autointent/modules/prediction/argmax.py
@@ -10,6 +10,8 @@
 
 
 class ArgmaxPredictor(PredictionModule):
+    name = "argmax"
+
     def fit(self, context: Context) -> None:
         if context.data_handler.has_oos_samples():
             logger = logging.getLogger(__name__)
diff --git a/autointent/modules/prediction/jinoos.py b/autointent/modules/prediction/jinoos.py
index ca981793..da3004e5 100644
--- a/autointent/modules/prediction/jinoos.py
+++ b/autointent/modules/prediction/jinoos.py
@@ -21,6 +21,7 @@ class JinoosPredictorDumpMetadata(TypedDict):
 
 class JinoosPredictor(PredictionModule):
     metadata_dict_name = "metadata.json"
+    name = "jinoos"
 
     def __init__(self, search_space: list[float] | None = None) -> None:
         self.search_space = np.array(search_space) if search_space is not None else default_search_space
diff --git a/autointent/modules/prediction/threshold.py b/autointent/modules/prediction/threshold.py
index c2cc428b..d4047a22 100644
--- a/autointent/modules/prediction/threshold.py
+++ b/autointent/modules/prediction/threshold.py
@@ -23,6 +23,7 @@ class ThresholdPredictor(PredictionModule):
     metadata_dict_name: str = "metadata.json"
     multilabel: bool
     tags: list[Tag]
+    name = "threshold"
 
     def __init__(self, thresh: float | npt.NDArray[Any]) -> None:
         self.thresh = thresh
diff --git a/autointent/modules/prediction/tunable.py b/autointent/modules/prediction/tunable.py
index 095a7ce4..cf0f9af9 100644
--- a/autointent/modules/prediction/tunable.py
+++ b/autointent/modules/prediction/tunable.py
@@ -24,6 +24,7 @@ class TunablePredictorDumpMetadata(TypedDict):
 
 class TunablePredictor(PredictionModule):
     metadata_dict_name: str = "metadata.json"
+    name = "tunable"
 
     def __init__(self, n_trials: int | None = None) -> None:
         self.n_trials = n_trials
diff --git a/autointent/modules/regexp.py b/autointent/modules/regexp.py
index 09396aab..694c025f 100644
--- a/autointent/modules/regexp.py
+++ b/autointent/modules/regexp.py
@@ -10,6 +10,8 @@
 
 
 class RegExp(Module):
+    name = "regexp"
+
     def fit(self, context: Context) -> None:
         self.regexp_patterns = [
             {
diff --git a/autointent/modules/retrieval/vectordb.py b/autointent/modules/retrieval/vectordb.py
index 07be507e..5c1bb270 100644
--- a/autointent/modules/retrieval/vectordb.py
+++ b/autointent/modules/retrieval/vectordb.py
@@ -11,6 +11,7 @@
 
 class VectorDBModule(RetrievalModule):
     vector_index: VectorIndex
+    name = "vector_db"
 
     def __init__(self, k: int, model_name: str) -> None:
         self.model_name = model_name
diff --git a/autointent/modules/scoring/description/description.py b/autointent/modules/scoring/description/description.py
index bf2edb4c..b4de121e 100644
--- a/autointent/modules/scoring/description/description.py
+++ b/autointent/modules/scoring/description/description.py
@@ -24,6 +24,7 @@ class DescriptionScorer(ScoringModule):
     metadata_dict_name: str = "metadata.json"
     weights_file_name: str = "description_vectors.npy"
     _vector_index: VectorIndex
+    name = "description"
 
     def __init__(self, temperature: float = 1.0) -> None:
         self.temperature = temperature
diff --git a/autointent/modules/scoring/dnnc/dnnc.py b/autointent/modules/scoring/dnnc/dnnc.py
index e53f65b6..5be3554c 100644
--- a/autointent/modules/scoring/dnnc/dnnc.py
+++ b/autointent/modules/scoring/dnnc/dnnc.py
@@ -33,6 +33,8 @@ class DNNCScorer(ScoringModule):
     - inspect batch size of model.predict?
     """
 
+    name = "dnnc"
+
     metadata_dict_name: str = "metadata.json"
     crossencoder_subdir: str = "crossencoder"
     model: CrossEncoder | CrossEncoderWithLogreg
diff --git a/autointent/modules/scoring/knn/knn.py b/autointent/modules/scoring/knn/knn.py
index 27689aa4..2f24d1ef 100644
--- a/autointent/modules/scoring/knn/knn.py
+++ b/autointent/modules/scoring/knn/knn.py
@@ -25,6 +25,7 @@ class KNNScorer(ScoringModule):
     weights: WEIGHT_TYPES
     metadata_dict_name: str = "metadata.json"
     _vector_index: VectorIndex
+    name = "knn"
 
     def __init__(self, k: int, weights: WEIGHT_TYPES) -> None:
         """
diff --git a/autointent/modules/scoring/linear.py b/autointent/modules/scoring/linear.py
index eeb8b9e7..8cea8e9b 100644
--- a/autointent/modules/scoring/linear.py
+++ b/autointent/modules/scoring/linear.py
@@ -39,6 +39,7 @@ class LinearScorer(ScoringModule):
     metadata_dict_name: str = "metadata.json"
     classifier_file_name: str = "classifier.joblib"
     embedding_model_subdir: str = "embedding_model"
+    name = "linear"
 
     def __init__(self, cv: int = 3, n_jobs: int = 1) -> None:
         self.cv = cv
diff --git a/autointent/modules/scoring/mlknn/mlknn.py b/autointent/modules/scoring/mlknn/mlknn.py
index eb98db1f..77912e01 100644
--- a/autointent/modules/scoring/mlknn/mlknn.py
+++ b/autointent/modules/scoring/mlknn/mlknn.py
@@ -26,6 +26,7 @@ class MLKnnScorer(ScoringModule):
     _cond_prob_false: NDArray[np.float64]
     metadata_dict_name: str = "metadata.json"
     arrays_filename: str = "probs.npz"
+    name = "mlknn"
 
     def __init__(self, k: int, s: float = 1.0, ignore_first_neighbours: int = 0) -> None:
         self.k = k
diff --git a/autointent/nodes/nodes_info/__init__.py b/autointent/nodes/nodes_info/__init__.py
index cc0cf4d6..affb3272 100644
--- a/autointent/nodes/nodes_info/__init__.py
+++ b/autointent/nodes/nodes_info/__init__.py
@@ -1,12 +1,14 @@
+from autointent.custom_types import NodeType
+
 from .base import NodeInfo
 from .prediction import PredictionNodeInfo
 from .retrieval import RetrievalNodeInfo
 from .scoring import ScoringNodeInfo
 
 NODES_INFO: dict[str, NodeInfo] = {
-    "retrieval": RetrievalNodeInfo(),
-    "scoring": ScoringNodeInfo(),
-    "prediction": PredictionNodeInfo(),
+    NodeType.retrieval: RetrievalNodeInfo(),
+    NodeType.scoring: ScoringNodeInfo(),
+    NodeType.prediction: PredictionNodeInfo(),
 }
 
 __all__ = ["NodeInfo", "PredictionNodeInfo", "RetrievalNodeInfo", "ScoringNodeInfo", "NODES_INFO"]
diff --git a/autointent/nodes/nodes_info/base.py b/autointent/nodes/nodes_info/base.py
index 411cc061..829bc45b 100644
--- a/autointent/nodes/nodes_info/base.py
+++ b/autointent/nodes/nodes_info/base.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from typing import ClassVar
 
+from autointent.custom_types import NodeType
 from autointent.metrics import METRIC_FN
 from autointent.modules import Module
 
@@ -8,4 +9,4 @@
 class NodeInfo:
     metrics_available: ClassVar[Mapping[str, METRIC_FN]]
     modules_available: ClassVar[Mapping[str, type[Module]]]
-    node_type: str
+    node_type: NodeType
diff --git a/autointent/nodes/nodes_info/prediction.py b/autointent/nodes/nodes_info/prediction.py
index b68feeb8..0ebccb05 100644
--- a/autointent/nodes/nodes_info/prediction.py
+++ b/autointent/nodes/nodes_info/prediction.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from typing import ClassVar
 
+from autointent.custom_types import NodeType
 from autointent.metrics import PREDICTION_METRICS_MULTICLASS, PREDICTION_METRICS_MULTILABEL, PredictionMetricFn
 from autointent.modules import PREDICTION_MODULES_MULTICLASS, PREDICTION_MODULES_MULTILABEL, Module
 
@@ -14,4 +15,4 @@ class PredictionNodeInfo(NodeInfo):
 
     modules_available: ClassVar[dict[str, type[Module]]] = PREDICTION_MODULES_MULTICLASS | PREDICTION_MODULES_MULTILABEL
 
-    node_type = "prediction"
+    node_type = NodeType.prediction
diff --git a/autointent/nodes/nodes_info/regexp.py b/autointent/nodes/nodes_info/regexp.py
index 2b292fa9..411a778b 100644
--- a/autointent/nodes/nodes_info/regexp.py
+++ b/autointent/nodes/nodes_info/regexp.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from typing import ClassVar
 
+from autointent.custom_types import NodeType
 from autointent.metrics import regexp_partial_accuracy, regexp_partial_precision
 from autointent.metrics.regexp import RegexpMetricFn
 from autointent.modules import Module, RegExp
@@ -15,6 +16,6 @@ class RegExpNode(NodeInfo):
         regexp_partial_precision,
     )
 
-    modules_available: ClassVar[Mapping[str, type[Module]]] = {"regexp": RegExp}  # type: ignore[type-abstract]
+    modules_available: ClassVar[Mapping[str, type[Module]]] = {NodeType.regexp: RegExp}  # type: ignore[type-abstract]
 
-    node_type = "regexp"
+    node_type = NodeType.regexp
diff --git a/autointent/nodes/nodes_info/retrieval.py b/autointent/nodes/nodes_info/retrieval.py
index 1f2e76cc..401da5ab 100644
--- a/autointent/nodes/nodes_info/retrieval.py
+++ b/autointent/nodes/nodes_info/retrieval.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from typing import ClassVar
 
+from autointent.custom_types import NodeType
 from autointent.metrics import (
     RETRIEVAL_METRICS_MULTICLASS,
     RETRIEVAL_METRICS_MULTILABEL,
@@ -20,4 +21,4 @@ class RetrievalNodeInfo(NodeInfo):
         RETRIEVAL_MODULES_MULTICLASS | RETRIEVAL_MODULES_MULTILABEL
     )
 
-    node_type = "retrieval"
+    node_type = NodeType.retrieval
diff --git a/autointent/nodes/nodes_info/scoring.py b/autointent/nodes/nodes_info/scoring.py
index 256e41d1..c31872fd 100644
--- a/autointent/nodes/nodes_info/scoring.py
+++ b/autointent/nodes/nodes_info/scoring.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from typing import ClassVar
 
+from autointent.custom_types import NodeType
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL, ScoringMetricFn
 from autointent.modules import SCORING_MODULES_MULTICLASS, SCORING_MODULES_MULTILABEL, ScoringModule
 
@@ -14,4 +15,4 @@ class ScoringNodeInfo(NodeInfo):
         SCORING_MODULES_MULTICLASS | SCORING_MODULES_MULTILABEL
     )
 
-    node_type = "scoring"
+    node_type = NodeType.scoring
diff --git a/autointent/pipeline/inference/inference_pipeline.py b/autointent/pipeline/inference/inference_pipeline.py
index db135b4a..6a95e552 100644
--- a/autointent/pipeline/inference/inference_pipeline.py
+++ b/autointent/pipeline/inference/inference_pipeline.py
@@ -3,7 +3,7 @@
 from hydra.utils import instantiate
 
 from autointent.configs.inference_pipeline import InferencePipelineConfig
-from autointent.custom_types import LABEL_TYPE
+from autointent.custom_types import LABEL_TYPE, NodeType
 from autointent.nodes.inference import InferenceNode
 
 
@@ -16,5 +16,5 @@ def from_dict_config(cls, config: dict[str, Any]) -> "InferencePipeline":
         return instantiate(InferencePipelineConfig, **config)  # type: ignore[no-any-return]
 
     def predict(self, utterances: list[str]) -> list[LABEL_TYPE]:
-        scores = self.nodes["scoring"].module.predict(utterances)
-        return self.nodes["prediction"].module.predict(scores)  # type: ignore[return-value]
+        scores = self.nodes[NodeType.scoring].module.predict(utterances)
+        return self.nodes[NodeType.prediction].module.predict(scores)  # type: ignore[return-value]
diff --git a/autointent/pipeline/optimization/pipeline_optimizer.py b/autointent/pipeline/optimization/pipeline_optimizer.py
index d0532164..dc632de4 100644
--- a/autointent/pipeline/optimization/pipeline_optimizer.py
+++ b/autointent/pipeline/optimization/pipeline_optimizer.py
@@ -12,6 +12,7 @@
 from autointent.nodes import NodeOptimizer
 
 from .utils import NumpyEncoder
+from autointent.custom_types import NodeType
 
 
 class PipelineOptimizer:
@@ -70,7 +71,7 @@ def dump(self, logs_dir: str | Path | None) -> None:
             yaml.dump(inference_config, file)
 
 
-def make_report(logs: dict[str, Any], nodes: list[str]) -> str:
+def make_report(logs: dict[str, Any], nodes: list[NodeType]) -> str:
     ids = [np.argmax(logs["metrics"][node]) for node in nodes]
     configs = []
     for i, node in zip(ids, nodes, strict=False):
diff --git a/pyproject.toml b/pyproject.toml
index 95f78baa..8cc3977a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -138,5 +138,6 @@ module = [
     "hydra.*",
     "transformers",
     "faiss",
+    "joblib",
 ]
 ignore_missing_imports = true

From 52d1f492f6a0d8ee4a355a5cc4d82eb4b174cea4 Mon Sep 17 00:00:00 2001
From: Roman Solomatin <36135455+Samoed@users.noreply.github.com>
Date: Tue, 5 Nov 2024 21:01:17 +0300
Subject: [PATCH 2/2] lint

---
 autointent/modules/prediction/argmax.py                | 2 +-
 autointent/modules/scoring/description/description.py  | 3 +--
 autointent/modules/scoring/dnnc/dnnc.py                | 1 -
 autointent/pipeline/inference/inference_pipeline.py    | 3 +--
 autointent/pipeline/optimization/pipeline_optimizer.py | 2 +-
 5 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/autointent/modules/prediction/argmax.py b/autointent/modules/prediction/argmax.py
index c4502962..5b76a665 100644
--- a/autointent/modules/prediction/argmax.py
+++ b/autointent/modules/prediction/argmax.py
@@ -13,10 +13,10 @@
 
 class ArgmaxPredictor(PredictionModule):
     metadata = {}  # noqa: RUF012
+    name = "argmax"
 
     def __init__(self) -> None:
         pass
-    name = "argmax"
 
 
     @classmethod
diff --git a/autointent/modules/scoring/description/description.py b/autointent/modules/scoring/description/description.py
index 96ce7517..b8a7b8a3 100644
--- a/autointent/modules/scoring/description/description.py
+++ b/autointent/modules/scoring/description/description.py
@@ -10,8 +10,7 @@
 
 from autointent.context import Context
 from autointent.context.embedder import Embedder
-from autointent.context.vector_index_client import VectorIndex
-from autointent.context.vector_index_client import VectorIndexClient
+from autointent.context.vector_index_client import VectorIndex, VectorIndexClient
 from autointent.context.vector_index_client.cache import get_db_dir
 from autointent.custom_types import LabelType
 from autointent.modules.scoring.base import ScoringModule
diff --git a/autointent/modules/scoring/dnnc/dnnc.py b/autointent/modules/scoring/dnnc/dnnc.py
index 3f8addc8..69e90f0a 100644
--- a/autointent/modules/scoring/dnnc/dnnc.py
+++ b/autointent/modules/scoring/dnnc/dnnc.py
@@ -37,7 +37,6 @@ class DNNCScorer(ScoringModule):
 
     name = "dnnc"
 
-    metadata_dict_name: str = "metadata.json"
     crossencoder_subdir: str = "crossencoder"
     model: CrossEncoder | CrossEncoderWithLogreg
     prebuilt_index: bool = False
diff --git a/autointent/pipeline/inference/inference_pipeline.py b/autointent/pipeline/inference/inference_pipeline.py
index df40f958..c17e24f3 100644
--- a/autointent/pipeline/inference/inference_pipeline.py
+++ b/autointent/pipeline/inference/inference_pipeline.py
@@ -3,8 +3,7 @@
 from hydra.utils import instantiate
 
 from autointent.configs.inference_pipeline import InferencePipelineConfig
-from autointent.custom_types import LabelType
-from autointent.custom_types import NodeType
+from autointent.custom_types import LabelType, NodeType
 from autointent.nodes.inference import InferenceNode
 
 
diff --git a/autointent/pipeline/optimization/pipeline_optimizer.py b/autointent/pipeline/optimization/pipeline_optimizer.py
index dc632de4..396e174e 100644
--- a/autointent/pipeline/optimization/pipeline_optimizer.py
+++ b/autointent/pipeline/optimization/pipeline_optimizer.py
@@ -9,10 +9,10 @@
 
 from autointent import Context
 from autointent.configs.pipeline_optimizer import PipelineOptimizerConfig
+from autointent.custom_types import NodeType
 from autointent.nodes import NodeOptimizer
 
 from .utils import NumpyEncoder
-from autointent.custom_types import NodeType
 
 
 class PipelineOptimizer: