Upload 2 files

configuration.py

@@ -101,6 +101,7 @@ class ExtendedMptAttentionConfig(PretrainedConfig):
         sim_threshold=0.25,
         tokenizer_all_special_ids=[0, 50278],
         remove_special_ids=False,
+        use_external_mind_by_layer: list[bool] = [True for _ in range(32)],
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -121,6 +122,7 @@ class ExtendedMptAttentionConfig(PretrainedConfig):
         self.sim_threshold = sim_threshold
         self.tokenizer_all_special_ids = tokenizer_all_special_ids
         self.remove_special_ids = remove_special_ids
+        self.use_external_mind_by_layer = use_external_mind_by_layer
 
         if attn_type not in ["multihead_attention", "multiquery_attention"]:
             raise ValueError(
@@ -245,7 +247,6 @@ class ExtendedMptConfig(PretrainedConfig):
         n_layers: int = 32,
         expansion_ratio: int = 4,
         max_seq_len_inference: int = 2048,
-        max_seq_len_train: int = 2048,
         vocab_size: int = 50432,
         resid_pdrop: float = 0.0,
         layer_norm_epsilon: float = 1e-5,
@@ -261,11 +262,12 @@ class ExtendedMptConfig(PretrainedConfig):
         use_cache: bool = False,
         initializer_range=0.02,
         use_external_mind: bool = True,
-        use_external_mind_by_layer: list[bool] = [True for _ in range(32)],
         **kwargs,
     ):
         if attn_config is None:
-            self.attn_config = ExtendedMptAttentionConfig()
+            self.attn_config = ExtendedMptAttentionConfig(
+                use_external_mind_by_layer=[True for _ in range(n_layers)]
+            )
         elif not isinstance(attn_config, ExtendedMptAttentionConfig):
             self.attn_config = ExtendedMptAttentionConfig(**attn_config)
         else:
@@ -275,7 +277,6 @@ class ExtendedMptConfig(PretrainedConfig):
         self.n_layers = n_layers
         self.expansion_ratio = expansion_ratio
         self.max_seq_len = max_seq_len_inference
-        self.max_seq_len_train = max_seq_len_train
         self.vocab_size = vocab_size
         self.resid_pdrop = resid_pdrop
         self.emb_pdrop = emb_pdrop

modeling.py

@@ -920,7 +920,7 @@ class ExtendedMptForCausalLM(MptPreTrainedModel):
 
     _tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(self, config: ExtendedMptConfig, external_memories=None):
+    def __init__(self, config: ExtendedMptConfig, external_memories:list=None):
         super().__init__(config)
         self.transformer: ExtendedMptModel = ExtendedMptModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
@@ -1016,8 +1016,9 @@ class ExtendedMptForCausalLM(MptPreTrainedModel):
         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(
-                self.memory_ids, 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: