22# Copyright (C) 2025-2026 Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
33
44from enum import Enum
5- from typing import Dict
5+ from typing import Dict , Type , TypeAlias , TypeVar
66
7- from pydantic import BaseModel , Field
7+ from optuna import Trial
8+ from pydantic import BaseModel , Field , PrivateAttr , RootModel
89from pydantic_settings import (
910 BaseSettings ,
1011 CliSettingsSource ,
1112 EnvSettingsSource ,
1213 PydanticBaseSettingsSource ,
1314 TomlConfigSettingsSource ,
1415)
16+ from typing_extensions import TypeAliasType
1517
1618
1719class QuantizationMethod (str , Enum ):
@@ -61,6 +63,63 @@ class DatasetSpecification(BaseModel):
6163 )
6264
6365
66+ class CategoricalParamSpecification (RootModel [list [str ]]):
67+ _name : str | None = PrivateAttr (None )
68+
69+ def __len__ (self ) -> int :
70+ return len (self .root )
71+
72+ def __contains__ (self , item : str ) -> bool :
73+ return item in self .root
74+
75+ def set_name (self , name : str ):
76+ self ._name = name
77+
78+ def suggest_categorical (self , trial : Trial ) -> str :
79+ if self ._name is None :
80+ raise TypeError ("Parameter name is not set." )
81+ if len (self .root ) == 1 :
82+ return self .root [0 ]
83+ return trial .suggest_categorical (self ._name , self .root )
84+
85+
86+ class FloatParamSpecification (BaseModel ):
87+ low : float = Field (
88+ description = "Lower endpoint of the range of suggested values (inclusive)." ,
89+ )
90+
91+ high : float = Field (
92+ description = "Upper endpoint of the range of suggested values (inclusive)." ,
93+ )
94+
95+ log : bool = Field (
96+ default = False ,
97+ description = "If true, the value is sampled from the range in the log domain." ,
98+ )
99+
100+ _name : str | None = PrivateAttr (None )
101+
102+ def set_name (self , name : str ):
103+ self ._name = name
104+
105+ def suggest_float (self , trial : Trial ) -> float :
106+ if self ._name is None :
107+ raise TypeError ("Parameter name is not set." )
108+ if self .low == self .high :
109+ return self .low
110+ return trial .suggest_float (self ._name , self .low , self .high , log = self .log )
111+
112+
113+ ParamSpecification : TypeAlias = CategoricalParamSpecification | FloatParamSpecification
114+ ParamSpecificationRecursive = TypeAliasType (
115+ "ParamSpecificationRecursive" ,
116+ "ParamSpecification | dict[str, ParamSpecificationRecursive]" ,
117+ )
118+ ParamSpecificationType = TypeVar (
119+ "ParamSpecificationType" , CategoricalParamSpecification , FloatParamSpecification
120+ )
121+
122+
64123class Settings (BaseSettings ):
65124 model : str = Field (description = "Hugging Face model ID, or path to model on disk." )
66125
@@ -313,6 +372,13 @@ class Settings(BaseSettings):
313372 description = "Dataset of prompts that tend to result in refusals (used for evaluating model performance)." ,
314373 )
315374
375+ parameter_ranges : dict [str , ParamSpecificationRecursive ] = Field (
376+ default = {},
377+ description = "Parameter ranges, per parameter or per module and parameter." ,
378+ )
379+
380+ _parameter_range_defaults : dict [str , ParamSpecification ] | None = PrivateAttr (None )
381+
316382 @classmethod
317383 def settings_customise_sources (
318384 cls ,
@@ -335,3 +401,55 @@ def settings_customise_sources(
335401 file_secret_settings ,
336402 TomlConfigSettingsSource (settings_cls , toml_file = "config.toml" ),
337403 )
404+
405+ def set_parameter_range_defaults (self , defaults : dict [str , ParamSpecification ]):
406+ self ._parameter_range_defaults = defaults
407+
408+ def _param_spec (
409+ self , name : str , spec_type : Type [ParamSpecificationType ]
410+ ) -> ParamSpecificationType :
411+ """
412+ Finds the setting with the longest matching suffix.
413+ For example, if we're looking for setting `foo.bar.baz` and
414+ both `baz` and `bar.baz` are available, return `bar.baz`.
415+ """
416+ if self ._parameter_range_defaults is None :
417+ raise ValueError (
418+ "Parameter defaults are not initialized. "
419+ "Call 'set_parameter_range_defaults()' first."
420+ )
421+ split_name = name .split ("." )
422+ # Get the default parameter specification.
423+ most_specific_spec = self ._parameter_range_defaults [split_name [- 1 ]]
424+ if most_specific_spec is None :
425+ raise TypeError (f"No parameter default found for { split_name [- 1 ]} )
426+ if not isinstance (most_specific_spec , spec_type ):
427+ raise TypeError (
428+ f"Expected parameter default to be an instance of { spec_type .__name__ }
429+ f"but got { type (most_specific_spec ).__name__ }
430+ )
431+ # Look for an override from the config.
432+ split_len = len (split_name )
433+ for i in range (split_len - 1 , - 1 , - 1 ):
434+ spec = self .parameter_ranges .get (split_name [i ], None )
435+ for j in range (i + 1 , split_len ):
436+ if spec is None :
437+ break
438+ if not isinstance (spec , dict ):
439+ # Found a setting for the prefix, but didn't match suffix.
440+ spec = None
441+ break
442+ spec = spec .get (split_name [j ], None )
443+ if spec is not None and isinstance (spec , spec_type ):
444+ most_specific_spec = spec
445+ # Set the specific name (for suggest_categorical() and suggest_float()).
446+ most_specific_spec .set_name (name )
447+ # most_specific_spec can only be ParamSpecificationType,
448+ # but ty's support for narrowing generics is still limited.
449+ return most_specific_spec # ty:ignore[invalid-return-type]
450+
451+ def categorical_spec (self , name : str ) -> CategoricalParamSpecification :
452+ return self ._param_spec (name , CategoricalParamSpecification )
453+
454+ def float_spec (self , name : str ) -> FloatParamSpecification :
455+ return self ._param_spec (name , FloatParamSpecification )
0 commit comments