|
from typing import Any, Dict, List, Optional, Union |
|
|
|
import torch |
|
from transformers import PretrainedConfig |
|
|
|
|
|
class XLMRobertaFlashConfig(PretrainedConfig): |
|
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, |
|
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). |
|
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.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 |
|
|