tarzanagh commited on
Commit
650c615
·
verified ·
1 Parent(s): d1764e3

Remove DPR fallback

Browse files
Files changed (1) hide show
  1. src/passage_entity/retriever.py +12 -6
src/passage_entity/retriever.py CHANGED
@@ -387,12 +387,18 @@ class PassageEntityRetriever:
387
  self.rerank_time += time.time() - rerank_t0
388
 
389
  if len(top_facts) == 0:
390
- logger.info("No facts after reranking -> fallback to DPR")
391
- sorted_ids, sorted_scores = self._dense_passage_retrieval(q)
392
- else:
393
- sorted_ids, sorted_scores = self._graph_search(
394
- q, fact_scores, top_facts, top_indices
395
- )
 
 
 
 
 
 
396
 
397
  top_docs = [
398
  self.chunk_store.get_row(self.passage_node_keys[idx])["content"]
 
387
  self.rerank_time += time.time() - rerank_t0
388
 
389
  if len(top_facts) == 0:
390
+ # Reranker rejected all facts use top facts by embedding score
391
+ logger.info("No facts after reranking -> using top facts by embedding score")
392
+ link_top_k = self.config.linking_top_k
393
+ if len(fact_scores) > 0:
394
+ top_indices = np.argsort(fact_scores)[-link_top_k:][::-1].tolist()
395
+ real_ids = [self.fact_node_keys[i] for i in top_indices]
396
+ rows = self.fact_store.get_rows(real_ids)
397
+ top_facts = [eval(rows[rid]["content"]) for rid in real_ids]
398
+
399
+ sorted_ids, sorted_scores = self._graph_search(
400
+ q, fact_scores, top_facts, top_indices
401
+ )
402
 
403
  top_docs = [
404
  self.chunk_store.get_row(self.passage_node_keys[idx])["content"]