AmenRa / retriv

A Python Search Engine for Humans 🥸
MIT License
174 stars 20 forks source link

ANN_Searcher not dealing with -1 returned by faiss_index.search() #31

Open jacobvsdanniel opened 8 months ago

jacobvsdanniel commented 8 months ago
Traceback (most recent call last):
  File "/lib/python3.8/site-packages/retriv/dense_retriever/dense_retriever.py", line 251, in search
    doc_ids = self.map_internal_ids_to_original_ids(doc_ids)
  File "/lib/python3.8/site-packages/retriv/base_retriever.py", line 80, in map_internal_ids_to_original_ids
    return [self.id_mapping[doc_id] for doc_id in doc_ids]
  File "/lib/python3.8/site-packages/retriv/base_retriever.py", line 80, in <listcomp>
    return [self.id_mapping[doc_id] for doc_id in doc_ids]
KeyError: -1

Update "/lib/python3.8/site-packages/retriv/base_retriever.py", line 80 to

    return [self.id_mapping[doc_id] for doc_id in doc_ids if doc_id != -1]

would fix the problem.