diff --git a/.github/workflows/test-inference.yaml b/.github/workflows/test-inference.yaml
index ca4c45ff0..a68ef07e8 100644
--- a/.github/workflows/test-inference.yaml
+++ b/.github/workflows/test-inference.yaml
@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:
diff --git a/.github/workflows/test-nodes.yaml b/.github/workflows/test-nodes.yaml
index 99507571b..b10161724 100644
--- a/.github/workflows/test-nodes.yaml
+++ b/.github/workflows/test-nodes.yaml
@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:
diff --git a/.github/workflows/test-optimization.yaml b/.github/workflows/test-optimization.yaml
index ea1cf861e..4625f39d7 100644
--- a/.github/workflows/test-optimization.yaml
+++ b/.github/workflows/test-optimization.yaml
@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:
diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml
index 3612d561f..5883080eb 100644
--- a/.github/workflows/unit-tests.yaml
+++ b/.github/workflows/unit-tests.yaml
@@ -5,8 +5,6 @@ on:
     branches:
       - dev
   pull_request:
-    branches:
-      - dev
 
 jobs:
   test:
diff --git a/autointent/configs/_optimization.py b/autointent/configs/_optimization.py
index fb8e2068b..1d85081a4 100644
--- a/autointent/configs/_optimization.py
+++ b/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."""
 
 
diff --git a/autointent/custom_types.py b/autointent/custom_types.py
index ae5f9e6e2..b0318ee38 100644
--- a/autointent/custom_types.py
+++ b/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."""
diff --git a/autointent/modules/decision/_adaptive.py b/autointent/modules/decision/_adaptive.py
index e0d155a6e..806a35444 100644
--- a/autointent/modules/decision/_adaptive.py
+++ b/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.
 
diff --git a/autointent/modules/decision/_jinoos.py b/autointent/modules/decision/_jinoos.py
index f266d66ec..14676c2b6 100644
--- a/autointent/modules/decision/_jinoos.py
+++ b/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.
 
diff --git a/autointent/modules/decision/_threshold.py b/autointent/modules/decision/_threshold.py
index 76bf0f281..42dcbca20 100644
--- a/autointent/modules/decision/_threshold.py
+++ b/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.
 
diff --git a/autointent/modules/decision/_tunable.py b/autointent/modules/decision/_tunable.py
index 34d0c2ce9..82f348b99 100644
--- a/autointent/modules/decision/_tunable.py
+++ b/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.
 
diff --git a/autointent/modules/embedding/_logreg.py b/autointent/modules/embedding/_logreg.py
index 64283ee9c..496b245ab 100644
--- a/autointent/modules/embedding/_logreg.py
+++ b/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.
diff --git a/autointent/modules/embedding/_retrieval.py b/autointent/modules/embedding/_retrieval.py
index 5063a1b00..2ca62eff1 100644
--- a/autointent/modules/embedding/_retrieval.py
+++ b/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":
         """
diff --git a/autointent/modules/scoring/_description/description.py b/autointent/modules/scoring/_description/description.py
index 79ade2aa7..ff55ae334 100644
--- a/autointent/modules/scoring/_description/description.py
+++ b/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":
         """
