Upload ModularStarEncoder

config.json

@@ -11,7 +11,6 @@
   "conditional_size": 4,
   "embedding_dropout": 0.1,
   "eos_token_id": 0,
-  "pad_token_id": 0,
   "hidden_act": "gelu_pytorch_tanh",
   "hidden_size": 1024,
   "initializer_range": 0.018042,

modularStarEncoder.py

@@ -1,6 +1,6 @@
 from transformers import  Starcoder2Model
 import sys
-from .config import ModularStarEncoderConfig
+from config import ModularStarEncoderConfig
 import os
 from dataclasses import dataclass
 from typing import Optional, Tuple, Union
@@ -92,10 +92,11 @@ class ModularStarEncoderOutput(ModelOutput):
             heads.
     """
 
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     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
 
 
@@ -316,12 +317,16 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
             output = (prediction_scores, seq_relationship_score) + outputs[2:]
             return ((total_loss,) + output) if total_loss is not None else output
 
+
+        last_hidden_state= outputs.hidden_states[-1]
+
         return ModularStarEncoderOutput(
-            loss=total_loss,
-            prediction_logits=prediction_scores,
-            seq_relationship_logits=seq_relationship_score,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
+            last_hidden_state = last_hidden_state,
+            hidden_states = outputs.hidden_states,
+            loss = total_loss,
+            prediction_logits = prediction_scores,
+            seq_relationship_logits = seq_relationship_score,
+            attentions = outputs.attentions,
         )