Feat/optuna

autointent/_pipeline/_pipeline.py

@@ -10,10 +10,9 @@
 
 from autointent import Context, Dataset
 from autointent.configs import CrossEncoderConfig, EmbedderConfig, InferenceNodeConfig, LoggingConfig, VectorIndexConfig
-from autointent.custom_types import ListOfGenericLabels, NodeType, ValidationScheme
+from autointent.custom_types import ListOfGenericLabels, NodeType, TuningType, ValidationScheme
 from autointent.metrics import PREDICTION_METRICS_MULTILABEL
 from autointent.nodes import InferenceNode, NodeOptimizer
-from autointent.nodes.schemes import OptimizationConfig
 from autointent.utils import load_default_search_space, load_search_space
 
 from ._schemas import InferencePipelineOutput, InferencePipelineUtteranceOutput
@@ -77,7 +76,7 @@ def from_search_space(cls, search_space: list[dict[str, Any]] | Path | str, seed
         """
         if isinstance(search_space, Path | str):
             search_space = load_search_space(search_space)
-        validated_search_space = OptimizationConfig(search_space).model_dump()  # type: ignore[arg-type]
+        validated_search_space = search_space  # OptimizationConfig(search_space).model_dump()  # type: ignore[arg-type]
         nodes = [NodeOptimizer(**node) for node in validated_search_space]
         return cls(nodes=nodes, seed=seed)
 
@@ -93,7 +92,7 @@ def default_optimizer(cls, multilabel: bool, seed: int = 42) -> "Pipeline":
         """
         return cls.from_search_space(search_space=load_default_search_space(multilabel), seed=seed)
 
-    def _fit(self, context: Context) -> None:
+    def _fit(self, context: Context, tuning: TuningType = "brute") -> None:
         """
         Optimize the pipeline.
 
@@ -108,7 +107,7 @@ def _fit(self, context: Context) -> None:
         for node_type in NodeType:
             node_optimizer = self.nodes.get(node_type, None)
             if node_optimizer is not None:
-                node_optimizer.fit(context)  # type: ignore[union-attr]
+                node_optimizer.fit(context, tuning)  # type: ignore[union-attr]
         if not context.vector_index_config.save_db:
             self._logger.info("removing vector database from file system...")
             # TODO clear cache from appdirs
@@ -123,7 +122,12 @@ def _is_inference(self) -> bool:
         return isinstance(self.nodes[NodeType.scoring], InferenceNode)
 
     def fit(
-        self, dataset: Dataset, scheme: ValidationScheme = "ho", n_folds: int = 3, refit_after: bool = False
+        self,
+        dataset: Dataset,
+        scheme: ValidationScheme = "ho",
+        n_folds: int = 3,
+        refit_after: bool = False,
+        tuning: TuningType = "brute",
     ) -> Context:
         """
         Optimize the pipeline from dataset.
@@ -141,7 +145,7 @@ def fit(
         context.configure_vector_index(self.vector_index_config, self.embedder_config)
         context.configure_cross_encoder(self.cross_encoder_config)
         self.validate_modules(dataset)
-        self._fit(context)
+        self._fit(context, tuning)
 
         if context.is_ram_to_clear():
             nodes_configs = context.optimization_info.get_inference_nodes_config()

autointent/configs/_optimization.py

@@ -4,7 +4,7 @@
 
 from pydantic import BaseModel, Field
 
-from autointent.custom_types import ValidationScheme
+from autointent.custom_types import TuningType, ValidationScheme
 
 from ._name import get_run_name
 
@@ -25,6 +25,7 @@ class TaskConfig(BaseModel):
 
     search_space_path: Path | None = None
     """Path to the search space configuration file. If None, the default search space will be used"""
+    sampler: TuningType = "brute"
 
 
 class LoggingConfig(BaseModel):

autointent/custom_types.py

@@ -71,4 +71,5 @@ class Split:
     INTENTS = "intents"
 
 
+TuningType = Literal["brute", "bayes"]
 ValidationScheme = Literal["ho", "cv"]

autointent/nodes/_optimization/_node_optimizer.py

@@ -4,17 +4,43 @@
 import itertools as it
 import logging
 from copy import deepcopy
+from functools import partial
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal, TypedDict
 
+import optuna
 import torch
+from optuna.trial import Trial
 
 from autointent import Dataset
 from autointent.context import Context
-from autointent.custom_types import NodeType
+from autointent.custom_types import NodeType, TuningType
 from autointent.nodes._nodes_info import NODES_INFO
 
 
+class ParamSpaceCat(TypedDict):
+    choices: list[Any]
+
+
+class ParamSpaceInt(TypedDict, total=False):
+    low: int
+    high: int
+    step: int
+    log: bool
+
+
+class ParamSpaceFloat(TypedDict, total=False):
+    low: float
+    high: float
+    step: float
+    log: bool
+
+
+class ParamSpace(TypedDict):
+    type: Literal["cat", "int", "float"]
+    content: ParamSpaceCat | ParamSpaceInt | ParamSpaceFloat
+
+
 class NodeOptimizer:
     """Node optimizer class."""
 
@@ -43,30 +69,42 @@ def __init__(
         self.modules_search_spaces = search_space
         self._logger = logging.getLogger(__name__)  # TODO solve duplicate logging messages problem
 
-    def fit(self, context: Context) -> None:
+    def fit(self, context: Context, tuning: TuningType = "brute") -> None:
         """
         Fit the node optimizer.
 
         :param context: Context
         """
         self._logger.info("starting %s node optimization...", self.node_info.node_type)
 
+        if tuning == "brute":
+            self._fit_brute(context)
+        elif tuning == "bayes":
+            self._fit_bayes(context)
+        else:
+            msg = f"Unexepected tuning type: {tuning}"
+            raise ValueError(msg)
+
+        self._logger.info("%s node optimization is finished!", self.node_info.node_type)
+
+    def _fit_brute(self, context: Context) -> None:
         for search_space in deepcopy(self.modules_search_spaces):
             module_name = search_space.pop("module_name")
 
             for j_combination, params_combination in enumerate(it.product(*search_space.values())):
                 module_kwargs = dict(zip(search_space.keys(), params_combination, strict=False))
 
                 self._logger.debug("initializing %s module...", module_name)
-                context.callback_handler.start_module(
-                    module_name=module_name, num=j_combination, module_kwargs=module_kwargs
-                )
                 module = self.node_info.modules_available[module_name].from_context(context, **module_kwargs)
 
                 embedder_name = module.get_embedder_name()
                 if embedder_name is not None:
                     module_kwargs["embedder_name"] = embedder_name
 
+                context.callback_handler.start_module(
+                    module_name=module_name, num=j_combination, module_kwargs=module_kwargs
+                )
+
                 self._logger.debug("scoring %s module...", module_name)
                 metrics_score = module.score(context, metrics=self.metrics)
                 metric_value = metrics_score[self.target_metric]
@@ -98,7 +136,79 @@ def fit(self, context: Context) -> None:
                     gc.collect()
                     torch.cuda.empty_cache()
 
-        self._logger.info("%s node optimization is finished!", self.node_info.node_type)
+    def _fit_bayes(self, context: Context, seed: int = 42, n_trials: int = 10) -> None:
+        for search_space in deepcopy(self.modules_search_spaces):
+            self._counter = 0
+            study = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=seed))
+            optuna.logging.set_verbosity(optuna.logging.WARNING)
+            module_name = search_space.pop("module_name")
+            obj = partial(self.objective, module_name=module_name, search_space=search_space, context=context)
+            study.optimize(obj, n_trials=n_trials)
+
+    def objective(
+        self, trial: Trial, module_name: str, search_space: dict[str, ParamSpace | list[Any]], context: Context
+    ) -> float:
+        config = self.suggest(trial, search_space)
+
+        self._logger.debug("initializing %s module...", module_name)
+        module = self.node_info.modules_available[module_name].from_context(context, **config)
+
+        embedder_name = module.get_embedder_name()
+        if embedder_name is not None:
+            config["embedder_name"] = embedder_name
+
+        context.callback_handler.start_module(module_name=module_name, num=self._counter, module_kwargs=config)
+
+        self._logger.debug("scoring %s module...", module_name)
+        all_metrics = module.score(context, metrics=self.metrics)
+        target_metric = all_metrics[self.target_metric]
+
+        context.callback_handler.log_metrics(all_metrics)
+        context.callback_handler.end_module()
+
+        dump_dir = context.get_dump_dir()
+
+        if dump_dir is not None:
+            module_dump_dir = self.get_module_dump_dir(dump_dir, module_name, self._counter)
+            module.dump(module_dump_dir)
+        else:
+            module_dump_dir = None
+
+        context.optimization_info.log_module_optimization(
+            self.node_info.node_type,
+            module_name,
+            config,
+            target_metric,
+            self.target_metric,
+            module.get_assets(),  # retriever name / scores / predictions
+            module_dump_dir,
+            module=module if not context.is_ram_to_clear() else None,
+        )
+
+        if context.is_ram_to_clear():
+            module.clear_cache()
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        self._counter += 1
+
+        return target_metric
+
+    def suggest(self, trial: Trial, search_space: dict[str, ParamSpace | list[Any]]) -> dict[str, Any]:
+        res: dict[str, Any] = {}
+        for param_name, param_space in search_space.items():
+            if isinstance(param_space, list):
+                res[param_name] = trial.suggest_categorical(param_name, choices=param_space)
+            elif param_space["type"] == "cat":
+                res[param_name] = trial.suggest_categorical(param_name, **param_space["content"])
+            elif param_space["type"] == "int":
+                res[param_name] = trial.suggest_int(param_name, **param_space["content"])
+            elif param_space["type"] == "float":
+                res[param_name] = trial.suggest_float(param_name, **param_space["content"])
+            else:
+                msg = f"Unsupported type of param search space: {param_space}"
+                raise TypeError(msg)
+        return res
 
     def get_module_dump_dir(self, dump_dir: Path, module_name: str, j_combination: int) -> str:
         """

tests/assets/configs/bayes.yaml

@@ -0,0 +1,32 @@
+- node_type: embedding
+  target_metric: retrieval_hit_rate
+  search_space:
+    - module_name: retrieval
+      k: [10]
+      embedder_name:
+        - sentence-transformers/all-MiniLM-L6-v2
+        - avsolatorio/GIST-small-Embedding-v0
+- node_type: scoring
+  target_metric: scoring_roc_auc
+  search_space:
+    - module_name: knn
+      k:
+        type: "int"
+        content:
+          low: 5
+          high: 10
+          step: 1
+      weights: [uniform, distance, closest]
+    - module_name: linear
+- node_type: decision
+  target_metric: decision_accuracy
+  search_space:
+    - module_name: threshold
+      thresh:
+        type: float
+        content:
+          low: 0.1
+          high: 0.9
+    - module_name: tunable
+    - module_name: argmax
+    - module_name: jinoos

tests/conftest.py

@@ -27,7 +27,7 @@ def dataset_unsplitted():
     return Dataset.from_json(path)
 
 
-TaskType = Literal["multiclass", "multilabel", "description"]
+TaskType = Literal["multiclass", "multilabel", "description", "bayes"]
 
 
 def get_search_space_path(task_type: TaskType):

tests/pipeline/test_optimization.py

@@ -1,5 +1,4 @@
 import os
-from typing import Literal
 
 import pytest
 
@@ -11,7 +10,18 @@
 )
 from tests.conftest import get_search_space, setup_environment
 
-TaskType = Literal["multiclass", "multilabel", "description"]
+
+def test_bayes(dataset):
+    project_dir = setup_environment()
+    search_space = get_search_space("bayes")
+
+    pipeline_optimizer = Pipeline.from_search_space(search_space)
+
+    pipeline_optimizer.set_config(LoggingConfig(project_dir=project_dir, dump_modules=True, clear_ram=True))
+    pipeline_optimizer.set_config(VectorIndexConfig())
+    pipeline_optimizer.set_config(EmbedderConfig(batch_size=16, max_length=32, device="cpu"))
+
+    pipeline_optimizer.fit(dataset, scheme="cv", refit_after=True, tuning="bayes")
 
 
 @pytest.mark.parametrize(