Feat/optuna (#128)

* define interface

* basic ho iterator

* move obtaining data for train from node optimizer to modules themselves

* stage progress

* implement cv iterator

* minor bug fix

* implement cv iterator for decision node

* move cv iteration to base module definition

* implement cv iterator for embedding node

* add training to `score_ho` of each node

* properly define base module

* fix codestyle

* remove regexp node

* remove regexp validator

* fix typing problems (except `DataHandler._split_cv`)

* add ingore oos decorator

* fix codestyle

* fix typing

* add oos handling to cv iterator

* remove `DataHandler.dump()`

* minor bug fix

* implement splitting to cv folds

* fix codestyle

* remove regex tests

* bug fix

* bug fix

* update tests

* fix typing

* big fix

* basic test on cv folding

* add tests for metrics to ignore oos samples

* add tests for cv iterator

* fix codestyle

* minor bug fix

* fix codestyle

* add test for cv

* bug fix

* implement cv iterator for description scorer

* refactor cv iterator for description node

* fix typing

* add cache cleaning before refitting

* bug fix

* implement refitting the whole pipeline with all train data

* fix typing

* bug fix

* fix typing

* respond to samoed

* create `ValidationType` in `autointent.custom_types`

* fix docstring

* properly expose `n_folds` argument

* implement `_fit_bayes`

* add typing to param spaces

* minor bug fix

* minor bug fix

* fix codestyle

* add tuning selection to pipeline

* add test on bayes

* disable search space validation for now

* fix codestyle

* remove `ParamSpaceCat` (it's redundant)

* move to optuna entirely

* refactor yaml format a little bit

* add test for random sampler

* rename some variables

* add config validation for optuna (#132)

* add config validation

* add validation for union types

* remove debug code

* remove comment

* run tests on pr for all branches

* fix mlknn

* fix type validation

* return CI config back to normal

* fix default value for step in `ParamSpaceFloat`

* update schema

* update callback test

* change CI config

* update search space configs for testing

* enable validation back

* remove TunableDecision from search spaces

* upd schema

---------

Co-authored-by: Roman Solomatin <[email protected]>
Co-authored-by: Roman Solomatin <[email protected]>

Author: voorhs

.github/workflows/test-inference.yaml

@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:

.github/workflows/test-nodes.yaml

@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:

.github/workflows/test-optimization.yaml

@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:

.github/workflows/unit-tests.yaml

@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:

autointent/_pipeline/_pipeline.py

@@ -10,7 +10,7 @@
 
 from autointent import Context, Dataset
 from autointent.configs import InferenceNodeConfig, LoggingConfig, VectorIndexConfig
-from autointent.custom_types import ListOfGenericLabels, NodeType, ValidationScheme
+from autointent.custom_types import ListOfGenericLabels, NodeType, SamplerType, ValidationScheme
 from autointent.metrics import PREDICTION_METRICS_MULTILABEL
 from autointent.nodes import InferenceNode, NodeOptimizer
 from autointent.nodes.schemes import OptimizationConfig
@@ -87,7 +87,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, sampler: SamplerType = "brute") -> None:
         """
         Optimize the pipeline.
 
@@ -102,7 +102,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, sampler)  # 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
@@ -117,7 +117,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,
+        sampler: SamplerType = "brute",
     ) -> Context:
         """
         Optimize the pipeline from dataset.
@@ -135,7 +140,7 @@ def fit(
         context.configure_vector_index(self.vector_index_config)
 
         self.validate_modules(dataset)
-        self._fit(context)
+        self._fit(context, sampler)
 
         if context.is_ram_to_clear():
             nodes_configs = context.optimization_info.get_inference_nodes_config()

autointent/configs/_optimization.py

@@ -2,9 +2,9 @@
 
 from pathlib import Path
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PositiveInt
 
-from autointent.custom_types import ValidationScheme
+from autointent.custom_types import SamplerType, ValidationScheme
 
 from ._name import get_run_name
 
@@ -16,7 +16,7 @@ class DataConfig(BaseModel):
     """Path to the training data. Can be local path or HF repo."""
     scheme: ValidationScheme
     """Hold-out or cross-validation."""
-    n_folds: int = 3
+    n_folds: PositiveInt = 3
     """Number of folds in cross-validation."""
 
 
@@ -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: SamplerType = "brute"
 
 
 class LoggingConfig(BaseModel):

autointent/custom_types.py

@@ -5,7 +5,9 @@
 """
 
 from enum import Enum
-from typing import Literal, TypeAlias
+from typing import Annotated, Literal, TypeAlias
+
+from annotated_types import Interval
 
 
 class LogLevel(Enum):
@@ -71,4 +73,9 @@ class Split:
     INTENTS = "intents"
 
 
+SamplerType = Literal["brute", "tpe", "random"]
 ValidationScheme = Literal["ho", "cv"]
+
+
+FloatFromZeroToOne = Annotated[float, Interval(ge=0, le=1)]
+"""Float value between 0 and 1, inclusive."""

autointent/modules/decision/_adaptive.py

@@ -7,7 +7,7 @@
 import numpy.typing as npt
 
 from autointent import Context
-from autointent.custom_types import ListOfGenericLabels, ListOfLabelsWithOOS, MultiLabel
+from autointent.custom_types import FloatFromZeroToOne, ListOfGenericLabels, ListOfLabelsWithOOS, MultiLabel
 from autointent.exceptions import MismatchNumClassesError
 from autointent.metrics import decision_f1
 from autointent.modules.abc import DecisionModule
@@ -58,7 +58,7 @@ class AdaptiveDecision(DecisionModule):
     supports_oos = False
     name = "adaptive"
 
-    def __init__(self, search_space: list[float] | None = None) -> None:
+    def __init__(self, search_space: list[FloatFromZeroToOne] | None = None) -> None:
         """
         Initialize the AdaptiveDecision.
 
@@ -68,7 +68,7 @@ def __init__(self, search_space: list[float] | None = None) -> None:
         self.search_space = search_space if search_space is not None else default_search_space
 
     @classmethod
-    def from_context(cls, context: Context, search_space: list[float] | None = None) -> "AdaptiveDecision":
+    def from_context(cls, context: Context, search_space: list[FloatFromZeroToOne] | None = None) -> "AdaptiveDecision":
         """
         Create an AdaptiveDecision instance using a Context object.
 

autointent/modules/decision/_jinoos.py

@@ -6,7 +6,7 @@
 import numpy.typing as npt
 
 from autointent import Context
-from autointent.custom_types import ListOfGenericLabels
+from autointent.custom_types import FloatFromZeroToOne, ListOfGenericLabels
 from autointent.exceptions import MismatchNumClassesError
 from autointent.modules.abc import DecisionModule
 from autointent.schemas import Tag
@@ -55,7 +55,7 @@ class JinoosDecision(DecisionModule):
 
     def __init__(
         self,
-        search_space: list[float] | None = None,
+        search_space: list[FloatFromZeroToOne] | None = None,
     ) -> None:
         """
         Initialize Jinoos predictor.
@@ -65,7 +65,7 @@ def __init__(
         self.search_space = np.array(search_space) if search_space is not None else default_search_space
 
     @classmethod
-    def from_context(cls, context: Context, search_space: list[float] | None = None) -> "JinoosDecision":
+    def from_context(cls, context: Context, search_space: list[FloatFromZeroToOne] | None = None) -> "JinoosDecision":
         """
         Initialize from context.
 

autointent/modules/decision/_threshold.py

@@ -7,7 +7,7 @@
 import numpy.typing as npt
 
 from autointent import Context
-from autointent.custom_types import ListOfGenericLabels, MultiLabel
+from autointent.custom_types import FloatFromZeroToOne, ListOfGenericLabels, MultiLabel
 from autointent.exceptions import MismatchNumClassesError
 from autointent.modules.abc import DecisionModule
 from autointent.schemas import Tag
@@ -75,7 +75,7 @@ class ThresholdDecision(DecisionModule):
 
     def __init__(
         self,
-        thresh: float | list[float],
+        thresh: FloatFromZeroToOne | list[FloatFromZeroToOne],
     ) -> None:
         """
         Initialize threshold predictor.
@@ -85,7 +85,9 @@ def __init__(
         self.thresh = thresh if isinstance(thresh, float) else np.array(thresh)
 
     @classmethod
-    def from_context(cls, context: Context, thresh: float | list[float] = 0.5) -> "ThresholdDecision":
+    def from_context(
+        cls, context: Context, thresh: FloatFromZeroToOne | list[FloatFromZeroToOne] = 0.5
+    ) -> "ThresholdDecision":
         """
         Initialize from context.