Nattapong Tapachoom
Update activation function to swiglu and enhance model registration with GenerationMixin
f4dad28 | """ | |
| Modern Thai GPT Model for Hugging Face Transformers | |
| This file allows loading the custom ModernThaiGPT model on HF Hub | |
| """ | |
| import math | |
| from typing import Optional, Tuple, List, Dict | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.nn import CrossEntropyLoss | |
| from transformers import PreTrainedModel, PretrainedConfig | |
| from transformers.modeling_utils import find_pruneable_heads_and_indices, prune_linear_layer | |
| from transformers.generation import GenerationMixin | |
| from transformers.modeling_outputs import CausalLMOutputWithPast | |
| class ZombitThaiGPTConfig(PretrainedConfig): | |
| """Configuration class for Zombit Thai GPT model""" | |
| model_type = "ZombitThaiGPT" | |
| keys_to_ignore_at_inference = ["past_key_values"] | |
| def __init__( | |
| self, | |
| vocab_size=50000, | |
| n_positions=131072, # 128K context length | |
| n_embd=1024, | |
| n_layer=24, | |
| n_head=16, | |
| n_inner=None, | |
| activation_function="swiglu", | |
| resid_pdrop=0.1, | |
| embd_pdrop=0.1, | |
| attn_pdrop=0.1, | |
| layer_norm_epsilon=1e-5, | |
| initializer_range=0.02, | |
| gradient_checkpointing=False, | |
| use_cache=True, | |
| bos_token_id=1, | |
| eos_token_id=2, | |
| tie_word_embeddings=False, | |
| num_kv_heads=8, | |
| rope_theta=10000.0, | |
| rope_scaling=None, # For dynamic RoPE scaling | |
| **kwargs | |
| ): | |
| super().__init__( | |
| bos_token_id=bos_token_id, | |
| eos_token_id=eos_token_id, | |
| tie_word_embeddings=tie_word_embeddings, | |
| **kwargs | |
| ) | |
| self.vocab_size = vocab_size | |
| self.n_positions = n_positions | |
| self.n_embd = n_embd | |
| self.n_layer = n_layer | |
| self.n_head = n_head | |
| self.n_inner = n_inner | |
| self.activation_function = activation_function | |
| self.resid_pdrop = resid_pdrop | |
| self.embd_pdrop = embd_pdrop | |
| self.attn_pdrop = attn_pdrop | |
| self.layer_norm_epsilon = layer_norm_epsilon | |
| self.initializer_range = initializer_range | |
| self.gradient_checkpointing = gradient_checkpointing | |
| self.use_cache = use_cache | |
| self.num_kv_heads = num_kv_heads | |
| self.rope_theta = rope_theta | |
| self.rope_scaling = rope_scaling | |
| class RoPE(nn.Module): | |
| """Rotary Position Embedding with Dynamic Scaling for Long Contexts""" | |
| def __init__(self, dim: int, max_seq_len: int = 131072, rope_theta: float = 10000.0, rope_scaling: Optional[Dict] = None): | |
| super().__init__() | |
| self.dim = dim | |
| self.max_seq_len = max_seq_len | |
| self.rope_theta = rope_theta | |
| self.rope_scaling = rope_scaling | |
| # Initialize inv_freq based on scaling configuration | |
| if rope_scaling is not None: | |
| self._init_scaled_inv_freq() | |
| else: | |
| inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2).float() / dim)) | |
| self.register_buffer("inv_freq", inv_freq) | |
| self._cache: Optional[torch.Tensor] = None | |
| def _init_scaled_inv_freq(self): | |
| """Initialize inv_freq with dynamic scaling for long contexts""" | |
| if self.rope_scaling is None: | |
| inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, self.dim, 2).float() / self.dim)) | |
| self.register_buffer("inv_freq", inv_freq) | |
| return | |
| scaling_type = self.rope_scaling.get("type", "linear") | |
| scaling_factor = self.rope_scaling.get("factor", 1.0) | |
| if scaling_type == "linear": | |
| inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, self.dim, 2).float() / self.dim)) | |
| inv_freq = inv_freq / scaling_factor | |
| elif scaling_type == "dynamic": | |
| # Dynamic NTK scaling for better long context handling | |
| inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, self.dim, 2).float() / self.dim)) | |
| # Apply dynamic scaling based on sequence length | |
| inv_freq = inv_freq * (scaling_factor ** (torch.arange(0, self.dim, 2).float() / self.dim)) | |
| else: | |
| # Default linear scaling | |
| inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, self.dim, 2).float() / self.dim)) | |
| self.register_buffer("inv_freq", inv_freq) | |
| def _get_cos_sin(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: | |
| # Check if we need to update scaling for this sequence length | |
| if self.rope_scaling is not None and self.rope_scaling.get("type") == "dynamic": | |
| # Dynamic scaling based on sequence length | |
| if seq_len > 2048: # Apply scaling for longer sequences | |
| scaling_factor = (seq_len / 2048.0) ** 0.1 # Gentle scaling | |
| inv_freq_scaled = self.inv_freq * scaling_factor | |
| else: | |
| inv_freq_scaled = self.inv_freq | |
| else: | |
| inv_freq_scaled = self.inv_freq | |
| if self._cache is not None and self._cache.size(0) >= seq_len: | |
| return self._cache[:seq_len].to(device) | |
| t = torch.arange(seq_len, device=device).float() | |
| freqs = torch.einsum("i , j -> i j", t, inv_freq_scaled.to(device)) | |
| emb = torch.cat((freqs, freqs), dim=-1) | |
| cos_cached = emb.cos() | |
| sin_cached = emb.sin() | |
| self._cache = torch.stack([cos_cached, sin_cached], dim=0) | |
| return cos_cached, sin_cached | |
| def _apply_rotary_emb(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| if x.dim() == 3: | |
| batch, seq_len, dim = x.shape | |
| x1, x2 = x[..., ::2], x[..., 1::2] | |
| cos = cos[:seq_len, :dim//2].unsqueeze(0) | |
| sin = sin[:seq_len, :dim//2].unsqueeze(0) | |
| return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) | |
| elif x.dim() == 4: | |
| batch, num_heads, seq_len, head_dim = x.shape | |
| x1, x2 = x[..., ::2], x[..., 1::2] | |
| cos = cos[:seq_len, :head_dim//2].unsqueeze(0).unsqueeze(0) | |
| sin = sin[:seq_len, :head_dim//2].unsqueeze(0).unsqueeze(0) | |
| return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) | |
| return x | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| seq_len = x.size(-2) | |
| cos, sin = self._get_cos_sin(seq_len, x.device) | |
| return self._apply_rotary_emb(x, cos, sin) | |
| class RMSNorm(nn.Module): | |
| """Root Mean Square Layer Normalization""" | |
| def __init__(self, dim: int, eps: float = 1e-6): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps) | |
| return self.weight * (x / rms) | |
| class SwiGLU(nn.Module): | |
| """SwiGLU activation function""" | |
| def __init__(self, dim_in: int, dim_out: int): | |
| super().__init__() | |
| hidden_dim = dim_out | |
| self.w1 = nn.Linear(dim_in, hidden_dim) | |
| self.w2 = nn.Linear(dim_in, hidden_dim) | |
| self.w3 = nn.Linear(hidden_dim, dim_in) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.w3(F.silu(self.w1(x)) * self.w2(x)) | |
| class GroupedQueryAttention(nn.Module): | |
| """Grouped Query Attention""" | |
| def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, max_seq_len: int, dropout: float = 0.1, rope_theta: float = 10000.0, rope_scaling: Optional[Dict] = None): | |
| super().__init__() | |
| self.d_model = d_model | |
| self.num_heads = num_heads | |
| self.num_kv_heads = num_kv_heads | |
| self.d_k = d_model // num_heads | |
| self.heads_per_group = num_heads // num_kv_heads | |
| self.W_q = nn.Linear(d_model, d_model) | |
| self.W_k = nn.Linear(d_model, num_kv_heads * self.d_k) | |
| self.W_v = nn.Linear(d_model, num_kv_heads * self.d_k) | |
| self.W_o = nn.Linear(d_model, d_model) | |
| self.rope = RoPE(self.d_k, max_seq_len, rope_theta=rope_theta, rope_scaling=rope_scaling) | |
| self.dropout = nn.Dropout(dropout) | |
| self.kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None | |
| def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, use_cache: bool = False) -> torch.Tensor: | |
| batch_size, seq_len, _ = x.size() | |
| q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) | |
| k = self.W_k(x).view(batch_size, seq_len, self.num_kv_heads, self.d_k).transpose(1, 2) | |
| v = self.W_v(x).view(batch_size, seq_len, self.num_kv_heads, self.d_k).transpose(1, 2) | |
| q = self.rope(q.transpose(1, 2)).transpose(1, 2) | |
| k = self.rope(k.transpose(1, 2)).transpose(1, 2) | |
| if use_cache and self.kv_cache is not None: | |
| k_cache, v_cache = self.kv_cache | |
| k = torch.cat([k_cache, k], dim=2) | |
| v = torch.cat([v_cache, v], dim=2) | |
| self.kv_cache = (k, v) | |
| if use_cache and self.kv_cache is None: | |
| self.kv_cache = (k, v) | |
| if self.heads_per_group > 1: | |
| k = k.repeat_interleave(self.heads_per_group, dim=1) | |
| v = v.repeat_interleave(self.heads_per_group, dim=1) | |
| scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) | |
| if attention_mask is not None: | |
| # Handle attention_mask from transformers generation pipeline | |
| try: | |
| if attention_mask.dim() == 4: # [batch, 1, seq_len, seq_len] from transformers | |
| # Extract the actual mask and expand to match scores | |
| attention_mask = attention_mask.squeeze(1) # [batch, seq_len, seq_len] | |
| attention_mask = attention_mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1) | |
| elif attention_mask.dim() == 3: # [batch, seq_len, seq_len] from model | |
| # Check if dimensions match scores | |
| if attention_mask.size(-1) != scores.size(-1) or attention_mask.size(-2) != scores.size(-2): | |
| # Create a mask that matches the current sequence length | |
| current_seq_len = scores.size(-1) | |
| mask_seq_len = attention_mask.size(-1) | |
| if current_seq_len <= mask_seq_len: | |
| # Take the relevant portion of the mask | |
| attention_mask = attention_mask[..., :current_seq_len, :current_seq_len] | |
| else: | |
| # Pad the mask if needed (shouldn't happen in normal generation) | |
| padding = current_seq_len - mask_seq_len | |
| attention_mask = F.pad(attention_mask, (0, padding, 0, padding), value=float('-inf')) | |
| attention_mask = attention_mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1) | |
| elif attention_mask.dim() == 2: # [seq_len, seq_len] causal mask | |
| # Check if dimensions match scores | |
| current_seq_len = scores.size(-1) | |
| mask_seq_len = attention_mask.size(-1) | |
| if current_seq_len <= mask_seq_len: | |
| # Take the relevant portion of the mask | |
| attention_mask = attention_mask[:current_seq_len, :current_seq_len] | |
| else: | |
| # Pad the mask if needed | |
| padding = current_seq_len - mask_seq_len | |
| attention_mask = F.pad(attention_mask, (0, padding, 0, padding), value=float('-inf')) | |
| attention_mask = attention_mask.unsqueeze(0).unsqueeze(0).expand(batch_size, self.num_heads, -1, -1) | |
| scores = scores + attention_mask | |
| except RuntimeError as e: | |
| # If there's still a dimension mismatch, skip attention mask processing | |
| if "size of tensor" in str(e): | |
| print(f"⚠️ Skipping attention mask due to dimension mismatch: {e}") | |
| pass # Skip attention mask processing | |
| else: | |
| raise e | |
| attention = F.softmax(scores, dim=-1) | |
| attention = self.dropout(attention) | |
| context = torch.matmul(attention, v) | |
| # CRITICAL FIX: Proper handling of grouped query attention dimensions | |
| # context shape: [batch_size, num_heads, seq_len, d_k] | |
| # After grouped query attention, we have num_heads=16, d_k=64 | |
| # We need to get back to d_model=1024 = 16*64 | |
| # Transpose to [batch_size, seq_len, num_heads, d_k] | |
| context = context.transpose(1, 2) | |
| # The key insight: we should only use the first num_kv_heads worth of output | |
| # because that's what we actually computed, the rest is just repeated | |
| if self.heads_per_group > 1: | |
| # We repeated the KV heads, so we need to "un-repeat" the output | |
| # Take only the first group's worth of heads | |
| heads_to_keep = self.num_kv_heads * self.heads_per_group | |
| if context.size(2) > heads_to_keep: | |
| context = context[:, :, :heads_to_keep, :] | |
| # Now reshape to [batch_size, seq_len, d_model] | |
| # Since d_model = num_heads * d_k = 16 * 64 = 1024 | |
| final_dim = min(context.size(2) * context.size(3), self.d_model) | |
| context = context.contiguous().view(batch_size, seq_len, -1) | |
| # Ensure exactly d_model dimensions | |
| if context.size(-1) != self.d_model: | |
| if context.size(-1) > self.d_model: | |
| context = context[..., :self.d_model] | |
| else: | |
| padding = self.d_model - context.size(-1) | |
| context = F.pad(context, (0, padding)) | |
| output = self.W_o(context) | |
| return output | |
| def clear_cache(self): | |
| self.kv_cache = None | |
| class ModernTransformerBlock(nn.Module): | |
| """Enhanced Transformer Block""" | |
| def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, d_ff: int, max_seq_len: int, dropout: float = 0.1, rope_theta: float = 10000.0, rope_scaling: Optional[Dict] = None): | |
| super().__init__() | |
| self.attention = GroupedQueryAttention(d_model, num_heads, num_kv_heads, max_seq_len, dropout, rope_theta, rope_scaling) | |
| self.norm1 = RMSNorm(d_model) | |
| self.norm2 = RMSNorm(d_model) | |
| self.feed_forward = SwiGLU(d_model, d_ff) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, use_cache: bool = False) -> torch.Tensor: | |
| attn_output = self.attention(x, attention_mask, use_cache) | |
| x = self.norm1(x + self.dropout(attn_output)) | |
| ff_output = self.feed_forward(x) | |
| x = self.norm2(x + self.dropout(ff_output)) | |
| return x | |
| class ZombitThaiGPTPreTrainedModel(PreTrainedModel): | |
| """Base class for Zombit Thai GPT models""" | |
| config_class = ZombitThaiGPTConfig | |
| base_model_prefix = "transformer" | |
| _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.masked_bias", r"h\.\d+\.attn\.bias"] | |
| def __init__(self, *inputs, **kwargs): | |
| super().__init__(*inputs, **kwargs) | |
| def _init_weights(self, module): | |
| """Initialize the weights""" | |
| if isinstance(module, (nn.Linear, nn.Embedding)): | |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) | |
| if isinstance(module, nn.Linear) and module.bias is not None: | |
| module.bias.data.zero_() | |
| elif isinstance(module, nn.LayerNorm): | |
| module.bias.data.zero_() | |
| module.weight.data.fill_(1.0) | |
| class ZombitThaiGPTModel(ZombitThaiGPTPreTrainedModel): | |
| """Zombit Thai GPT Model""" | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.embed_dim = config.n_embd | |
| self.vocab_size = config.vocab_size | |
| self.wte = nn.Embedding(config.vocab_size, self.embed_dim) | |
| self.drop = nn.Dropout(config.embd_pdrop) | |
| self.gradient_checkpointing = config.gradient_checkpointing | |
| self.h = nn.ModuleList([ | |
| ModernTransformerBlock( | |
| self.embed_dim, | |
| config.n_head, | |
| config.num_kv_heads, | |
| config.n_inner or 4 * self.embed_dim, | |
| config.n_positions, | |
| config.resid_pdrop, | |
| config.rope_theta, | |
| config.rope_scaling | |
| ) | |
| for _ in range(config.n_layer) | |
| ]) | |
| self.ln_f = RMSNorm(self.embed_dim, eps=config.layer_norm_epsilon) | |
| self.init_weights() | |
| def get_input_embeddings(self): | |
| return self.wte | |
| def set_input_embeddings(self, new_embeddings): | |
| self.wte = new_embeddings | |
| def forward( | |
| self, | |
| input_ids: Optional[torch.Tensor] = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, | |
| inputs_embeds: Optional[torch.Tensor] = None, | |
| use_cache: Optional[bool] = None, | |
| output_attentions: Optional[bool] = None, | |
| output_hidden_states: Optional[bool] = None, | |
| return_dict: Optional[bool] = None, | |
| ): | |
| if input_ids is not None and inputs_embeds is not None: | |
| raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") | |
| elif input_ids is not None: | |
| input_shape = input_ids.size() | |
| input_ids = input_ids.view(-1, input_shape[-1]) | |
| batch_size = input_ids.shape[0] | |
| elif inputs_embeds is not None: | |
| input_shape = inputs_embeds.size()[:-1] | |
| batch_size = inputs_embeds.shape[0] | |
| else: | |
| raise ValueError("You have to specify either input_ids or inputs_embeds") | |
| device = input_ids.device if input_ids is not None else inputs_embeds.device | |
| if past_key_values is None or len(past_key_values) == 0: | |
| past_length = 0 | |
| past_key_values = tuple([None] * len(self.h)) | |
| else: | |
| # Check if past_key_values has valid content (not just None values) | |
| if past_key_values[0] is not None and isinstance(past_key_values[0], tuple) and len(past_key_values[0]) > 0 and past_key_values[0][0] is not None: | |
| past_length = past_key_values[0][0].size(-2) | |
| else: | |
| past_length = 0 | |
| past_key_values = tuple([None] * len(self.h)) | |
| if inputs_embeds is None: | |
| inputs_embeds = self.wte(input_ids) | |
| hidden_states = self.drop(inputs_embeds) | |
| # Create causal mask accounting for past key values | |
| seq_length = input_shape[-1] | |
| # Total sequence length includes past context | |
| total_seq_length = seq_length + past_length | |
| causal_mask = self._create_causal_mask(total_seq_length, device) | |
| if attention_mask is not None: | |
| # attention_mask from transformers is [batch_size, seq_len] where 1=attend, 0=don't attend | |
| # Convert to causal mask format [seq_len, seq_len] with proper broadcasting | |
| attention_mask = attention_mask.to(dtype=hidden_states.dtype) | |
| # Pad attention_mask to match total sequence length if needed | |
| if attention_mask.size(-1) < total_seq_length: | |
| padding = total_seq_length - attention_mask.size(-1) | |
| attention_mask = F.pad(attention_mask, (0, padding), value=1.0) | |
| # Expand attention_mask to [batch_size, total_seq_len, total_seq_len] | |
| attention_mask = attention_mask.unsqueeze(-1).expand(-1, -1, total_seq_length) | |
| # Apply to causal mask: where attention_mask is 0, set to -inf | |
| causal_mask = attention_mask + causal_mask.unsqueeze(0) | |
| causal_mask = torch.where(causal_mask == 0, torch.tensor(float('-inf'), device=device, dtype=hidden_states.dtype), causal_mask) | |
| # For generation, we only need the mask for the current sequence | |
| if past_length > 0: | |
| causal_mask = causal_mask[..., past_length:, :] | |
| presents = () if use_cache else None | |
| all_self_attentions = () if output_attentions else None | |
| all_hidden_states = () if output_hidden_states else None | |
| for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): | |
| if output_hidden_states: | |
| all_hidden_states = all_hidden_states + (hidden_states,) | |
| if self.gradient_checkpointing and self.training: | |
| if use_cache: | |
| logger.warning("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...") | |
| use_cache = False | |
| def create_custom_forward(module): | |
| def custom_forward(*inputs): | |
| return module(*inputs, use_cache, output_attentions) | |
| return custom_forward | |
| outputs = torch.utils.checkpoint.checkpoint( | |
| create_custom_forward(block), | |
| hidden_states, | |
| causal_mask, | |
| ) | |
| else: | |
| outputs = block( | |
| hidden_states, | |
| attention_mask=causal_mask, | |
| use_cache=use_cache, | |
| ) | |
| hidden_states = outputs | |
| if use_cache: | |
| presents = presents + (outputs[1:],) | |
| hidden_states = self.ln_f(hidden_states) | |
| hidden_states = hidden_states.view(*input_shape, -1) | |
| if output_hidden_states: | |
| all_hidden_states = all_hidden_states + (hidden_states,) | |
| if not return_dict: | |
| return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) | |
| return { | |
| "last_hidden_state": hidden_states, | |
| "past_key_values": presents, | |
| "hidden_states": all_hidden_states, | |
| "attentions": all_self_attentions, | |
| } | |
| def _create_causal_mask(self, seq_length: int, device: torch.device) -> torch.Tensor: | |
| """Create causal mask for autoregressive generation""" | |
| causal_mask = torch.triu(torch.ones(seq_length, seq_length, device=device), diagonal=1) | |
| causal_mask = causal_mask.masked_fill(causal_mask == 1, float('-inf')) | |
| return causal_mask | |
| class ZombitThaiGPTLMHeadModel(ZombitThaiGPTPreTrainedModel, GenerationMixin): | |
| """Zombit Thai GPT Language Model with LM Head""" | |
| _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.masked_bias", r"h\.\d+\.attn\.bias", r"lm_head\.weight"] | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.transformer = ZombitThaiGPTModel(config) | |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) | |
| # Tie weights | |
| self.lm_head.weight = self.transformer.wte.weight | |
| self.init_weights() | |
| def get_output_embeddings(self): | |
| return self.lm_head | |
| def set_output_embeddings(self, new_embeddings): | |
| self.lm_head = new_embeddings | |
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs): | |
| if past_key_values: | |
| input_ids = input_ids[:, -1:] | |
| attention_mask = kwargs.get("attention_mask", None) | |
| if attention_mask is not None and past_key_values is not None: | |
| attention_mask = attention_mask[:, -input_ids.shape[1]:] | |
| if inputs_embeds is not None and past_key_values is not None: | |
| model_inputs = {"inputs_embeds": inputs_embeds} | |
| else: | |
| model_inputs = {"input_ids": input_ids} | |
| model_inputs.update({ | |
| "past_key_values": past_key_values, | |
| "use_cache": kwargs.get("use_cache"), | |
| "attention_mask": attention_mask, | |
| }) | |
| return model_inputs | |
| def forward( | |
| self, | |
| input_ids: Optional[torch.Tensor] = None, | |
| past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| inputs_embeds: Optional[torch.Tensor] = None, | |
| labels: Optional[torch.Tensor] = None, | |
| use_cache: Optional[bool] = None, | |
| output_attentions: Optional[bool] = None, | |
| output_hidden_states: Optional[bool] = None, | |
| return_dict: Optional[bool] = None, | |
| ): | |
| r""" | |
| labels (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*): | |
| Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set | |
| `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` | |
| are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` | |
| """ | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
| transformer_outputs = self.transformer( | |
| input_ids, | |
| past_key_values=past_key_values, | |
| attention_mask=attention_mask, | |
| inputs_embeds=inputs_embeds, | |
| use_cache=use_cache, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict, | |
| ) | |
| hidden_states = transformer_outputs["last_hidden_state"] | |
| lm_logits = self.lm_head(hidden_states) | |
| loss = None | |
| if labels is not None: | |
| # Shift so that tokens < n predict n | |
| shift_logits = lm_logits[..., :-1, :].contiguous() | |
| shift_labels = labels[..., 1:].contiguous() | |
| # Flatten the tokens | |
| loss_fct = CrossEntropyLoss() | |
| loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) | |
| if not return_dict: | |
| output = (lm_logits,) + transformer_outputs[1:] | |
| return ((loss,) + output) if loss is not None else output | |
| return CausalLMOutputWithPast( | |
| loss=loss, | |
| logits=lm_logits, | |
| past_key_values=transformer_outputs.get("past_key_values"), | |
| hidden_states=transformer_outputs.get("hidden_states"), | |
| attentions=transformer_outputs.get("attentions"), | |
| ) | |
| def _reorder_cache(past_key_values, beam_idx): | |
| reordered_past = () | |
| for layer_past in past_key_values: | |
| reordered_past += ( | |
| tuple(past_state.index_select(0, beam_idx) for past_state in layer_past), | |
| ) | |
| return reordered_past |