Upload ModularStarEncoder

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, List
@@ -347,12 +347,18 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
             
             pooled_and_normalized.append(normalized_source_embedding)
 
-
-        return ModularStarEncoderOutput(
-            projected_pooled_normalized = pooled_and_normalized,
-            raw_hidden_states=source_embedding.hidden_states,
-            attentions=source_embedding.attentions,
-        )
+        if not self.till_layer:
+            return ModularStarEncoderOutput(
+                projected_pooled_normalized = pooled_and_normalized,
+                raw_hidden_states=source_embedding.hidden_states,
+                attentions=source_embedding.attentions,
+            )
+        else:
+            return ModularStarEncoderOutput(
+                projected_pooled_normalized = pooled_and_normalized[0],
+                raw_hidden_states=source_embedding.hidden_states,
+                attentions=source_embedding.attentions,
+            )