rapidfireai 0.0.1__py3-none-any.whl → 0.9.10__py3-none-any.whl

rapidfireai/__init__.py

@@ -1,11 +1,17 @@
 """
-RapidFire AI - Coming Soon
+RapidFire AI
 """
 
-__version__ = "0.0.1"
-__author__ = "Pradyumna Sridhara"
-__email__ = "pradyumna@rapidfire.ai"
+from .version import __version__, __version_info__
+
+__author__ = "RapidFire AI Inc."
+__email__ = "support@rapidfire.ai"
+
+from rapidfireai.experiment import Experiment
+
 
 def coming_soon():
     """Placeholder function - full functionality coming soon."""
-    return "RapidFire AI package is under development. Stay tuned!"
+    return "RapidFire AI package is under development. Stay tuned!"
+
+__all__ = ["Experiment"]

rapidfireai/automl/__init__.py

@@ -0,0 +1,20 @@
+"""AutoML module for hyperparameter optimization."""
+
+from .base import AutoMLAlgorithm
+from .datatypes import List, Range
+from .grid_search import RFGridSearch
+from .model_config import RFDPOConfig, RFGRPOConfig, RFLoraConfig, RFModelConfig, RFSFTConfig
+from .random_search import RFRandomSearch
+
+__all__ = [
+    "List",
+    "Range",
+    "RFGridSearch",
+    "RFRandomSearch",
+    "AutoMLAlgorithm",
+    "RFModelConfig",
+    "RFLoraConfig",
+    "RFSFTConfig",
+    "RFDPOConfig",
+    "RFGRPOConfig",
+]

rapidfireai/automl/base.py

