rapidfireai 0.10.3rc1__py3-none-any.whl → 0.11.1rc2__py3-none-any.whl

rapidfireai/automl/grid_search.py

@@ -1,8 +1,7 @@
 """Grid search implementation for AutoML training configurations."""
 
 from itertools import product
-from typing import Any, Dict
-from typing import List as ListType
+from typing import Any
 
 from rapidfireai.automl.base import AutoMLAlgorithm
 from rapidfireai.automl.datatypes import List
@@ -15,7 +14,7 @@ def recursive_expand_gridsearch(item: Any):
         keys = list(item.keys())
         value_lists = [list(recursive_expand_gridsearch(item[k])) for k in keys]
         for values in product(*value_lists):
-            yield dict(zip(keys, values))
+            yield dict(zip(keys, values, strict=False))
     elif isinstance(item, List):
         for value in item.values:
             yield from recursive_expand_gridsearch(value)
@@ -26,7 +25,7 @@ def recursive_expand_gridsearch(item: Any):
 class RFGridSearch(AutoMLAlgorithm):
     """Grid search algorithm that generates all hyperparameter combinations."""
 
-    def get_runs(self, seed: int) -> ListType[Dict[str, Any]]:
+    def get_runs(self, seed: int) -> list[dict[str, Any]]:
         """Generate all possible hyperparameter combinations for grid search."""
         if not isinstance(seed, int) or seed < 0:
             raise AutoMLException("seed must be a non-negative integer")
@@ -122,4 +121,4 @@ class RFGridSearch(AutoMLAlgorithm):
             return runs
 
         except Exception as e:
-            raise AutoMLException(f"Error generating runs: {e}") from e
+            raise AutoMLException(f"Error generating runs: {e}") from e

rapidfireai/automl/model_config.py

@@ -1,16 +1,18 @@
 """Model configuration for AutoML training."""
 
-import inspect
 import copy
+import inspect
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Type, Union, get_type_hints
+from typing import Any, get_type_hints
 
 from peft import LoraConfig
 from trl import DPOConfig, GRPOConfig, SFTConfig
 
 from rapidfireai.automl.datatypes import List, Range
 
-def _create_rf_class(base_class: Type, class_name: str):
+
+def _create_rf_class(base_class: type, class_name: str):
     """Creating a RF class that dynamically inherits all constructor parameters and supports singleton, list, and Range values."""
     if not inspect.isclass(base_class):
         raise ValueError(f"base_class must be a class, got {type(base_class)}")
@@ -28,45 +30,47 @@ def _create_rf_class(base_class: Type, class_name: str):
     def __init__(self, **kwargs):
         self._user_params = copy.deepcopy(kwargs)
         self._constructor_params = constructor_params
-        self._initializing = True  
-        
+        self._initializing = True
+
         parent_kwargs = {}
         for key, value in kwargs.items():
             if not isinstance(value, (List, Range)):
                 parent_kwargs[key] = value
-        
+
         base_class.__init__(self, **parent_kwargs)
-        
-        self._initializing = False  
+
+        self._initializing = False
+
     def copy_config(self):
         """Create a deep copy of the configuration."""
-        copied_params = copy.deepcopy(self._user_params)        
+        copied_params = copy.deepcopy(self._user_params)
         new_instance = self.__class__(**copied_params)
-        
+
         return new_instance
-    
+
     def __setattr__(self, name, value):
         """Override setattr to update _user_params when constructor parameters are modified."""
-        
-        if (hasattr(self, '_constructor_params') and 
-            name in self._constructor_params and 
-            hasattr(self, '_user_params') and
-            name in self._user_params and
-            not getattr(self, '_initializing', True)):  # Don't update during init
+
+        if (
+            hasattr(self, "_constructor_params")
+            and name in self._constructor_params
+            and hasattr(self, "_user_params")
+            and name in self._user_params
+            and not getattr(self, "_initializing", True)
+        ):  # Don't update during init
             self._user_params[name] = value
-        
+
         base_class.__setattr__(self, name, value)
-        
 
     return type(
         class_name,
         (base_class,),
         {
-            "__doc__": f"RF version of {base_class.__name__}", 
-            "__annotations__": new_type_hints, 
+            "__doc__": f"RF version of {base_class.__name__}",
+            "__annotations__": new_type_hints,
             "__init__": __init__,
             "copy": copy_config,
-            "__setattr__": __setattr__
+            "__setattr__": __setattr__,
         },
     )
 
@@ -83,20 +87,20 @@ class RFModelConfig:
     """Model configuration for AutoML training."""
 
     model_name: str = None
-    tokenizer: Optional[str] = None
-    tokenizer_kwargs: Optional[dict[str, Any]] = None
-    formatting_func: Optional[Union[Callable, List]] = None
-    compute_metrics: Optional[Union[Callable, List]] = None
-    peft_config: Optional[Union[RFLoraConfig, List]] = None
-    training_args: Optional[Union[RFSFTConfig, RFDPOConfig, RFGRPOConfig]] = None
-    model_type: Optional[str] = "causal_lm"
-    model_kwargs: Optional[dict[str, Any]] = None
-    ref_model_name: Optional[str] = None
-    ref_model_type: Optional[str] = None
-    ref_model_kwargs: Optional[dict[str, Any]] = None
-    reward_funcs: Optional[Union[str, List, Callable, Any]] = None
-    generation_config: Optional[dict[str, Any]] = None
-
-    def copy(self):#FIXME: Handle similar to create_rf_class
+    tokenizer: str | None = None
+    tokenizer_kwargs: dict[str, Any] | None = None
+    formatting_func: Callable | List | None = None
+    compute_metrics: Callable | List | None = None
+    peft_config: RFLoraConfig | List | None = None
+    training_args: RFSFTConfig | RFDPOConfig | RFGRPOConfig | None = None
+    model_type: str | None = "causal_lm"
+    model_kwargs: dict[str, Any] | None = None
+    ref_model_name: str | None = None
+    ref_model_type: str | None = None
+    ref_model_kwargs: dict[str, Any] | None = None
+    reward_funcs: str | List | Callable | Any | None = None
+    generation_config: dict[str, Any] | None = None
+
+    def copy(self):  # FIXME: Handle similar to create_rf_class
         """Create a deep copy of the RFModelConfig."""
         return copy.deepcopy(self)

rapidfireai/automl/random_search.py

@@ -1,10 +1,7 @@
 """Random search implementation for AutoML hyperparameter optimization."""
 
 import random
-import json
-from itertools import product
-from typing import Any, Dict
-from typing import List as ListType
+from typing import Any
 
 from rapidfireai.automl.base import AutoMLAlgorithm
 from rapidfireai.automl.datatypes import List, Range
@@ -15,38 +12,36 @@ from rapidfireai.utils.serialize import encode_payload
 def recursive_expand_randomsearch(item: Any):
     if isinstance(item, dict):
         return {k: recursive_expand_randomsearch(v) for k, v in item.items()}
-    elif isinstance(item, List):
-        return item.sample()
-    elif isinstance(item, Range):
+    elif isinstance(item, List) or isinstance(item, Range):
         return item.sample()
     else:
         return item
 
+
 class RFRandomSearch(AutoMLAlgorithm):
     """Random search algorithm that samples num_runs hyperparameter combinations."""
 
-    def get_runs(self, seed: int=42) -> ListType[Dict[str, Any]]:
+    def get_runs(self, seed: int = 42) -> list[dict[str, Any]]:
         """Generate num_runs random hyperparameter combinations."""
         if seed is not None and (not isinstance(seed, int) or seed < 0):
             raise AutoMLException("seed must be a non-negative integer")
-        
+
         if not isinstance(self.num_runs, int) or self.num_runs <= 0:
             raise AutoMLException("num_runs must be a positive integer")
 
-        
         random.seed(seed)
-        
+
         try:
             runs = []
             seen_configs = set()
-            max_attempts = self.num_runs * 10 
+            max_attempts = self.num_runs * 10
             attempts = 0
-            
+
             while len(runs) < self.num_runs and attempts < max_attempts:
                 attempts += 1
 
                 config = List(self.configs).sample()
-                
+
                 if config.peft_config is None:
                     selected_peft_config = None
                 elif isinstance(config.peft_config, list):
@@ -57,32 +52,26 @@ class RFRandomSearch(AutoMLAlgorithm):
                     selected_peft_config = config.peft_config
 
                 peft_params = (
-                    {} if selected_peft_config is None 
+                    {}
+                    if selected_peft_config is None
                     else recursive_expand_randomsearch(selected_peft_config._user_params)
                 )
 
-                
                 # Sample other parameters
                 training_params = (
-                    {} if config.training_args is None 
+                    {}
+                    if config.training_args is None
                     else recursive_expand_randomsearch(config.training_args._user_params)
                 )
-                
-                model_kwargs = (
-                    {} if config.model_kwargs is None 
-                    else recursive_expand_randomsearch(config.model_kwargs)
-                )
-                
+
+                model_kwargs = {} if config.model_kwargs is None else recursive_expand_randomsearch(config.model_kwargs)
+
                 ref_model_kwargs = (
-                    {} if config.ref_model_kwargs is None 
-                    else recursive_expand_randomsearch(config.ref_model_kwargs)
-                )
-                
-                reward_funcs = (
-                    {} if config.reward_funcs is None 
-                    else recursive_expand_randomsearch(config.reward_funcs)
+                    {} if config.ref_model_kwargs is None else recursive_expand_randomsearch(config.ref_model_kwargs)
                 )
 
+                reward_funcs = {} if config.reward_funcs is None else recursive_expand_randomsearch(config.reward_funcs)
+
                 # FIXME:  avoid hardcoding the excluded attributes
                 excluded_attrs = {
                     "model_name",
@@ -101,8 +90,7 @@ class RFRandomSearch(AutoMLAlgorithm):
                     k: v for k, v in config.__dict__.items() if k not in excluded_attrs and v is not None
                 }
                 additional_kwargs_sampled = (
-                    {} if not additional_kwargs 
-                    else recursive_expand_randomsearch(additional_kwargs)
+                    {} if not additional_kwargs else recursive_expand_randomsearch(additional_kwargs)
                 )
 
                 leaf = {
@@ -123,7 +111,7 @@ class RFRandomSearch(AutoMLAlgorithm):
                         "model_type": config.ref_model_type,
                         "model_kwargs": ref_model_kwargs,
                     }
-                    #FIXME: correct ref args
+                    # FIXME: correct ref args
                 elif self.trainer_type == "GRPO":
                     leaf["reward_funcs"] = reward_funcs