Question about memory storage and retrieval

[namespace-Pt]

Greetings! I have some questions regarding memory storage and retrieval in the provided code. Based on your code, it appears that you have stored the per-head embeddings of all chunks in the batch and subsequently retrieved the chunk embeddings for queries in the batch.

However, the paper states that each sample in the same batch originates from a different document. Thus, I find it puzzling that all chunks in the batch are mixed together and then retrieved. I have been unable to locate any specific code for reorganizing the retrieved chunks based on their original documents. As far as I understand, this reorganization may not be possible with faiss. Could you kindly provide some clarification on this matter? Thank you.

[nampdn]

same question here, hopefully author can help we clear this concern