@@ -0,0 +1,48 @@
+"""Base classes and configurations for AutoML algorithms."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from rapidfireai.automl.datatypes import List
+from rapidfireai.automl.model_config import RFModelConfig
+from rapidfireai.utils.exceptions import AutoMLException
+
+
+class AutoMLAlgorithm(ABC):
+    """Base class for AutoML algorithms."""
+
+    VALID_TRAINER_TYPES = {"SFT", "DPO", "GRPO"}
+
+    def __init__(self, configs=None, create_model_fn=None, trainer_type: str = "SFT", num_runs: int = 1):
+        """Initialize AutoML algorithm with configurations and trainer type."""
+        try:
+            self.configs = self._normalize_configs(configs)
+            self.trainer_type = trainer_type.upper()
+            self.num_runs = num_runs
+
+            if self.trainer_type not in self.VALID_TRAINER_TYPES:
+                raise AutoMLException(f"trainer_type must be one of {self.VALID_TRAINER_TYPES}")
+
+            self._validate_configs()
+        except Exception as e:
+            raise AutoMLException(f"Error initializing {self.__class__.__name__}: {e}") from e
+
+    def _normalize_configs(self, configs):
+        """Normalize configs to list format."""
+        if isinstance(configs, List):
+            return configs.values
+        elif isinstance(configs, list):
+            return configs
+        return [configs] if configs else []
+
+    def _validate_configs(self):
+        """Validate all configs are RFModelConfig instances."""
+        for config in self.configs:
+            if not isinstance(config, RFModelConfig):
+                raise AutoMLException(f"All configs must be RFModelConfig instances, got {type(config)}")
+
+    @abstractmethod
+    def get_runs(self, seed: int) -> list[dict[str, Any]]:
+        """Generate hyperparameter combinations for different training configurations."""
+        if not isinstance(seed, int) or seed < 0:
+            raise AutoMLException("seed must be a non-negative integer")

rapidfireai/automl/datatypes.py

@@ -0,0 +1,42 @@
+"""Contains classes for representing hyperparameter data types."""
+
+import random
+
+# TODO: need to set seed for random module.
+# TODO: List.sample() will not work for nested lists.
+# TODO: add support for sampling methods like 'uniform' and 'loguniform'.
+
+
+class Range:
+    """Represents a range of values for a hyperparameter."""
+
+    def __init__(self, start, end, dtype: str | None = None):
+        if dtype is None:
+            self.dtype = "int" if isinstance(start, int) and isinstance(end, int) else "float"
+        else:
+            if dtype not in ("int", "float"):
+                raise ValueError("dtype must be either 'int' or 'float'.")
+            self.dtype = dtype
+        if not (isinstance(start, (int, float)) and isinstance(end, (int, float))):
+            raise ValueError("start and end must be either int or float.")
+        self.start = start
+        self.end = end
+
+    def sample(self):
+        """Sample a value from the range [self.start, self.end]."""
+        if self.dtype == "int":
+            return random.randint(self.start, self.end)
+        return random.uniform(self.start, self.end)
+
+
+class List:
+    """Represents a list of values for a hyperparameter."""
+
+    def __init__(self, values):
+        if not isinstance(values, list):
+            raise ValueError("List expects a list of values.")
+        self.values = values
+
+    def sample(self):
+        """Sample a value from the list."""
+        return random.choice(self.values)

rapidfireai/automl/grid_search.py

@@ -0,0 +1,125 @@
+"""Grid search implementation for AutoML training configurations."""
+
+from itertools import product
+from typing import Any, Dict
+from typing import List as ListType
+
+from rapidfireai.automl.base import AutoMLAlgorithm
+from rapidfireai.automl.datatypes import List
+from rapidfireai.utils.exceptions import AutoMLException
+
+
+def recursive_expand_gridsearch(item: Any):
+    """Recursively expand nested structures with List datatypes into all combinations."""
+    if isinstance(item, dict):
+        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))
+    elif isinstance(item, List):
+        for value in item.values:
+            yield from recursive_expand_gridsearch(value)
+    else:
+        yield item
+
+
+class RFGridSearch(AutoMLAlgorithm):
+    """Grid search algorithm that generates all hyperparameter combinations."""
+
+    def get_runs(self, seed: int) -> ListType[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")
+
+        try:
+            runs = []
+            for config in self.configs:
+                if config.peft_config is None:
+                    peft_configs = [None]
+                elif isinstance(config.peft_config, List):
+                    peft_configs = config.peft_config.values
+                elif isinstance(config.peft_config, list):
+                    peft_configs = config.peft_config
+                else:
+                    peft_configs = [config.peft_config]
+
+                for peft_config in peft_configs:
+                    peft_instances = (
+                        [{}] if peft_config is None else list(recursive_expand_gridsearch(peft_config._user_params))
+                    )
+                    training_instances = (
+                        [{}]
+                        if config.training_args is None
+                        else list(recursive_expand_gridsearch(config.training_args._user_params))
+                    )
+                    model_kwargs_instances = (
+                        [{}] if config.model_kwargs is None else list(recursive_expand_gridsearch(config.model_kwargs))
+                    )
+                    ref_model_kwargs_instances = (
+                        [{}]
+                        if config.ref_model_kwargs is None
+                        else list(recursive_expand_gridsearch(config.ref_model_kwargs))
+                    )
+                    reward_funcs_instances = (
+                        [{}] if config.reward_funcs is None else list(recursive_expand_gridsearch(config.reward_funcs))
+                    )
+
+                    # Get additional kwargs for Trainer
+                    # FIXME: this is a hack to get the additional kwargs, we should find a better way to do this
+                    excluded_attrs = {
+                        "model_name",
+                        "tokenizer",
+                        "tokenizer_kwargs",
+                        "model_type",
+                        "model_kwargs",
+                        "peft_config",
+                        "training_args",
+                        "ref_model_name",
+                        "ref_model_type",
+                        "ref_model_kwargs",
+                        "reward_funcs",
+                    }
+                    # excluded_attrs = set(config.__dict__.keys()) - set(config.__annotations__.keys())
+                    additional_kwargs = {
+                        k: v for k, v in config.__dict__.items() if k not in excluded_attrs and v is not None
+                    }
+                    additional_kwargs_instances = (
+                        [{}] if not additional_kwargs else list(recursive_expand_gridsearch(additional_kwargs))
+                    )
+
+                    # Generate gridsearch combinations
+                    for peft_params in peft_instances:
+                        for training_params in training_instances:
+                            for model_kwargs in model_kwargs_instances:
+                                for additional_kwargs in additional_kwargs_instances:
+                                    leaf = {
+                                        "trainer_type": self.trainer_type,
+                                        "training_args": training_params,
+                                        "peft_params": peft_params,
+                                        "model_name": config.model_name,
+                                        "tokenizer": config.tokenizer,
+                                        "tokenizer_kwargs": config.tokenizer_kwargs,
+                                        "model_type": config.model_type,
+                                        "model_kwargs": model_kwargs,
+                                        "additional_kwargs": additional_kwargs,
+                                    }
+
+                                    if self.trainer_type == "DPO":
+                                        leaf["ref_model_config"] = {
+                                            "model_name": config.ref_model_name,
+                                            "model_type": config.ref_model_type,
+                                        }
+                                        for ref_model_kwargs in ref_model_kwargs_instances:
+                                            leaf["ref_model_config"]["model_kwargs"] = ref_model_kwargs
+                                            runs.append(leaf)
+                                    elif self.trainer_type == "GRPO":
+                                        for reward_func in reward_funcs_instances:
+                                            leaf["reward_funcs"] = reward_func
+                                            runs.append(leaf)
+                                    else:
+                                        runs.append(leaf)
+
+            return runs
+
+        except Exception as e:
+            raise AutoMLException(f"Error generating runs: {e}") from e

rapidfireai/automl/model_config.py

@@ -0,0 +1,102 @@
+"""Model configuration for AutoML training."""
+
+import inspect
+import copy
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Type, Union, 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):
+    """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)}")
+
+    sig = inspect.signature(base_class.__init__)
+    constructor_params = [p for p in sig.parameters.keys() if p != "self"]
+
+    type_hints = get_type_hints(base_class)
+    new_type_hints = {}
+
+    for param_name, param_type in type_hints.items():
+        if param_name in constructor_params:
+            new_type_hints[param_name] = param_type | List | Range
+
+    def __init__(self, **kwargs):
+        self._user_params = copy.deepcopy(kwargs)
+        self._constructor_params = constructor_params
+        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  
+    def copy_config(self):
+        """Create a deep copy of the configuration."""
+        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
+            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, 
+            "__init__": __init__,
+            "copy": copy_config,
+            "__setattr__": __setattr__
+        },
+    )
+
+
+# Create RF wrapper classes for external libraries
+RFLoraConfig = _create_rf_class(LoraConfig, "RFLoraConfig")
+RFSFTConfig = _create_rf_class(SFTConfig, "RFSFTConfig")
+RFDPOConfig = _create_rf_class(DPOConfig, "RFDPOConfig")
+RFGRPOConfig = _create_rf_class(GRPOConfig, "RFGRPOConfig")
+
+
+@dataclass
+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
+        """Create a deep copy of the RFModelConfig."""
+        return copy.deepcopy(self)

rapidfireai/automl/random_search.py

@@ -0,0 +1,145 @@
+"""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 rapidfireai.automl.base import AutoMLAlgorithm
+from rapidfireai.automl.datatypes import List, Range
+from rapidfireai.utils.exceptions import AutoMLException
+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):
+        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]]:
+        """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 
+            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):
+                    selected_peft_config = List(config.peft_config).sample()
+                elif isinstance(config.peft_config, List):
+                    selected_peft_config = config.peft_config.sample()
+                else:
+                    selected_peft_config = config.peft_config
+
+                peft_params = (
+                    {} 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 
+                    else recursive_expand_randomsearch(config.training_args._user_params)
+                )
+                
+                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)
+                )
+
+                # FIXME:  avoid hardcoding the excluded attributes
+                excluded_attrs = {
+                    "model_name",
+                    "tokenizer",
+                    "tokenizer_kwargs",
+                    "model_type",
+                    "model_kwargs",
+                    "peft_config",
+                    "training_args",
+                    "ref_model_name",
+                    "ref_model_type",
+                    "ref_model_kwargs",
+                    "reward_funcs",
+                }
+                additional_kwargs = {
+                    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)
+                )
+
+                leaf = {
+                    "trainer_type": self.trainer_type,
+                    "training_args": training_params,
+                    "peft_params": peft_params,
+                    "model_name": config.model_name,
+                    "tokenizer": config.tokenizer,
+                    "tokenizer_kwargs": config.tokenizer_kwargs,
+                    "model_type": config.model_type,
+                    "model_kwargs": model_kwargs,
+                    "additional_kwargs": additional_kwargs_sampled,
+                }
+
+                if self.trainer_type == "DPO":
+                    leaf["ref_model_config"] = {
+                        "model_name": config.ref_model_name,
+                        "model_type": config.ref_model_type,
+                        "model_kwargs": ref_model_kwargs,
+                    }
+                    #FIXME: correct ref args
+                elif self.trainer_type == "GRPO":
+                    leaf["reward_funcs"] = reward_funcs
+
+                # Check for duplicates using hashable representation
+                config_hash = encode_payload(leaf)
+                if config_hash not in seen_configs:
+                    seen_configs.add(config_hash)
+                    runs.append(leaf)
+
+            if len(runs) < self.num_runs:
+                raise AutoMLException(
+                    f"Could not generate {self.num_runs} unique configurations. "
+                    f"Generated {len(runs)} unique configs after {attempts} attempts. "
+                )
+
+            return runs
+
+        except Exception as e:
+            raise AutoMLException(f"Error generating runs: {e}") from e

