| 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 |