| from typing import Any, Dict, List, Optional, Union |
|
|
| import torch |
| from transformers import PretrainedConfig |
|
|
|
|
| class XLMRobertaFlashConfig(PretrainedConfig): |
|
|
| model_type = "xlm-roberta" |
|
|
| def __init__( |
| self, |
| vocab_size: int = 250002, |
| hidden_size: int = 1024, |
| num_hidden_layers: int = 24, |
| num_attention_heads: int = 16, |
| intermediate_size: int = 4096, |
| hidden_act: str = "gelu", |
| hidden_dropout_prob: float = 0.1, |
| attention_probs_dropout_prob: float = 0.1, |
| max_position_embeddings: int = 8194, |
| type_vocab_size: int = 1, |
| initializer_range: float = 0.02, |
| layer_norm_eps: float = 1e-05, |
| pad_token_id: int = 1, |
| bos_token_id: int = 0, |
| eos_token_id: int = 2, |
| position_embedding_type: str = "rotary", |
| rotary_emb_base: float = 10000.0, |
| use_cache: bool = True, |
| use_reentrant: bool = False, |
| classifier_dropout: Optional[float] = None, |
| lora_adaptations: Optional[List[str]] = None, |
| task_instructions: Optional[Dict[str, str]] = None, |
| lora_rank: int = 4, |
| lora_dropout_p: float = 0.0, |
| lora_alpha: int = 1, |
| lora_main_params_trainable: bool = False, |
| load_trained_adapters: bool = False, |
| use_flash_attn: bool = True, |
| torch_dtype: Optional[Union[str, torch.dtype]] = None, |
| emb_pooler: Optional[str] = None, |
| matryoshka_dimensions: Optional[List[int]] = None, |
| truncate_dim: Optional[int] = None, |
| **kwargs: Dict[str, Any], |
| ): |
| """ |
| Initialize the XLMRobertaFlashConfig configuration. |
| |
| Args: |
| vocab_size (int): Size of the vocabulary. |
| hidden_size (int): Dimensionality of the encoder layers and the pooler layer. |
| num_hidden_layers (int): Number of hidden layers in the Transformer encoder. |
| num_attention_heads (int): Number of attention heads for each attention layer in the Transformer encoder. |
| intermediate_size (int): Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer. |
| hidden_act (str): The activation function to use. |
| hidden_dropout_prob (float): The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. |
| attention_probs_dropout_prob (float): The dropout ratio for the attention probabilities. |
| max_position_embeddings (int): The maximum length of the position embeddings. |
| type_vocab_size (int): The vocabulary size of the token type ids. |
| initializer_range (float): The standard deviation for initializing all weight matrices. |
| layer_norm_eps (float): The epsilon used by the layer normalization layers. |
| pad_token_id (int): The ID of the padding token. |
| bos_token_id (int): The ID of the beginning-of-sequence token. |
| eos_token_id (int): The ID of the end-of-sequence token. |
| position_embedding_type (str): Type of position embeddings. Options are 'absolute', 'alibi', or 'rotary'. |
| rotary_emb_base (float): Base for rotary embeddings. |
| use_cache (bool): Whether or not the model should return the last key/values attentions (not used by all models). |
| use_reentrant (bool): Whether or not the model should enable the 'use_reentrant' flag in gradient checkpointing. |
| classifier_dropout (Optional[float]): The dropout ratio for the classification head. |
| lora_adaptations (Optional[List[str]]): LoRA adaptations configuration. |
| lora_prompts (Optional[Dict[str, str]]): LoRA prompts configuration. |
| lora_rank (int): Rank for LoRA adaptations. |
| lora_dropout_p (float): Dropout probability for LoRA adaptations. |
| lora_alpha (int): Alpha parameter for LoRA. |
| lora_main_params_trainable (bool): Whether to make the main model parameters trainable when using LoRA. |
| load_trained_adapters (bool): Whether to load trained adapters. |
| use_flash_attn (bool): Whether to use FlashAttention. |
| torch_dtype (Optional[Union[str, torch.dtype]]): Data type for the tensors. |
| emb_pooler (Optional[str]): Pooling layer configuration. |
| matryoshka_dimensions (Optional[List[int]]): Configuration for matryoshka dimension reduction. |
| truncate_dim (Optional[int]): Dimension to truncate embeddings to, if any. |
| **kwargs (Dict[str, Any]): Additional keyword arguments passed to the configuration. |
| """ |
|
|
| super().__init__( |
| pad_token_id=pad_token_id, |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| **kwargs, |
| ) |
|
|
| self.vocab_size = vocab_size |
| self.hidden_size = hidden_size |
| self.num_hidden_layers = num_hidden_layers |
| self.num_attention_heads = num_attention_heads |
| self.hidden_act = hidden_act |
| self.intermediate_size = intermediate_size |
| self.hidden_dropout_prob = hidden_dropout_prob |
| self.attention_probs_dropout_prob = attention_probs_dropout_prob |
| self.max_position_embeddings = max_position_embeddings |
| self.type_vocab_size = type_vocab_size |
| self.initializer_range = initializer_range |
| self.layer_norm_eps = layer_norm_eps |
| self.position_embedding_type = position_embedding_type |
| self.rotary_emb_base = rotary_emb_base |
| self.use_cache = use_cache |
| self.use_reentrant = use_reentrant |
| self.classifier_dropout = classifier_dropout |
| self.load_trained_adapters = load_trained_adapters |
| self.lora_adaptations = lora_adaptations |
| self.task_instructions = task_instructions |
| self.lora_rank = lora_rank |
| self.lora_dropout_p = lora_dropout_p |
| self.lora_alpha = lora_alpha |
| self.lora_main_params_trainable = lora_main_params_trainable |
| self.use_flash_attn = use_flash_attn |
| self.emb_pooler = emb_pooler |
| self.matryoshka_dimensions = matryoshka_dimensions |
| self.truncate_dim = truncate_dim |
| if ( |
| torch_dtype |
| and hasattr(torch, torch_dtype) |
| and type(getattr(torch, torch_dtype)) is torch.dtype |
| ): |
| self.torch_dtype = getattr(torch, torch_dtype) |
| else: |
| self.torch_dtype = torch_dtype |
| if not self.use_flash_attn or not torch.cuda.is_available(): |
| self.torch_dtype = torch.float32 |
