Upload ModularStarEncoder

config.json

@@ -3,6 +3,10 @@
     "ModularStarEncoder"
   ],
   "attention_dropout": 0.1,
+  "auto_map": {
+    "AutoConfig": "config.ModularStarEncoderConfig",
+    "AutoModel": "modularStarEncoder.ModularStarEncoder"
+  },
   "bos_token_id": 0,
   "conditional_size": 4,
   "embedding_dropout": 0.1,

config.py

@@ -0,0 +1,81 @@
+from transformers import PretrainedConfig
+from typing import List
+
+
+#STARCODER2_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+class ModularStarEncoderConfig(PretrainedConfig):
+    model_type = "ModularStarEncoder"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        attention_dropout= 0.1,
+        residual_dropout=  0.1,
+        embedding_dropout=  0.1,
+        bos_token_id=  0,
+        eos_token_id= 0,
+        hidden_act= "gelu_pytorch_tanh",
+        _attn_implementation="flash_attention_2",
+        hidden_size= 1024,
+        conditional_size= 4,
+        initializer_range= 0.018042,
+        intermediate_size= 12288,
+        max_position_embeddings= 2048,
+        mlp_type= "default",
+        model_type= "starcoder2",
+        torch_dtype= "bfloat16",
+        layer_matryoshka_loss= True,
+        matryoshka_layers= [4,9,18,27,36],
+        norm_epsilon= 1e-05,
+        layer_norm_eps=1e-05,
+        norm_type= "layer_norm",
+        num_attention_heads= 16, 
+        num_hidden_layers= 36, 
+        num_key_value_heads= 4,
+        rope_theta= 999999.4420358813,
+        sliding_window= None,
+        transformers_version= "4.39.3",
+        use_bias= True,
+        use_cache= False,
+        vocab_size= 49156,
+        pad_token_id=0,
+        **kwargs,
+    ):
+        if _attn_implementation not in ["flash_attention_2", "sdpa"]:
+            raise ValueError(f"`_attn_implementation` must be 'flash_attention_2', 'sdpa', got {_attn_implementation}.")
+
+        self.attention_dropout=attention_dropout ,
+        self.residual_dropout=  residual_dropout,
+        self.embedding_dropout=  embedding_dropout,
+        self.bos_token_id=  bos_token_id,
+        self.eos_token_id= eos_token_id,
+        self.hidden_act= hidden_act,
+        self._attn_implementation=_attn_implementation,
+        self.hidden_size= hidden_size,
+        self.conditional_size= conditional_size,
+        self.initializer_range= initializer_range,
+        self.intermediate_size= intermediate_size,
+        self.max_position_embeddings= max_position_embeddings,
+        self.mlp_type= mlp_type,
+        self.model_type= model_type,
+        self.torch_dtype= torch_dtype,
+        self.layer_matryoshka_loss= layer_matryoshka_loss,
+        self.matryoshka_layers= matryoshka_layers,
+        self.norm_epsilon= norm_epsilon,
+        self.layer_norm_eps=layer_norm_eps,
+        self.norm_type= norm_type,
+        self.num_attention_heads= num_attention_heads, 
+        self.num_hidden_layers= num_hidden_layers, 
+        self.num_key_value_heads= num_key_value_heads,
+        self.rope_theta= rope_theta,
+        self.sliding_window= sliding_window,
+        self.transformers_version= transformers_version,
+        self.use_bias= use_bias,
+        self.use_cache= use_cache,
+        self.vocab_size= vocab_size,
+        self.pad_token_id=pad_token_id,
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs)

modularStarEncoder.py

@@ -0,0 +1,373 @@
+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,
+        )
+
+
+
+
+