Upload 2 files

modeling.py

@@ -654,7 +654,7 @@ class ExtendedLlamaAttention(nn.Module):
         if not output_attentions:
             attn_weights = None
 
-        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
 
@@ -1568,6 +1568,7 @@ class ExtendedLlamaForCausalLM(LlamaPreTrainedModel):
                 "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