Re-using same GenericDataLoader for different splits raises KeyError

[j-rossi-nl]

Here is the script:

gdl = GenericDataLoader(data_folder='.local/datasets/scifact')

corpus, queries, qrels = gdl.load(split="test")
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5183/5183 [00:00<00:00, 195184.62it/s]

corpus, queries, qrels = gdl.load(split="train")
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Input In [24], in <cell line: 1>()
----> 1 corpus, queries, qrels = gdl.load(split="train")

File ~/.cache/pypoetry/virtualenvs/sparse-attention-05J_nGF--py3.9/lib/python3.9/site-packages/beir/datasets/data_loader.py:78, in GenericDataLoader.load(self, split)
     76 if os.path.exists(self.qrels_file):
     77     self._load_qrels()
---> 78     self.queries = {qid: self.queries[qid] for qid in self.qrels}
     79     logger.info("Loaded %d %s Queries.", len(self.queries), split.upper())
     80     logger.info("Query Example: %s", list(self.queries.values())[0])

File ~/.cache/pypoetry/virtualenvs/sparse-attention-05J_nGF--py3.9/lib/python3.9/site-packages/beir/datasets/data_loader.py:78, in <dictcomp>(.0)
     76 if os.path.exists(self.qrels_file):
     77     self._load_qrels()
---> 78     self.queries = {qid: self.queries[qid] for qid in self.qrels}
     79     logger.info("Loaded %d %s Queries.", len(self.queries), split.upper())
     80     logger.info("Query Example: %s", list(self.queries.values())[0])

KeyError: '0'

The workaround at the moment is:

gdl.queries = {}

Obviously the code does not deal with reloading. Since the corpus is using so much memory, it would be a pity to create a new DataLoader for each split.

I would be happy to make a PR, but there will be a delay. I'll check again soon.