Refactor/only one pipeline should live (#62)

Author: voorhs

autointent/__init__.py

@@ -1,6 +1,6 @@
 from ._embedder import Embedder
 from .context import Context
 from .context.data_handler import Dataset
-from .pipeline import InferencePipeline, PipelineOptimizer
+from .pipeline import Pipeline
 
-__all__ = ["Context", "Dataset", "Embedder", "InferencePipeline", "PipelineOptimizer"]
+__all__ = ["Context", "Dataset", "Embedder", "Pipeline"]

autointent/configs/__init__.py

@@ -1,5 +1,3 @@
-from ._inference_cli import InferenceConfig
-from ._inference_pipeline import InferencePipelineConfig
 from ._node import InferenceNodeConfig, NodeOptimizerConfig
 from ._optimization_cli import (
     AugmentationConfig,
@@ -10,19 +8,15 @@
     TaskConfig,
     VectorIndexConfig,
 )
-from ._pipeline_optimizer import PipelineOptimizerConfig
 
 __all__ = [
     "AugmentationConfig",
     "DataConfig",
     "EmbedderConfig",
-    "InferenceConfig",
     "InferenceNodeConfig",
-    "InferencePipelineConfig",
     "LoggingConfig",
     "NodeOptimizerConfig",
     "OptimizationConfig",
-    "PipelineOptimizerConfig",
     "TaskConfig",
     "VectorIndexConfig",
 ]

autointent/configs/_inference_cli.py

@@ -5,14 +5,14 @@
 
 from omegaconf import MISSING
 
-from autointent.custom_types import NodeType, NodeTypeType
+from autointent.custom_types import NodeType
 
 
 @dataclass
 class InferenceNodeConfig:
     """Configuration for the inference node."""
 
-    node_type: NodeTypeType = MISSING
+    node_type: NodeType = MISSING
     """Type of the node. Should be one of the NODE_TYPES"""
     module_type: str = MISSING  # TODO: add custom type
     """Type of the module. Should be one of the Module"""

autointent/configs/_inference_pipeline.py

@@ -11,7 +11,7 @@
 from numpy.typing import NDArray
 
 from autointent.configs import InferenceNodeConfig
-from autointent.custom_types import NODE_TYPES, NodeType, NodeTypeType
+from autointent.custom_types import NodeType
 
 from ._data_models import Artifact, Artifacts, RetrieverArtifact, ScorerArtifact, Trial, Trials, TrialsIds
 from ._logger import get_logger
@@ -188,13 +188,13 @@ def get_best_oos_scores(self, split: Literal["train", "validation", "test"]) ->
             return best_scorer_artifact.oos_scores[split]
         return best_scorer_artifact.oos_scores
 
-    def dump_evaluation_results(self) -> dict[str, dict[str, list[float]]]:
+    def dump_evaluation_results(self) -> dict[str, Any]:
         """
         Dump evaluation results for all nodes.
 
         :return: Dictionary containing metrics and configurations for all nodes.
         """
-        node_wise_metrics = {node_type: self._get_metrics_values(node_type) for node_type in NODE_TYPES}
+        node_wise_metrics = {node_type: self._get_metrics_values(node_type) for node_type in NodeType}
         return {
             "metrics": node_wise_metrics,
             "configs": self.trials.model_dump(),
@@ -206,15 +206,15 @@ def get_inference_nodes_config(self) -> list[InferenceNodeConfig]:
 
         :return: List of `InferenceNodeConfig` objects for inference nodes.
         """
-        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 NodeType]
         res = []
-        for idx, node_type in zip(trial_ids, NODE_TYPES, strict=True):
+        for idx, node_type in zip(trial_ids, NodeType, strict=True):
             if idx is None:
                 continue
             trial = self.trials.get_trial(node_type, idx)
             res.append(
                 InferenceNodeConfig(
-                    node_type=node_type,  # type: ignore[arg-type]
+                    node_type=node_type,
                     module_type=trial.module_type,
                     module_config=trial.module_params,
                     load_path=trial.module_dump_dir,
@@ -234,11 +234,11 @@ def _get_best_module(self, node_type: str) -> "Module | None":
             return self.modules.get(node_type)[idx]
         return None
 
-    def get_best_modules(self) -> dict[NodeTypeType, "Module"]:
+    def get_best_modules(self) -> dict[NodeType, "Module"]:
         """
         Retrieve the best modules for all node types.
 
         :return: Dictionary of the best modules for each node type.
         """
-        res = {nt: self._get_best_module(nt) for nt in NODE_TYPES}
-        return {nt: m for nt, m in res.items() if m is not None}  # type: ignore[misc]
+        res = {nt: self._get_best_module(nt) for nt in NodeType}
+        return {nt: m for nt, m in res.items() if m is not None}

autointent/configs/_node.py

@@ -52,11 +52,7 @@ class BaseMetadataDict(TypedDict):
 class NodeType(str, Enum):
     """Enumeration of node types in the AutoIntent pipeline."""
 
+    regexp = "regexp"
     retrieval = "retrieval"
-    prediction = "prediction"
     scoring = "scoring"
-    regexp = "regexp"
-
-
-NODE_TYPES = [NodeType.retrieval.value, NodeType.prediction.value, NodeType.scoring.value, NodeType.regexp.value]
-NodeTypeType = Literal["retrieval", "prediction", "scoring", "regexp"]
+    prediction = "prediction"

autointent/configs/_pipeline_optimizer.py

@@ -1,28 +1,27 @@
 # TODO: make up a better and more versatile config
-nodes:
-  - node_type: retrieval
-    metric: retrieval_hit_rate
-    search_space:
-      - module_type: vector_db
-        k: [10]
-        embedder_name:
-          - avsolatorio/GIST-small-Embedding-v0
-          - infgrad/stella-base-en-v2
-  - node_type: scoring
-    metric: scoring_roc_auc
-    search_space:
-      - module_type: knn
-        k: [1, 3, 5, 10]
-        weights: ["uniform", "distance", "closest"]
-      - module_type: linear
-      - module_type: dnnc
-        cross_encoder_name:
-          - BAAI/bge-reranker-base
-          - cross-encoder/ms-marco-MiniLM-L-6-v2
-        k: [1, 3, 5, 10]
-  - node_type: prediction
-    metric: prediction_accuracy
-    search_space:
-      - module_type: threshold
-        thresh: [0.5]
-      - module_type: argmax
+- node_type: retrieval
+  metric: retrieval_hit_rate
+  search_space:
+    - module_type: vector_db
+      k: [10]
+      embedder_name:
+        - avsolatorio/GIST-small-Embedding-v0
+        - infgrad/stella-base-en-v2
+- node_type: scoring
+  metric: scoring_roc_auc
+  search_space:
+    - module_type: knn
+      k: [1, 3, 5, 10]
+      weights: ["uniform", "distance", "closest"]
+    - module_type: linear
+    - module_type: dnnc
+      cross_encoder_name:
+        - BAAI/bge-reranker-base
+        - cross-encoder/ms-marco-MiniLM-L-6-v2
+      k: [1, 3, 5, 10]
+- node_type: prediction
+  metric: prediction_accuracy
+  search_space:
+    - module_type: threshold
+      thresh: [0.5]
+    - module_type: argmax

autointent/context/optimization_info/_optimization_info.py

@@ -1,22 +1,21 @@
 # TODO: make up a better and more versatile config
-nodes:
-  - node_type: retrieval
-    metric: retrieval_hit_rate_intersecting
-    search_space:
-      - module_type: vector_db
-        k: [10]
-        embedder_name:
-          - deepvk/USER-bge-m3
-  - node_type: scoring
-    metric: scoring_roc_auc
-    search_space:
-      - module_type: knn
-        k: [3]
-        weights: ["uniform", "distance", "closest"]
-      - module_type: linear
-  - node_type: prediction
-    metric: prediction_accuracy
-    search_space:
-      - module_type: threshold
-        thresh: [0.5]
-      - module_type: adaptive
+- node_type: retrieval
+  metric: retrieval_hit_rate_intersecting
+  search_space:
+    - module_type: vector_db
+      k: [10]
+      embedder_name:
+        - deepvk/USER-bge-m3
+- node_type: scoring
+  metric: scoring_roc_auc
+  search_space:
+    - module_type: knn
+      k: [3]
+      weights: ["uniform", "distance", "closest"]
+    - module_type: linear
+- node_type: prediction
+  metric: prediction_accuracy
+  search_space:
+    - module_type: threshold
+      thresh: [0.5]
+    - module_type: adaptive