import math import os from dataclasses import dataclass from typing import Optional, Tuple, Union import torch import torch.utils.checkpoint from torch import nn from torch.nn import CrossEntropyLoss, BCEWithLogitsLoss, MSELoss from transformers.modeling_outputs import ( BaseModelOutput, CausalLMOutput, SequenceClassifierOutput ) from transformers.modeling_utils import PreTrainedModel from transformers.utils import logging from .rita_configuration import RITAConfig import torch.nn.functional as F logger = logging.get_logger(__name__) @torch.jit.script def RITA_gelu(hidden_states): return hidden_states * 0.5 * (1.0 + torch.tanh(0.79788456 * hidden_states * (1 + 0.044715 * hidden_states * hidden_states))) class RITAGELU(nn.Module): def __init__(self): super().__init__() def forward(self, hidden_states): return RITA_gelu(hidden_states) def rotate_half(x): x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=x1.ndim - 1) class RotaryEmbedding(nn.Module): def __init__(self, config): super().__init__() assert config.d_model % config.num_heads == 0 self.d_model = config.d_model self.num_heads = config.num_heads self.max_seq_len = config.max_seq_len head_dim = self.d_model // self.num_heads inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim)) self.register_buffer('inv_freq', inv_freq) self.seq_len_cached = None self.cos_cached = None self.sin_cached = None def forward(self, x: torch.FloatTensor, seq_dim=1) -> torch.FloatTensor: seq_len = x.shape[seq_dim] if seq_len != self.seq_len_cached: self.seq_len_cached = seq_len t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1).to(x.device) self.cos_cached = emb.cos()[None, None, :, :] self.sin_cached = emb.sin()[None, None, :, :] return self.cos_cached, self.sin_cached def apply_rotary_pos_emb(self, q, k, cos, sin): return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) class SelfAttention(nn.Module): """Implementation of MultiHeadAttention following `Karpathy's MinGPT `_. modified to use rotary embeddings. Parameters ---------- d_model: int, total dimension of the model. num_heads: int, number of parallel attention heads. num_layers: int, number of layers in the model, used for the Megatron-like init. rotaty_embedding: Optional[Block], default None, a RotaryEmbedding Block to add positionnal information in Queries and Keys dropout: float, default 0.1, amount of dropout on the attention weights. sigma: float, default 0.02, standard deviation used for the init. trainable: bool, default True, if False, the Module parameters will be hidden from the optimizer. """ def __init__( self, d_model: int, num_heads: int, num_layers: int, rotary_embedding= None, dropout: float = 0.1, sigma=0.02, use_cache: bool = False, bias=True, ): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.head_dim = self.d_model // self.num_heads self.num_layers = num_layers self.dropout = dropout self.sigma = sigma self.bias = bias # key, query, value projections for all heads self.key = nn.Linear(d_model, d_model, bias=bias) self.query = nn.Linear(d_model, d_model, bias=bias) self.value = nn.Linear(d_model, d_model, bias=bias) # regularization self.attn_drop = nn.Dropout(dropout) self.resid_drop = nn.Dropout(dropout) # output projection self.proj = nn.Linear(d_model, d_model, bias=bias) self.rotary_embedding = rotary_embedding self.layer_id = None # will be set by the Transformer itself self.use_cache = use_cache self.qkv = None self.bias = bias def forward( self, x, attn_mask: Optional[torch.BoolTensor] = None, padding_mask: Optional[torch.BoolTensor] = None, ) -> Tuple[torch.FloatTensor, torch.FloatTensor]: N, L, D = x.size() # Batch_size, Context_size, d_model # calculate query, key, values for all heads in batch and move head forward to be the batch dim k = ( self.key(x).view(N, L, self.num_heads, D // self.num_heads).transpose(1, 2) ) # (N, nh, L, hs) q = ( self.query(x).view(N, L, self.num_heads, D // self.num_heads).transpose(1, 2) ) # (N, nh, L, hs) v = ( self.value(x).view(N, L, self.num_heads, D // self.num_heads).transpose(1, 2) ) # (N, nh, L, hs) if self.rotary_embedding is not None: cos, sin = self.rotary_embedding(x) q, k = self.rotary_embedding.apply_rotary_pos_emb(q, k, cos, sin) # causal self-attention; Self-attend: (N, nh, L, hs) x (N, nh, hs, L) -> (N, nh, L, L) att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) if attn_mask is not None: att[:,:,-L:, -L: ].masked_fill_(attn_mask.view(1, 1, L, L), float("-inf")) att = ( att.transpose(0, 2) .masked_fill(padding_mask.view(1, 1, N, L), float("-inf")) .transpose(0, 2) if padding_mask is not None else att ) att = F.softmax(att, dim=-1) att = self.attn_drop(att) y = att @ v # (N, nh, L, L) x (N, nh, L, hs) -> (N, nh, L, hs) y = ( y.transpose(1, 2).contiguous().view(N, L, D) ) # re-assemble all head outputs side by side # output projection y = self.resid_drop(self.proj(y)) return y class DecoderLayer(nn.Module): """Transformer block containing the self-attention module and the feedfoward module.""" def __init__( self, config ): super().__init__() self.self_attention = SelfAttention(config.d_model, config.num_heads, config.dropout, rotary_embedding=RotaryEmbedding(config)) self.attn_norm = nn.LayerNorm(config.d_model) self.attn_dropout = nn.Dropout(config.dropout) self.mlp = nn.Sequential( nn.Linear(config.d_model, config.d_feedforward, bias=True), RITAGELU(), nn.Linear(config.d_feedforward, config.d_model, bias=True), ) self.mlp_norm = nn.LayerNorm(config.d_model) self.mlp_dropout = nn.Dropout(config.dropout) def forward( self, x: torch.FloatTensor, attn_mask: torch.BoolTensor, padding_mask: Optional[torch.BoolTensor] = None, ) -> torch.FloatTensor: y = self.attn_norm(x) y = self.self_attention(y, attn_mask=attn_mask, padding_mask=padding_mask) x = x + self.attn_dropout(y) y = self.mlp_norm(x) y = self.mlp(y) x = x + self.mlp_dropout(y) return x class RITAModel(PreTrainedModel): config_class = RITAConfig base_model_prefix = "transformer" is_parallelizable = False def __init__( self, config ): super().__init__(config) self.embedding = nn.Embedding(config.vocab_size, config.d_model) self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_layers)]) self.final_norm = nn.LayerNorm(config.d_model) def forward( self, input_ids=None, past_key_values=None, # NOT USED attention_mask=None, token_type_ids=None, # NOT USED position_ids=None, # NOT USED head_mask=None, # NOT USED inputs_embeds=None, encoder_hidden_states=None, # NOT USED encoder_attention_mask=None, # NOT USED labels=None, use_cache=None, # NOT USED output_attentions=None, # NOT USED output_hidden_states=None, # NOT USED return_dict=None # NOT USED ) -> torch.FloatTensor: if inputs_embeds == None: x = self.embedding(input_ids) # N x L x D else: x = inputs_embeds if attention_mask == None: attention_mask = (torch.triu(torch.ones(input_ids.size(1), input_ids.size(1))) == 0).transpose(0, 1).contiguous().to(input_ids.device) for layer in self.layers: x = layer(x, attn_mask=attention_mask) x = self.final_norm(x) # N x L x D return BaseModelOutput( hidden_states=x, ) #Some common HF functions. def get_input_embeddings(self): return self.embedding def set_input_embeddings(self, new_embeddings): self.embedding = new_embeddings def _init_weights(self, module): """Initialize the weights.""" if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) class RITAModelForCausalLM(PreTrainedModel): config_class = RITAConfig base_model_prefix = "transformer" is_parallelizable = False def __init__( self, config ): super().__init__(config) self.transformer = RITAModel(config) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) def forward( self, input_ids=None, past_key_values=None, # NOT USED attention_mask=None, token_type_ids=None, # NOT USED position_ids=None, # NOT USED head_mask=None, # NOT USED inputs_embeds=None, encoder_hidden_states=None, # NOT USED encoder_attention_mask=None, # NOT USED labels=None, use_cache=None, # NOT USED output_attentions=None, # NOT USED output_hidden_states=None, # NOT USED return_dict=None # NOT USED ) -> torch.FloatTensor: transformer_outputs = self.transformer( input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) logits = self.lm_head(transformer_outputs.hidden_states) loss = None if labels is not None: # Shift so that tokens < n predict n shift_logits = 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)) return CausalLMOutput( loss=loss, logits=logits, hidden_states=transformer_outputs.hidden_states, ) #Some common HF functions. def get_input_embeddings(self): return self.transformer.embedding def set_input_embeddings(self, new_embeddings): self.transformer.embedding = new_embeddings def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, lm_head): self.lm_head = lm_head def _init_weights(self, module): """Initialize the weights.""" if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) class RITAModelForSequenceClassification(PreTrainedModel): config_class = RITAConfig base_model_prefix = "transformer" is_parallelizable = False def __init__(self, config): super().__init__(config) self.num_labels = config.num_labels self.transformer = RITAModel(config) self.score = nn.Linear(config.d_model, self.num_labels, bias=False) def forward( self, input_ids=None, past_key_values=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): r""" labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). """ 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, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_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[0] logits = self.score(hidden_states) if input_ids is not None: batch_size, sequence_length = input_ids.shape[:2] else: batch_size, sequence_length = inputs_embeds.shape[:2] assert ( self.config.pad_token_id is not None or batch_size == 1 ), "Cannot handle batch sizes > 1 if no padding token is defined." if self.config.pad_token_id is None: sequence_lengths = -1 else: if input_ids is not None: sequence_lengths = torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1 else: sequence_lengths = -1 logger.warning( f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " f"unexpected if using padding tokens in conjunction with `inputs_embeds.`" ) pooled_logits = logits[torch.arange(batch_size, device=self.device), sequence_lengths] loss = None if labels is not None: if self.config.problem_type is None: if self.num_labels == 1: self.config.problem_type = "regression" elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): self.config.problem_type = "single_label_classification" else: self.config.problem_type = "multi_label_classification" if self.config.problem_type == "regression": loss_fct = MSELoss() if self.num_labels == 1: loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) else: loss = loss_fct(pooled_logits, labels) elif self.config.problem_type == "single_label_classification": loss_fct = CrossEntropyLoss() loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) elif self.config.problem_type == "multi_label_classification": loss_fct = BCEWithLogitsLoss() loss = loss_fct(pooled_logits, labels) if not return_dict: output = (pooled_logits,) + transformer_outputs[1:] return ((loss,) + output) if loss is not None else output return SequenceClassifierOutput( loss=loss, logits=pooled_logits, ) def _init_weights(self, module): """Initialize the weights.""" if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0)