Upload 2 files

modeling.py

@@ -47,7 +47,7 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 
-from .configuration import ExtendedLlamaConfig
+from emts_clean.src.llama.configuration import ExtendedLlamaConfig
 
 logger = logging.get_logger(__name__)
 
@@ -1144,7 +1144,7 @@ class ExtendedLlamaForCausalLM(LlamaPreTrainedModel):
 
     _tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(self, config, external_memories=None):
+    def __init__(self, config, external_memories:list=None):
         super().__init__(config)
         self.model = ExtendedLlamaModel(config)
         self.vocab_size = config.vocab_size
@@ -1242,9 +1242,9 @@ class ExtendedLlamaForCausalLM(LlamaPreTrainedModel):
         if (
             self.memory_ids is not None and self.memories is None
         ): 
+            self.memory_ids = torch.tensor([self.memory_ids], device=self.device) if type(self.memory_ids)==list else self.memory_ids
             self.memories = self.generate_cache(
-                torch.tensor(self.memory_ids, device=self.device),
-                cache_type=self.memory_type,
+                self.memory_ids, cache_type=self.memory_type,
             )
             # EM: Remove special tokens from memory cache
             if self.remove_special_ids: