from transformers import AutoConfig, Starcoder2Model, Starcoder2Config from modularStarEncoder.config import ModularStarEncoderConfig import math import os import warnings from dataclasses import dataclass from typing import List, Optional, Tuple, Union import sys import torch import torch.utils.checkpoint from torch import nn from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers.activations import ACT2FN from transformers.modeling_outputs import ( BaseModelOutputWithPastAndCrossAttentions, BaseModelOutputWithPoolingAndCrossAttentions, CausalLMOutputWithCrossAttentions, MaskedLMOutput, MultipleChoiceModelOutput, NextSentencePredictorOutput, QuestionAnsweringModelOutput, SequenceClassifierOutput, TokenClassifierOutput, ) from transformers.modeling_utils import PreTrainedModel from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer from transformers.utils import ( ModelOutput, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings, ) logger = logging.get_logger(__name__) class StarEncoder2PreTrainedModel(PreTrainedModel): """ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. """ config_class = ModularStarEncoderConfig base_model_prefix = "ModularStarEncoder" model_type = "ModularStarEncoder" supports_gradient_checkpointing = True _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True # def __init__(self): # self._supports_flash_attn_2 = True # super().__init__() def _init_weights(self, module): """Initialize the weights""" if isinstance(module, nn.Linear): # Slightly different from the TF version which uses truncated_normal for initialization # cf https://github.com/pytorch/pytorch/pull/5617 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 StarEncoder2Pooler(nn.Module): def __init__(self, config): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.activation = nn.Tanh() def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # We "pool" the model by simply taking the hidden state corresponding # to the last token. last_token_tensor = hidden_states[:, -1] pooled_output = self.dense(last_token_tensor) pooled_output = self.activation(pooled_output) return pooled_output @dataclass class ModularStarEncoderOutput(ModelOutput): """ Output type of [`BertForPreTraining`]. Args: loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss. prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`): Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax). hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. """ loss: Optional[torch.FloatTensor] = None prediction_logits: torch.FloatTensor = None seq_relationship_logits: torch.FloatTensor = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None class StarEncoder2PredictionHeadTransform(nn.Module): def __init__(self, config): super().__init__() self.is_matryoshka = config.layer_matryoshka_loss if self.is_matryoshka: self.dense = nn.Linear(config.hidden_size + config.conditional_size, config.hidden_size + config.conditional_size) self.LayerNorm = nn.LayerNorm(config.hidden_size + config.conditional_size, eps=config.layer_norm_eps) else: self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if isinstance(config.hidden_act, str): self.transform_act_fn = ACT2FN[config.hidden_act] else: self.transform_act_fn = config.hidden_act def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.dense(hidden_states) hidden_states = self.transform_act_fn(hidden_states) hidden_states = self.LayerNorm(hidden_states) return hidden_states class StarEncoder2LMPredictionHead(nn.Module): def __init__(self, config): super().__init__() for element in dir(config): value = getattr(config, element) # Get the attribute value if isinstance(value, tuple) or isinstance(value, list): setattr(config, element, value[0] ) self.transform = StarEncoder2PredictionHeadTransform(config) # The output weights are the same as the input embeddings, but there is # an output-only bias for each token. self.is_matryoshka = config.layer_matryoshka_loss if self.is_matryoshka: self.decoder = nn.Linear(config.hidden_size + config.conditional_size, config.vocab_size, bias=False) else: self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.bias = nn.Parameter(torch.zeros(config.vocab_size)) # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings` self.decoder.bias = self.bias def forward(self, hidden_states): hidden_states = self.transform(hidden_states) hidden_states = self.decoder(hidden_states) return hidden_states class StarEncoder2PreTrainingHeads(nn.Module): def __init__(self, config): super().__init__() self.predictions = StarEncoder2LMPredictionHead(config) self.is_matryoshka = config.layer_matryoshka_loss if self.is_matryoshka: self.seq_relationship = nn.Linear(config.hidden_size + config.conditional_size, 2) self.conditional_embeddings = nn.Embedding(len(config.matryoshka_layers),config.conditional_size) else: self.seq_relationship = nn.Linear(config.hidden_size, 2) def forward(self, sequence_output, pooled_output,idx_layer: Optional[torch.Tensor] = None): if self.is_matryoshka: prediction_scores = self.predictions(torch.cat([sequence_output , self.conditional_embeddings(torch.tensor(idx_layer,device=sequence_output.get_device()).int()).expand(sequence_output.size()[0],sequence_output.size()[1],-1)],dim=-1)) seq_relationship_score = self.seq_relationship(torch.cat([pooled_output , self.conditional_embeddings(torch.tensor(idx_layer,device=pooled_output.get_device()).int()).expand(pooled_output.size()[0],-1)],dim=-1)) else: prediction_scores = self.predictions(sequence_output) seq_relationship_score = self.seq_relationship(pooled_output) return prediction_scores, seq_relationship_score class ModularStarEncoder(StarEncoder2PreTrainedModel): _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"] config_class = ModularStarEncoderConfig def __init__(self, config): super().__init__(config) self.model_type = "ModularStarEncoder" self.cls = StarEncoder2PreTrainingHeads(config) self.layer_matryoshka_loss = config.layer_matryoshka_loss self.matryoshka_layers = config.matryoshka_layers if self.layer_matryoshka_loss: config.sliding_window = None logger.warning_once( "The matryoshka loss is implemented without sliding_window, if you want to use the sliding window set sliding_window to True" ) if self.matryoshka_layers[-1] != config.num_hidden_layers: logger.warning_once( f"To get optimal results, the last layer on matryoshka layers, which now is {self.matryoshka_layers[-1]} " "must be set as the overall number of hidden layers." f"The overall number of hidden layers is now set to {config.num_hidden_layers}" ) sys.exit() self.starEncoder2 = Starcoder2Model(config) self.pooler = StarEncoder2Pooler(config) #setting off causal masking for layer in self.starEncoder2.layers: layer.self_attn.is_causal=False # Initialize weights and apply final processing self.post_init() # def get_output_embeddings(self): # return self.cls.predictions.decoder # def set_output_embeddings(self, new_embeddings): # self.cls.predictions.decoder = new_embeddings def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, #token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, next_sentence_label: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple[torch.Tensor], ModularStarEncoderOutput]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see `input_ids` docstring) Indices should be in `[0, 1]`: - 0 indicates sequence B is a continuation of sequence A, - 1 indicates sequence B is a random sequence. kwargs (`Dict[str, any]`, optional, defaults to *{}*): Used to hide legacy arguments that have been deprecated. Returns: Example: ```python >>> from transformers import AutoTokenizer, BertForPreTraining >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased") >>> model = BertForPreTraining.from_pretrained("google-bert/bert-base-uncased") >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") >>> outputs = model(**inputs) >>> prediction_logits = outputs.prediction_logits >>> seq_relationship_logits = outputs.seq_relationship_logits ``` """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.starEncoder2( input_ids, attention_mask=attention_mask, # token_type_ids=token_type_ids, position_ids=position_ids, # head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict, ) #TODO FIX FOR EFFICIENCY, COMPUTE FORWARD PASS JUST ON MATRYOSKA LAYERS #if layer matryoshka on, compute the scores for all the heads if self.layer_matryoshka_loss: prediction_scores = [] seq_relationship_score = [] #for layer in outputs.hidden_states: for counter,idx_layer in enumerate(self.matryoshka_layers): #pooling head to project last hidden states as CLS token is in the last position pooled_output = self.pooler(outputs.hidden_states[idx_layer]) #all the hidden states related to the last layer sequence_output = outputs.hidden_states[idx_layer] temp_prediction_scores, temp_seq_relationship_score = self.cls(sequence_output, pooled_output,counter) prediction_scores.append(temp_prediction_scores) seq_relationship_score.append(temp_seq_relationship_score) else: #pooling head to project last hidden states as CLS token is in the last position pooled_output = self.pooler(outputs.last_hidden_state) #all the hidden states related to the last layer sequence_output = outputs.last_hidden_state prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) total_loss = None if labels is not None and next_sentence_label is not None and not self.layer_matryoshka_loss: loss_fct = CrossEntropyLoss() masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1)) total_loss = masked_lm_loss + next_sentence_loss elif labels is not None and next_sentence_label is not None and self.layer_matryoshka_loss: loss_fct = CrossEntropyLoss() num_layers = len(prediction_scores) #for layer in self.matryoshka_layers: seq_relationship_score for index in range(num_layers): masked_lm_loss = loss_fct(prediction_scores[index].view(-1, self.config.vocab_size), labels.view(-1)) next_sentence_loss = loss_fct(seq_relationship_score[index].view(-1, 2), next_sentence_label.view(-1)) if total_loss: total_loss += (masked_lm_loss + next_sentence_loss) * ((index+1)/num_layers) else: total_loss = (masked_lm_loss + next_sentence_loss) * ((index+1)/num_layers) if not return_dict: output = (prediction_scores, seq_relationship_score) + outputs[2:] return ((total_loss,) + output) if total_loss is not None else output return ModularStarEncoderOutput( loss=total_loss, prediction_logits=prediction_scores, seq_relationship_logits=seq_relationship_score, hidden_states=outputs.hidden_states, attentions=outputs.attentions, )