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
@@ -307,10 +307,10 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
                 output_attentions=output_attentions,
                 output_hidden_states=True,
                 return_dict=True,
-            ).hidden_states
+            )
 
 
-        DEVICE = source_embedding[-1].get_device()
+        DEVICE = source_embedding.hidden_states[-1].get_device()
         if DEVICE<0:
             DEVICE = "cpu"
 
@@ -334,7 +334,7 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
 
         pooled_and_normalized = []
         for idx,matr_layer in enumerate(self.matryoshka_layers):
-            source_embedding_proj = projection_fn[idx](source_embedding[matr_layer])
+            source_embedding_proj = projection_fn[idx](source_embedding.hidden_states[matr_layer])
 
             normalized_source_embedding, embedding_norms = pool_and_normalize(
                         source_embedding_proj,