DenseRetrievalExactSearch: Out of memory

[memray]

I found BEIR consumes huge amount of memory when evaluating with large datasets such as HotpotQA and MSMARCO.

L73-L78 could be optimized since self.results will keep accumulating new scores, whereas only top K items will be used for computing metrics.

Consider changing to the code below:

for query_itr in range(len(query_embeddings)):
    query_id = query_ids[query_itr]                  
    for sub_corpus_id, score in zip(cos_scores_top_k_idx[query_itr], cos_scores_top_k_values[query_itr]):
        corpus_id = corpus_ids[corpus_start_idx+sub_corpus_id]
        if corpus_id != query_id:
            self.results[query_id][corpus_id] = score
    _results = {}
    for cid, score in sorted(self.results[query_id].items(), key=lambda k: k[1], reverse=True)[:top_k]:
         _results[cid] = score
    self.results[query_id] = _results

[rahmanidashti]

Hi @memray, thank you for sharing this, it is really useful to address the memory issue when working with large datasets.