add config validation for optuna

.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/configs/_optimization.py

@@ -2,7 +2,7 @@
 
 from pathlib import Path
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PositiveInt
 
 from autointent.custom_types import SamplerType, ValidationScheme
 
@@ -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."""
 
 

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):
@@ -73,3 +75,7 @@ class Split:
 
 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
@@ -59,7 +59,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.
 
@@ -69,7 +69,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.
 

autointent/modules/decision/_tunable.py

@@ -6,6 +6,7 @@
 import numpy.typing as npt
 import optuna
 from optuna.trial import Trial
+from pydantic import PositiveInt
 
 from autointent.context import Context
 from autointent.custom_types import ListOfGenericLabels
@@ -77,7 +78,7 @@ class TunableDecision(DecisionModule):
 
     def __init__(
         self,
-        n_trials: int = 320,
+        n_trials: PositiveInt = 320,
         seed: int = 0,
         tags: list[Tag] | None = None,
     ) -> None:
@@ -93,7 +94,7 @@ def __init__(
         self.tags = tags
 
     @classmethod
-    def from_context(cls, context: Context, n_trials: int = 320) -> "TunableDecision":
+    def from_context(cls, context: Context, n_trials: PositiveInt = 320) -> "TunableDecision":
         """
         Initialize from context.
 

autointent/modules/embedding/_logreg.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from pydantic import PositiveInt
 from sklearn.linear_model import LogisticRegression, LogisticRegressionCV
 from sklearn.multioutput import MultiOutputClassifier
 from sklearn.preprocessing import LabelEncoder
@@ -48,7 +49,7 @@ class LogregAimedEmbedding(EmbeddingModule):
     def __init__(
         self,
         embedder_name: str,
-        cv: int = 3,
+        cv: PositiveInt = 3,
         embedder_device: str = "cpu",
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
@@ -76,7 +77,7 @@ def from_context(
         cls,
         context: Context,
         embedder_name: str,
-        cv: int = 3,
+        cv: PositiveInt = 3,
     ) -> "LogregAimedEmbedding":
         """
         Create a LogregAimedEmbedding instance using a Context object.

autointent/modules/embedding/_retrieval.py

@@ -1,5 +1,7 @@
 """RetrievalAimedEmbedding class for a proxy optimization of embedding."""
 
+from pydantic import PositiveInt
+
 from autointent import Context, VectorIndex
 from autointent.context.optimization_info import RetrieverArtifact
 from autointent.custom_types import ListOfLabels
@@ -41,7 +43,7 @@ class RetrievalAimedEmbedding(EmbeddingModule):
 
     def __init__(
         self,
-        k: int,
+        k: PositiveInt,
         embedder_name: str,
         embedder_device: str = "cpu",
         embedder_batch_size: int = 32,
@@ -69,7 +71,7 @@ def __init__(
     def from_context(
         cls,
         context: Context,
-        k: int,
+        k: PositiveInt,
         embedder_name: str,
     ) -> "RetrievalAimedEmbedding":
         """

autointent/modules/scoring/_description/description.py

@@ -5,6 +5,7 @@
 import numpy as np
 import scipy
 from numpy.typing import NDArray
+from pydantic import PositiveFloat
 from sklearn.metrics.pairwise import cosine_similarity
 
 from autointent import Context, Embedder
@@ -37,7 +38,7 @@ class DescriptionScorer(ScoringModule):
     def __init__(
         self,
         embedder_name: str,
-        temperature: float = 1.0,
+        temperature: PositiveFloat = 1.0,
         embedder_device: str = "cpu",
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
@@ -64,7 +65,7 @@ def __init__(
     def from_context(
         cls,
         context: Context,
-        temperature: float,
+        temperature: PositiveFloat,
         embedder_name: str | None = None,
     ) -> "DescriptionScorer":
         """

autointent/modules/scoring/_dnnc/dnnc.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import numpy.typing as npt
+from pydantic import PositiveInt
 
 from autointent import Context, Ranker, VectorIndex
 from autointent.custom_types import ListOfLabels
@@ -77,7 +78,7 @@ def __init__(  # noqa: PLR0913
         self,
         cross_encoder_name: str,
         embedder_name: str,
-        k: int,
+        k: PositiveInt,
         embedder_device: str = "cpu",
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
@@ -118,7 +119,7 @@ def from_context(
         cls,
         context: Context,
         cross_encoder_name: str,
-        k: int,
+        k: PositiveInt,
         embedder_name: str | None = None,
         train_head: bool = False,
     ) -> "DNNCScorer":

autointent/modules/scoring/_knn/knn.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import numpy.typing as npt
+from pydantic import PositiveInt
 
 from autointent import Context, VectorIndex
 from autointent.custom_types import WEIGHT_TYPES, ListOfLabels
@@ -57,7 +58,7 @@ class KNNScorer(ScoringModule):
     def __init__(
         self,
         embedder_name: str,
-        k: int,
+        k: PositiveInt,
         weights: WEIGHT_TYPES = "distance",
         embedder_device: str = "cpu",
         embedder_batch_size: int = 32,
@@ -90,7 +91,7 @@ def __init__(
     def from_context(
         cls,
         context: Context,
-        k: int,
+        k: PositiveInt,
         weights: WEIGHT_TYPES,
         embedder_name: str | None = None,
     ) -> "KNNScorer":

autointent/modules/scoring/_mlknn/mlknn.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
+from pydantic import NonNegativeInt, PositiveFloat, PositiveInt
 
 from autointent import Context, VectorIndex
 from autointent.custom_types import ListOfLabels
@@ -57,7 +58,7 @@ class MLKnnScorer(ScoringModule):
 
     def __init__(
         self,
-        k: int,
+        k: PositiveInt,
         embedder_name: str,
         s: float = 1.0,
         ignore_first_neighbours: int = 0,
@@ -91,9 +92,9 @@ def __init__(
     def from_context(
         cls,
         context: Context,
-        k: int,
-        s: float = 1.0,
-        ignore_first_neighbours: int = 0,
+        k: PositiveInt,
+        s: PositiveFloat = 1.0,
+        ignore_first_neighbours: NonNegativeInt = 0,
         embedder_name: str | None = None,
     ) -> "MLKnnScorer":
         """