rename some variables

Author: voorhs

autointent/_pipeline/_pipeline.py

@@ -10,7 +10,7 @@
 
 from autointent import Context, Dataset
 from autointent.configs import CrossEncoderConfig, EmbedderConfig, InferenceNodeConfig, LoggingConfig, VectorIndexConfig
-from autointent.custom_types import ListOfGenericLabels, NodeType, TuningType, 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.utils import load_default_search_space, load_search_space
@@ -92,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, tuning: TuningType = "brute") -> None:
+    def _fit(self, context: Context, sampler: SamplerType = "brute") -> None:
         """
         Optimize the pipeline.
 
@@ -107,7 +107,7 @@ def _fit(self, context: Context, tuning: TuningType = "brute") -> None:
         for node_type in NodeType:
             node_optimizer = self.nodes.get(node_type, None)
             if node_optimizer is not None:
-                node_optimizer.fit(context, tuning)  # 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
@@ -127,7 +127,7 @@ def fit(
         scheme: ValidationScheme = "ho",
         n_folds: int = 3,
         refit_after: bool = False,
-        tuning: TuningType = "brute",
+        sampler: SamplerType = "brute",
     ) -> Context:
         """
         Optimize the pipeline from dataset.
@@ -145,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, tuning)
+        self._fit(context, sampler)
 
         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 TuningType, ValidationScheme
+from autointent.custom_types import SamplerType, ValidationScheme
 
 from ._name import get_run_name
 
@@ -25,7 +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"
+    sampler: SamplerType = "brute"
 
 
 class LoggingConfig(BaseModel):

autointent/custom_types.py

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

autointent/nodes/_optimization/_node_optimizer.py

@@ -13,7 +13,7 @@
 
 from autointent import Dataset
 from autointent.context import Context
-from autointent.custom_types import NodeType, TuningType
+from autointent.custom_types import NodeType, SamplerType
 from autointent.nodes._nodes_info import NODES_INFO
 
 
@@ -59,7 +59,7 @@ def __init__(
         self.modules_search_spaces = search_space
         self._logger = logging.getLogger(__name__)  # TODO solve duplicate logging messages problem
 
-    def fit(self, context: Context, tuning: TuningType = "brute") -> None:
+    def fit(self, context: Context, sampler: SamplerType = "brute") -> None:
         """
         Fit the node optimizer.
 
@@ -73,19 +73,19 @@ def fit(self, context: Context, tuning: TuningType = "brute") -> None:
             n_trials = None
             if "n_trials" in search_space:
                 n_trials = search_space.pop("n_trials")
-            if tuning == "bayes":
-                sampler = optuna.samplers.TPESampler(seed=context.seed)
+            if sampler == "tpe":
+                sampler_instance = optuna.samplers.TPESampler(seed=context.seed)
                 n_trials = n_trials or 10
-            elif tuning == "brute":
-                sampler = optuna.samplers.BruteForceSampler(seed=context.seed)  # type: ignore[assignment]
+            elif sampler == "brute":
+                sampler_instance = optuna.samplers.BruteForceSampler(seed=context.seed)  # type: ignore[assignment]
                 n_trials = None
-            elif tuning == "random":
-                sampler = optuna.samplers.RandomSampler(seed=context.seed)  # type: ignore[assignment]
+            elif sampler == "random":
+                sampler_instance = optuna.samplers.RandomSampler(seed=context.seed)  # type: ignore[assignment]
                 n_trials = n_trials or 10
             else:
-                msg = f"Unexpected sampler: {tuning}"
+                msg = f"Unexpected sampler: {sampler}"
                 raise ValueError(msg)
-            study = optuna.create_study(direction="maximize", sampler=sampler)
+            study = optuna.create_study(direction="maximize", sampler=sampler_instance)
             optuna.logging.set_verbosity(optuna.logging.WARNING)
             obj = partial(self.objective, module_name=module_name, search_space=search_space, context=context)
             study.optimize(obj, n_trials=n_trials)

tests/assets/configs/bayes.yaml tests/assets/configs/optuna.yamltests/assets/configs/bayes.yaml renamed to tests/assets/configs/optuna.yaml

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

tests/conftest.py

@@ -12,20 +12,20 @@
 
 
 @pytest.mark.parametrize(
-    "tuning",
-    ["bayes", "random"],
+    "sampler",
+    ["tpe", "random"],
 )
-def test_bayes(dataset, tuning):
+def test_bayes(dataset, sampler):
     project_dir = setup_environment()
-    search_space = get_search_space("bayes")
+    search_space = get_search_space("optuna")
 
     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="ho", refit_after=False, tuning=tuning)
+    pipeline_optimizer.fit(dataset, scheme="ho", refit_after=False, sampler=sampler)
 
 
 @pytest.mark.parametrize(