Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2024 Databricks Mosaic Research and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""DBRX model configuration"""16 17from typing import Any, Optional18 19from ...configuration_utils import PretrainedConfig20from ...utils import logging21 22 23logger = logging.get_logger(__name__)24 25 26class DbrxAttentionConfig(PretrainedConfig):27    """Configuration class for Dbrx Attention.28 29    [`DbrxAttention`] class. It is used to instantiate attention layers30    according to the specified arguments, defining the layers architecture.31 32    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the33    documentation from [`PretrainedConfig`] for more information.34 35    Args:36        attn_pdrop (`float`, *optional*, defaults to 0.0):37            The dropout probability for the attention layers.38        clip_qkv (`float`, *optional*):39            If set, clip the queries, keys, and values in the attention layer to this value.40        kv_n_heads (`int`, *optional*, defaults to 1): For grouped_query_attention only, allow user to specify number of kv heads.41        rope_theta (`float`, *optional*, defaults to 10000.0): The base frequency for rope.42    """43 44    base_config_key = "attn_config"45 46    def __init__(47        self,48        attn_pdrop: float = 0.0,49        clip_qkv: Optional[float] = None,50        kv_n_heads: int = 1,51        rope_theta: float = 10000.0,52        **kwargs: Any,53    ):54        super().__init__(**kwargs)55        self.attn_pdrop = attn_pdrop56        self.clip_qkv = clip_qkv57        self.kv_n_heads = kv_n_heads58        self.rope_theta = rope_theta59 60        for k in ["model_type", "attn_implementation", "transformers_version", "_commit_hash", "torch_dtype", "dtype"]:61            if k in kwargs:62                kwargs.pop(k)63        if len(kwargs) != 0:64            raise ValueError(f"Found unknown {kwargs=}")65 66 67class DbrxFFNConfig(PretrainedConfig):68    """Configuration class for Dbrx FFN.69 70    [`DbrxFFN`] class. It is used to instantiate feedforward layers according to71    the specified arguments, defining the layers architecture.72 73    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the74    documentation from [`PretrainedConfig`] for more information.75 76    Args:77        ffn_act_fn (`dict`, *optional*, defaults to `None`): A dict specifying activation function for the FFN.78            The dict should have a key 'name' with the value being the name of the activation function along with79            any additional keyword arguments. If `None`, then set to `{"name": "silu"}`.80        ffn_hidden_size (`int`, *optional*, defaults to 3584): The hidden size of the feedforward network.81        moe_num_experts (`int`, *optional*, defaults to 4): The number of experts in the mixture of experts layer.82        moe_top_k (`int`, *optional*, defaults to 1): The number of experts to use in the mixture of experts layer.83        moe_jitter_eps (`float`, *optional*, defaults to `None`): If not `None`, the jitter epsilon for the mixture of experts layer.84        moe_loss_weight (`float`, *optional*, defaults to 0.01): The loss weight for the mixture of experts layer.85        moe_normalize_expert_weights (`float`, *optional*, defaults to 1.0): The normalization factor for the expert weights.86    """87 88    base_config_key = "ffn_config"89 90    def __init__(91        self,92        ffn_act_fn: Optional[dict] = None,93        ffn_hidden_size: int = 3584,94        moe_num_experts: int = 4,95        moe_top_k: int = 1,96        moe_jitter_eps: Optional[float] = None,97        moe_loss_weight: float = 0.01,98        moe_normalize_expert_weights: Optional[float] = 1.0,99        **kwargs: Any,100    ):101        super().__init__()102        if ffn_act_fn is None:103            ffn_act_fn = {"name": "silu"}104        self.ffn_act_fn = ffn_act_fn105        self.ffn_hidden_size = ffn_hidden_size106        self.moe_num_experts = moe_num_experts107        self.moe_top_k = moe_top_k108        self.moe_jitter_eps = moe_jitter_eps109        self.moe_loss_weight = moe_loss_weight110        self.moe_normalize_expert_weights = moe_normalize_expert_weights111 112        for k in ["model_type", "attn_implementation", "transformers_version", "_commit_hash", "torch_dtype", "dtype"]:113            if k in kwargs:114                kwargs.pop(k)115        if len(kwargs) != 0:116            raise ValueError(f"Found unknown {kwargs=}")117 118 119class DbrxConfig(PretrainedConfig):120    r"""121 122    This is the configuration class to store the configuration of a [`DbrxModel`]. It is used to instantiate a Dbrx model according to the123    specified arguments, defining the model architecture. Instantiating a configuration with the124    defaults will yield a different configuration to that of the [databricks/dbrx-instruct](https://huggingface.co/databricks/dbrx-instruct) architecture.125 126    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the127    documentation from [`PretrainedConfig`] for more information.128 129 130    Args:131        d_model (`int`, *optional*, defaults to 2048):132            Dimensionality of the embeddings and hidden states.133        n_heads (`int`, *optional*, defaults to 16):134            Number of attention heads for each attention layer in the Transformer encoder.135        n_layers (`int`, *optional*, defaults to 24):136            Number of hidden layers in the Transformer encoder.137        max_seq_len (`int`, *optional*, defaults to 2048):138            The maximum sequence length of the model.139        vocab_size (`int`, *optional*, defaults to 32000):140            Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by141            the `inputs_ids` passed when calling [`DbrxModel`].142        resid_pdrop (`float`, *optional*, defaults to 0.0):143            The dropout probability applied to the attention output before combining with residual.144        emb_pdrop (`float`, *optional*, defaults to 0.0):145            The dropout probability for the embedding layer.146        attn_config (`dict`, *optional*):147            A dictionary used to configure the model's attention module.148        ffn_config (`dict`, *optional*):149            A dictionary used to configure the model's FFN module.150        use_cache (`bool`, *optional*, defaults to `True`):151            Whether or not the model should return the last key/values attentions (not used by all models).152        initializer_range (`float`, *optional*, defaults to 0.02):153            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.154        output_router_logits (`bool`, *optional*, defaults to `False`):155            Whether or not the router logits should be returned by the model. Enabling this will also156            allow the model to output the auxiliary loss. See [here]() for more details.157 158 159    Example:160    ```python161    >>> from transformers import DbrxConfig, DbrxModel162 163    >>> # Initializing a Dbrx configuration164    >>> configuration = DbrxConfig(n_layers=2, d_model=256, n_heads=8, vocab_size=128)165 166    >>> # Initializing a model (with random weights) from the configuration167    >>> model = DbrxModel(configuration)168 169    >>> # Accessing the model configuration170    >>> configuration = model.config171    ```172    """173 174    model_type = "dbrx"175    sub_configs = {"attn_config": DbrxAttentionConfig, "ffn_config": DbrxFFNConfig}176    attribute_map = {177        "num_attention_heads": "n_heads",178        "hidden_size": "d_model",179        "num_hidden_layers": "n_layers",180        "max_position_embeddings": "max_seq_len",181    }182 183    def __init__(184        self,185        d_model: int = 2048,186        n_heads: int = 16,187        n_layers: int = 24,188        max_seq_len: int = 2048,189        vocab_size: int = 32000,190        resid_pdrop: float = 0.0,191        emb_pdrop: float = 0.0,192        attn_config: Optional[DbrxAttentionConfig] = None,193        ffn_config: Optional[DbrxFFNConfig] = None,194        use_cache: bool = True,195        initializer_range: float = 0.02,196        output_router_logits: bool = False,197        **kwargs: Any,198    ):199        if attn_config is None:200            self.attn_config = DbrxAttentionConfig()201        elif isinstance(attn_config, dict):202            self.attn_config = DbrxAttentionConfig(**attn_config)203        else:204            self.attn_config = attn_config205 206        if ffn_config is None:207            self.ffn_config = DbrxFFNConfig()208        elif isinstance(ffn_config, dict):209            self.ffn_config = DbrxFFNConfig(**ffn_config)210        else:211            self.ffn_config = ffn_config212 213        self.d_model = d_model214        self.n_heads = n_heads215        self.n_layers = n_layers216        self.max_seq_len = max_seq_len217        self.vocab_size = vocab_size218        self.resid_pdrop = resid_pdrop219        self.emb_pdrop = emb_pdrop220        self.use_cache = use_cache221        self.initializer_range = initializer_range222        self.output_router_logits = output_router_logits223        self.num_key_value_heads = self.attn_config.kv_n_heads224 225        tie_word_embeddings = kwargs.pop("tie_word_embeddings", False)226        if tie_word_embeddings:227            raise ValueError("tie_word_embeddings is not supported for DBRX models.")228 229        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)230 231 232__all__ = ["DbrxConfig"]233