diff --git a/autointent/modules/scoring/_dnnc/dnnc.py b/autointent/modules/scoring/_dnnc/dnnc.py
index 0610785ae..84b51a496 100644
--- a/autointent/modules/scoring/_dnnc/dnnc.py
+++ b/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":
diff --git a/autointent/modules/scoring/_knn/knn.py b/autointent/modules/scoring/_knn/knn.py
index ab665dfe7..eaa5ce864 100644
--- a/autointent/modules/scoring/_knn/knn.py
+++ b/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":
diff --git a/autointent/modules/scoring/_mlknn/mlknn.py b/autointent/modules/scoring/_mlknn/mlknn.py
index b43763ab0..a25220f70 100644
--- a/autointent/modules/scoring/_mlknn/mlknn.py
+++ b/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":
         """
diff --git a/autointent/nodes/_optimization/_node_optimizer.py b/autointent/nodes/_optimization/_node_optimizer.py
index 9a4a2493a..bca970ebd 100644
--- a/autointent/nodes/_optimization/_node_optimizer.py
+++ b/autointent/nodes/_optimization/_node_optimizer.py
@@ -5,11 +5,12 @@
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
-from typing import Any, TypedDict
+from typing import Any
 
 import optuna
 import torch
 from optuna.trial import Trial
+from pydantic import BaseModel, Field
 
 from autointent import Dataset
 from autointent.context import Context
@@ -17,18 +18,18 @@
 from autointent.nodes._nodes_info import NODES_INFO
 
 
-class ParamSpaceInt(TypedDict, total=False):
-    low: int
-    high: int
-    step: int
-    log: bool
+class ParamSpaceInt(BaseModel):
+    low: int = Field(..., description="Low boundary of the search space.")
+    high: int = Field(..., description="High boundary of the search space.")
+    step: int = Field(1, description="Step of the search space.")
+    log: bool = Field(False, description="Whether to use a logarithmic scale.")
 
 
-class ParamSpaceFloat(TypedDict, total=False):
-    low: float
-    high: float
-    step: float
-    log: bool
+class ParamSpaceFloat(BaseModel):
+    low: float = Field(..., description="Low boundary of the search space.")
+    high: float = Field(..., description="High boundary of the search space.")
+    step: float = Field(0.1, description="Step of the search space.")
+    log: bool = Field(False, description="Whether to use a logarithmic scale.")
 
 
 class NodeOptimizer:
@@ -145,16 +146,24 @@ def objective(
 
         return target_metric
 
-    def suggest(
-        self, trial: Trial, search_space: dict[str, ParamSpaceInt | ParamSpaceFloat | list[Any]]
-    ) -> dict[str, Any]:
+    def suggest(self, trial: Trial, search_space: dict[str, Any | list[Any]]) -> dict[str, Any]:
         res: dict[str, Any] = {}
+
+        def is_valid_param_space(
+            param_space: dict[str, Any], space_type: type[ParamSpaceInt | ParamSpaceFloat]
+        ) -> bool:
+            try:
+                space_type(**param_space)
+                return True  # noqa: TRY300
+            except ValueError:
+                return False
+
         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 all(isinstance(v, int) for v in param_space.values()):
+            elif is_valid_param_space(param_space, ParamSpaceInt):
                 res[param_name] = trial.suggest_int(param_name, **param_space)
-            elif all(isinstance(v, float) for v in param_space.values()):
+            elif is_valid_param_space(param_space, ParamSpaceFloat):
                 res[param_name] = trial.suggest_float(param_name, **param_space)
             else:
                 msg = f"Unsupported type of param search space: {param_space}"
diff --git a/autointent/nodes/schemes.py b/autointent/nodes/schemes.py
index 58cba623b..acfc71fab 100644
--- a/autointent/nodes/schemes.py
+++ b/autointent/nodes/schemes.py
@@ -2,13 +2,58 @@
 
 import inspect
 from collections.abc import Iterator
-from typing import Any, Literal, TypeAlias, Union, get_type_hints
+from typing import Annotated, Any, Literal, TypeAlias, Union, get_args, get_origin, get_type_hints
 
-from pydantic import BaseModel, Field, RootModel
+from pydantic import BaseModel, Field, PositiveInt, RootModel
 
 from autointent.custom_types import NodeType
 from autointent.modules.abc import Module
 from autointent.nodes import DecisionNodeInfo, EmbeddingNodeInfo, ScoringNodeInfo
+from autointent.nodes._optimization._node_optimizer import ParamSpaceFloat, ParamSpaceInt
+
+
+def unwrap_annotated(tp: type) -> type:
+    """
+    Unwrap the Annotated type to get the actual type.
+
+    :param tp: Type to unwrap
+    :return: Unwrapped type
+    """
+    return get_args(tp)[0] if get_origin(tp) is Annotated else tp
+
+
+def type_matches(target: type, tp: type) -> bool:
+    """
+    Recursively check if the target type is present in the given type.
+
+    This function handles union types by unwrapping Annotated types where necessary.
+
+    :param target: Target type
+    :param tp: Given type
+    :return: If the target type is present in the given type
+    """
+    origin = get_origin(tp)
+
+    if origin is Union:  # float | list[float]
+        return any(type_matches(target, arg) for arg in get_args(tp))
+    return unwrap_annotated(tp) is target
+
+
+def get_optuna_class(param_type: type) -> type[ParamSpaceInt | ParamSpaceFloat] | None:
+    """
+    Get the Optuna class for the given parameter type.
+
+    If the (possibly annotated or union) type includes int or float, this function
+    returns the corresponding search space class.
+
+    :param param_type: Parameter type (could be a union, annotated type, or container)
+    :return: ParamSpaceInt if the type matches int, ParamSpaceFloat if it matches float, else None.
+    """
+    if type_matches(int, param_type):
+        return ParamSpaceInt
+    if type_matches(float, param_type):
+        return ParamSpaceFloat
+    return None
 
 
 def generate_models_and_union_type_for_classes(
@@ -20,9 +65,12 @@ def generate_models_and_union_type_for_classes(
     for cls in classes:
         init_signature = inspect.signature(cls.from_context)
         globalns = getattr(cls.from_context, "__globals__", {})
-        type_hints = get_type_hints(cls.from_context, globalns, None)  # Resolve forward refs
+        type_hints = get_type_hints(cls.from_context, globalns, None, include_extras=True)  # Resolve forward refs
 
-        fields = {"module_name": (Literal[cls.name], Field(...))}
+        fields = {
+            "module_name": (Literal[cls.name], Field(...)),
+            "n_trials": (PositiveInt | None, Field(None, description="Number of trials")),
+        }
 
         for param_name, param in init_signature.parameters.items():
             if param_name in ("self", "cls", "context"):
@@ -30,8 +78,11 @@ def generate_models_and_union_type_for_classes(
 
             param_type: TypeAlias = type_hints.get(param_name, Any)  # type: ignore[valid-type]  # noqa: PYI042
             field = Field(default=[param.default]) if param.default is not inspect.Parameter.empty else Field(...)
-
-            fields[param_name] = (list[param_type], field)  # type: ignore[assignment]
+            search_type = get_optuna_class(param_type)
+            if search_type is None:
+                fields[param_name] = (list[param_type], field)
+            else:
+                fields[param_name] = (list[param_type] | search_type, field)
 
         model_name = f"{cls.__name__}InitModel"
         models[cls.__name__] = type(
diff --git a/docs/optimizer_config.schema.json b/docs/optimizer_config.schema.json
index 73e79757e..6980b460f 100644
--- a/docs/optimizer_config.schema.json
+++ b/docs/optimizer_config.schema.json
@@ -7,6 +7,20 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "search_space": {
                     "default": [
                         null
@@ -15,6 +29,8 @@
                         "anyOf": [
                             {
                                 "items": {
+                                    "maximum": 1.0,
+                                    "minimum": 0.0,
                                     "type": "number"
                                 },
                                 "type": "array"
@@ -40,6 +56,20 @@
                     "const": "argmax",
                     "title": "Module Name",
                     "type": "string"
+                },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
                 }
             },
             "required": [
@@ -55,6 +85,20 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "cross_encoder_name": {
                     "items": {
                         "type": "string"
@@ -63,11 +107,19 @@
                     "type": "array"
                 },
                 "k": {
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "K",
-                    "type": "array"
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
+                    "title": "K"
                 },
                 "embedder_name": {
                     "default": [
@@ -183,12 +235,34 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "temperature": {
-                    "items": {
-                        "type": "number"
-                    },
-                    "title": "Temperature",
-                    "type": "array"
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0.0,
+                                "type": "number"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceFloat"
+                        }
+                    ],
+                    "title": "Temperature"
                 },
                 "embedder_name": {
                     "default": [
@@ -324,6 +398,20 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "search_space": {
                     "default": [
                         null
@@ -332,6 +420,8 @@
                         "anyOf": [
                             {
                                 "items": {
+                                    "maximum": 1.0,
+                                    "minimum": 0.0,
                                     "type": "number"
                                 },
                                 "type": "array"
@@ -358,12 +448,34 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "k": {
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "K",
-                    "type": "array"
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
+                    "title": "K"
                 },
                 "weights": {
                     "items": {
@@ -410,6 +522,20 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "embedder_name": {
                     "default": [
                         null
@@ -441,6 +567,20 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "embedder_name": {
                     "items": {
                         "type": "string"
@@ -449,14 +589,22 @@
                     "type": "array"
                 },
                 "cv": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
                     "default": [
                         3
                     ],
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "Cv",
-                    "type": "array"
+                    "title": "Cv"
                 }
             },
             "required": [
@@ -473,32 +621,70 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "k": {
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "K",
-                    "type": "array"
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
+                    "title": "K"
                 },
                 "s": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0.0,
+                                "type": "number"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceFloat"
+                        }
+                    ],
                     "default": [
                         1.0
                     ],
-                    "items": {
-                        "type": "number"
-                    },
-                    "title": "S",
-                    "type": "array"
+                    "title": "S"
                 },
                 "ignore_first_neighbours": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "minimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
                     "default": [
                         0
                     ],
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "Ignore First Neighbours",
-                    "type": "array"
+                    "title": "Ignore First Neighbours"
                 },
                 "embedder_name": {
                     "default": [
@@ -536,6 +722,70 @@
             "title": "NodeType",
             "type": "string"
         },
+        "ParamSpaceFloat": {
+            "properties": {
+                "low": {
+                    "description": "Low boundary of the search space.",
+                    "title": "Low",
+                    "type": "number"
+                },
+                "high": {
+                    "description": "High boundary of the search space.",
+                    "title": "High",
+                    "type": "number"
+                },
+                "step": {
+                    "default": 0.1,
+                    "description": "Step of the search space.",
+                    "title": "Step",
+                    "type": "number"
+                },
+                "log": {
+                    "default": false,
+                    "description": "Whether to use a logarithmic scale.",
+                    "title": "Log",
+                    "type": "boolean"
+                }
+            },
+            "required": [
+                "low",
+                "high"
+            ],
+            "title": "ParamSpaceFloat",
+            "type": "object"
+        },
+        "ParamSpaceInt": {
+            "properties": {
+                "low": {
+                    "description": "Low boundary of the search space.",
+                    "title": "Low",
+                    "type": "integer"
+                },
+                "high": {
+                    "description": "High boundary of the search space.",
+                    "title": "High",
+                    "type": "integer"
+                },
+                "step": {
+                    "default": 1,
+                    "description": "Step of the search space.",
+                    "title": "Step",
+                    "type": "integer"
+                },
+                "log": {
+                    "default": false,
+                    "description": "Whether to use a logarithmic scale.",
+                    "title": "Log",
+                    "type": "boolean"
+                }
+            },
+            "required": [
+                "low",
+                "high"
+            ],
+            "title": "ParamSpaceInt",
+            "type": "object"
+        },
         "RerankScorerInitModel": {
             "properties": {
                 "module_name": {
@@ -543,12 +793,33 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "k": {
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "K",
-                    "type": "array"
+                    "anyOf": [
+                        {
+                            "items": {
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
+                    "title": "K"
                 },
                 "weights": {
                     "items": {
@@ -597,38 +868,52 @@
                     "type": "array"
                 },
                 "m": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "anyOf": [
+                                    {
+                                        "type": "integer"
+                                    },
+                                    {
+                                        "type": "null"
+                                    }
+                                ]
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
                     "default": [
                         null
                     ],
-                    "items": {
-                        "anyOf": [
-                            {
-                                "type": "integer"
-                            },
-                            {
-                                "type": "null"
-                            }
-                        ]
-                    },
-                    "title": "M",
-                    "type": "array"
+                    "title": "M"
                 },
                 "rank_threshold_cutoff": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "anyOf": [
+                                    {
+                                        "type": "integer"
+                                    },
+                                    {
+                                        "type": "null"
+                                    }
+                                ]
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
                     "default": [
                         null
                     ],
-                    "items": {
-                        "anyOf": [
-                            {
-                                "type": "integer"
-                            },
-                            {
-                                "type": "null"
-                            }
-                        ]
-                    },
-                    "title": "Rank Threshold Cutoff",
-                    "type": "array"
+                    "title": "Rank Threshold Cutoff"
                 }
             },
             "required": [
@@ -647,12 +932,34 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "k": {
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "K",
-                    "type": "array"
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
+                    "title": "K"
                 },
                 "embedder_name": {
                     "items": {
@@ -764,6 +1071,20 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "clf_name": {
                     "default": [
                         "LogisticRegression"
@@ -822,25 +1143,50 @@
                     "title": "Module Name",
                     "type": "string"
                 },
+                "n_trials": {
+                    "anyOf": [
+                        {
+                            "exclusiveMinimum": 0,
+                            "type": "integer"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Number of trials",
+                    "title": "N Trials"
+                },
                 "thresh": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "anyOf": [
+                                    {
+                                        "maximum": 1.0,
+                                        "minimum": 0.0,
+                                        "type": "number"
+                                    },
+                                    {
+                                        "items": {
+                                            "maximum": 1.0,
+                                            "minimum": 0.0,
+                                            "type": "number"
+                                        },
+                                        "type": "array"
+                                    }
+                                ]
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceFloat"
+                        }
+                    ],
                     "default": [
                         0.5
                     ],
-                    "items": {
-                        "anyOf": [
-                            {
-                                "type": "number"
-                            },
-                            {
-                                "items": {
-                                    "type": "number"
-                                },
-                                "type": "array"
-                            }
-                        ]
-                    },
-                    "title": "Thresh",
-                    "type": "array"
+                    "title": "Thresh"
                 }
             },
             "required": [
@@ -857,14 +1203,22 @@
                     "type": "string"
                 },
                 "n_trials": {
+                    "anyOf": [
+                        {
+                            "items": {
+                                "exclusiveMinimum": 0,
+                                "type": "integer"
+                            },
+                            "type": "array"
+                        },
+                        {
+                            "$ref": "#/$defs/ParamSpaceInt"
+                        }
+                    ],
                     "default": [
                         320
                     ],
-                    "items": {
-                        "type": "integer"
-                    },
-                    "title": "N Trials",
-                    "type": "array"
+                    "title": "N Trials"
                 }
             },
             "required": [