| import math | |
| from typing import Optional, Tuple, Union | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import PreTrainedModel | |
| from transformers.modeling_outputs import MaskedLMOutput | |
| from .configuration_vminibert import VMiniBertConfig | |
| class RMSNorm(nn.Module): | |
| 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: | |
| norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return norm * self.weight | |
| class RotaryEmbedding(nn.Module): | |
| def __init__(self, dim: int, max_seq_len: int = 512, base: float = 10000.0): | |
| super().__init__() | |
| self.dim = dim | |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) | |
| self.register_buffer("inv_freq", inv_freq, persistent=False) | |
| def forward(self, x: torch.Tensor, seq_len: int): | |
| t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype) | |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq) | |
| emb = torch.cat((freqs, freqs), dim=-1) | |
| cos = emb.cos()[None, None, :, :] | |
| sin = emb.sin()[None, None, :, :] | |
| return cos, sin | |
| def rotate_half(x: torch.Tensor) -> torch.Tensor: | |
| x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:] | |
| return torch.cat((-x2, x1), dim=-1) | |
| def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): | |
| q_rot = (q * cos) + (rotate_half(q) * sin) | |
| k_rot = (k * cos) + (rotate_half(k) * sin) | |
| return q_rot, k_rot | |
| class VectorizedMultiHeadAttention(nn.Module): | |
| def __init__(self, config: VMiniBertConfig): | |
| super().__init__() | |
| self.num_heads = config.num_attention_heads | |
| self.head_dim = config.hidden_size // config.num_attention_heads | |
| self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) | |
| self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) | |
| self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) | |
| self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) | |
| self.dropout = nn.Dropout(config.attention_probs_dropout_prob) | |
| self.rotary_emb = RotaryEmbedding(self.head_dim, max_seq_len=config.max_position_embeddings) | |
| def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: | |
| B, S, H = x.shape | |
| Q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) | |
| K = self.k_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) | |
| V = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) | |
| cos, sin = self.rotary_emb(Q, S) | |
| Q, K = apply_rotary_pos_emb(Q, K, cos, sin) | |
| scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) | |
| if mask is not None: | |
| scores = scores.masked_fill(mask == 0, -1e4) | |
| attn_weights = self.dropout(F.softmax(scores, dim=-1)) | |
| out = torch.matmul(attn_weights, V).transpose(1, 2).contiguous().view(B, S, H) | |
| return self.out_proj(out) | |
| class GeGLU(nn.Module): | |
| def __init__(self, config: VMiniBertConfig): | |
| super().__init__() | |
| self.w_gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) | |
| self.w_up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) | |
| self.w_down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) | |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| gate = F.gelu(self.w_gate(x)) | |
| up = self.w_up(x) | |
| hidden = self.dropout(gate * up) | |
| return self.w_down(hidden) | |
| class EncoderLayer(nn.Module): | |
| def __init__(self, config: VMiniBertConfig): | |
| super().__init__() | |
| self.norm1 = RMSNorm(config.hidden_size, eps=config.layer_norm_eps) | |
| self.mha = VectorizedMultiHeadAttention(config) | |
| self.norm2 = RMSNorm(config.hidden_size, eps=config.layer_norm_eps) | |
| self.ffn = GeGLU(config) | |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) | |
| def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: | |
| normed_x = self.norm1(x) | |
| attn_out = self.mha(normed_x, mask=mask) | |
| x = x + self.dropout(attn_out) | |
| normed_x = self.norm2(x) | |
| ffn_out = self.ffn(normed_x) | |
| x = x + self.dropout(ffn_out) | |
| return x | |
| class VMiniBertForMaskedLM(PreTrainedModel): | |
| config_class = VMiniBertConfig | |
| base_model_prefix = "vminibert" | |
| _tied_weights_keys = {"mlm_head.weight": "token_embedding.weight"} | |
| def __init__(self, config: VMiniBertConfig): | |
| super().__init__(config) | |
| self.config = config | |
| self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) | |
| self.embed_norm = RMSNorm(config.hidden_size, eps=config.layer_norm_eps) | |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) | |
| self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.num_hidden_layers)]) | |
| self.final_norm = RMSNorm(config.hidden_size, eps=config.layer_norm_eps) | |
| self.transform = nn.Sequential( | |
| nn.Linear(config.hidden_size, config.hidden_size), | |
| nn.GELU(), | |
| RMSNorm(config.hidden_size, eps=config.layer_norm_eps) | |
| ) | |
| self.mlm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True) | |
| self.mlm_head.weight = self.token_embedding.weight | |
| self.post_init() | |
| def get_input_embeddings(self) -> nn.Module: | |
| return self.token_embedding | |
| def set_input_embeddings(self, value: nn.Module): | |
| self.token_embedding = value | |
| def get_output_embeddings(self) -> nn.Module: | |
| return self.mlm_head | |
| def set_output_embeddings(self, new_embeddings: nn.Module): | |
| self.mlm_head = new_embeddings | |
| def forward( | |
| self, | |
| input_ids: Optional[torch.Tensor] = None, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| labels: Optional[torch.Tensor] = None, | |
| return_dict: Optional[bool] = None, | |
| **kwargs | |
| ) -> Union[Tuple[torch.Tensor, ...], MaskedLMOutput]: | |
| return_dict = return_dict if return_dict is not None else True | |
| tok_embeds = self.token_embedding(input_ids) | |
| x = self.dropout(self.embed_norm(tok_embeds)) | |
| if attention_mask is not None: | |
| mask = attention_mask.unsqueeze(1).unsqueeze(2) | |
| else: | |
| mask = (input_ids != self.config.pad_token_id).unsqueeze(1).unsqueeze(2) | |
| for layer in self.layers: | |
| x = layer(x, mask=mask) | |
| x = self.final_norm(x) | |
| x = self.transform(x) | |
| logits = self.mlm_head(x) | |
| loss = None | |
| if labels is not None: | |
| loss_fct = nn.CrossEntropyLoss(ignore_index=-100) | |
| loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1)) | |
| if not return_dict: | |
| output = (logits,) | |
| return ((loss,) + output) if loss is not None else output | |
| return MaskedLMOutput( | |
| loss=loss, | |
| logits=logits, | |
| hidden_states=None, | |
| attentions=None, | |
| ) | |