Upload 2 files

modeling.py

@@ -356,7 +356,7 @@ class ExtendedMptAttention(nn.Module):
         )
         attn_output = self.out_proj(context_states)
 
-        if not output_retrieved_memory_idx:
+        if not output_retrieved_memory_idx or (long_range_past_key_value is None and faiss_indexes is None):
             reshaped_idx = None
 
         return attn_output, attn_weights, past_key_value, reshaped_idx
@@ -977,6 +977,7 @@ class ExtendedMptForCausalLM(MptPreTrainedModel):
                 "attention_mask": attention_mask,
                 "use_external_mind": kwargs.get("use_external_mind"), # EM: Add config here
                 "topk": kwargs.get("topk"),
+                "output_retrieved_memory_idx": kwargs.get("output_retrieved_memory_idx"),
             }
         )
         return model_inputs