rapidfireai/backend/chunks.py

@@ -0,0 +1,63 @@
+"""This module contains the DatasetChunker class which is responsible for chunking a PyTorch Dataset
+into chunks for distributed processing."""
+
+from datasets import Dataset
+
+
+class DatasetChunks:
+    """Chunks a HuggingFace Dataset into n_chunks for distributed processing."""
+
+    def __init__(self, dataset: Dataset, n_chunks: int):
+        self.dataset = dataset
+        self.n_chunks = n_chunks
+        self.dataset_size = len(dataset)
+
+        # Validate n_chunks
+        if n_chunks <= 0:
+            raise ValueError(f"n_chunks must be positive, got {n_chunks}")
+
+        # Calculate base size for even distribution (not chunk_size anymore)
+        self.base_size = self.dataset_size // n_chunks
+        self.extra_items = self.dataset_size % n_chunks
+        self.chunk_indices = self._create_chunk_indices()
+
+    def _create_chunk_indices(self):
+        """Create start/end index pairs for each chunk, distributing items as evenly as possible."""
+        chunks = {}
+
+        # Calculate base size and number of chunks that get an extra item
+        base_size = self.dataset_size // self.n_chunks
+        extra_items = self.dataset_size % self.n_chunks
+
+        current_idx = 0
+        for chunk_id in range(self.n_chunks):
+            # First 'extra_items' chunks get base_size + 1, rest get base_size
+            chunk_size = base_size + (1 if chunk_id < extra_items else 0)
+
+            if chunk_size > 0:  # Only create non-empty chunks
+                chunks[chunk_id] = (current_idx, current_idx + chunk_size)
+                current_idx += chunk_size
+
+        return chunks
+
+    def get_chunk(self, chunk_id: int) -> Dataset:
+        """Get a chunk as a HuggingFace Dataset subset."""
+        if chunk_id not in self.chunk_indices:
+            raise ValueError(f"Invalid chunk_id {chunk_id}. Valid range: 0-{len(self.chunk_indices) - 1}")
+
+        start_idx, end_idx = self.chunk_indices[chunk_id]
+        # Use HuggingFace Dataset's select method to create a subset
+        indices = list(range(start_idx, end_idx))
+        return self.dataset.select(indices)
+
+    def get_chunk_size(self, chunk_id: int) -> int:
+        """Get the size of a specific chunk."""
+        if chunk_id not in self.chunk_indices:
+            raise ValueError(f"Invalid chunk_id {chunk_id}")
+        start_idx, end_idx = self.chunk_indices[chunk_id]
+        return end_idx - start_idx
+
+    @property
+    def chunk_ids(self):
+        """Get all available chunk IDs."""
+        return list(self.chunk_indices.keys())