| from typing import List, Optional |
| from transformers import PretrainedConfig |
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
| class BranchyModelConfig(PretrainedConfig): |
| """ |
| Configuration class for BranchyModel. This class extends the PretrainedConfig class from the Transformers |
| library, providing configuration specific to models with branch functionality. |
| |
| Attributes: |
| branch_locations (List[int]): Specifies the indices of layers after which branches are added. These indices |
| start from 0, and each index represents a layer in the underlying transformer model. |
| penalty_weight (Optional[float]): The weight of the penalty term used in the "penalized_cross_entropy" loss. |
| This parameter is required and must be greater than 0 |
| window_size (int): Determines the number of tokens each branch considers from the input sequence. This allows |
| for reducing the computational load by limiting the context size each branch processes. |
| |
| Example: |
| config = BranchyModelConfig( |
| branch_locations=[2, 4, 6], |
| window_size=512 |
| ) |
| |
| Note: |
| This configuration class is specifically designed for use with the BranchyModel class, enabling flexible |
| and customizable branching within transformer models. |
| """ |
| model_type = "branchy" |
|
|
| def __init__( |
| self, |
| model_str: str = None, |
| head_thresholds: Optional[List[float]] = None, |
| confidence_metric: Optional[str] = "breaking_ties", |
| branch_locations: Optional[List[int]] = None, |
| branch_number: Optional[int] = 3, |
| penalty_weight: Optional[float] = 0, |
| head_window_size: int = 512, |
| copy_lm_head: Optional[bool] = False, |
| **kwargs |
| ): |
| """ |
| Initializes the BranchyModelConfig. |
| |
| Args: |
| model_str (str): The model string to be used for the model. From Huggingface's model hub. |
| branch_locations (List[int], optional): Locations of the branches. Defaults to None, indicating no branches. |
| branch_number (Optional[int], optional): Number of branches if branch_locations is not provided. Defaults to 3. |
| penalty_weight (Optional[float], optional): Weight for the penalty in loss calculation. |
| . Defaults to None. |
| head_window_size (int, optional): Number of tokens each branch can see. Defaults to 512. |
| """ |
| self.model_str = model_str |
| self.head_thresholds = head_thresholds |
| self.confidence_metric = confidence_metric |
| assert self.confidence_metric in ["breaking_ties", "max"], "confidence_metric must be 'breaking_ties' or 'max'. It should depend on how you found the thresholds." |
| self.branch_locations = branch_locations |
| self.penalty_weight = penalty_weight |
| self.head_window_size = head_window_size |
| if branch_locations is not None and branch_number is not None: |
| logger.warning("Both branch_locations and branch_number are provided. Using branch_locations.") |
| self.branch_number = branch_number if branch_locations is None else len(branch_locations) |
| self.copy_lm_head = copy_lm_head |
| |
| assert self.branch_number > 0, "branch_number must be a positive integer." |
| assert isinstance(self.penalty_weight, float) or isinstance(self.penalty_weight, int), "penalty_weight must be a float or an integer." |
| assert self.penalty_weight >= 0 and self.penalty_weight <= 1, "penalty_weight must be in the range [0, 1]." |
| if branch_locations is not None: |
| assert all([isinstance(loc, int) for loc in self.branch_locations]), "Branch locations must be integers." |
| assert all([loc >= 0 for loc in self.branch_locations]), "Branch locations must be non-negative." |
| if self.head_window_size is not None: |
| assert self.head_window_size > 0 , "head_window_size must be a positive integer or None." |
| if type(self.head_thresholds) == list: |
| assert len(self.head_thresholds) == self.branch_number, "Number of thresholds must match number of branches." |
| assert all([isinstance(threshold, float) for threshold in self.head_thresholds]), "Thresholds must be floats." |
| super().__init__(**kwargs) |
