Upload ModularStarEncoder

modularStarEncoder.py

@@ -1,39 +1,21 @@
 from transformers import AutoConfig, Starcoder2Model, Starcoder2Config
 import sys
+from config import ModularStarEncoderConfig
 import os
-from .config import ModularStarEncoderConfig
-import math
-import os
-import warnings
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 import sys
 import torch
 import torch.utils.checkpoint
 from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
+from torch.nn import  CrossEntropyLoss
 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__)
@@ -243,11 +225,7 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
         # 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
 
 
 
@@ -279,40 +257,20 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
             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 = []