Adrien commited on
Commit
3823e3e
·
1 Parent(s): 0090181

fix: reranking was done two times

Browse files
Files changed (1) hide show
  1. rag_demo/rag/retriever.py +5 -12
rag_demo/rag/retriever.py CHANGED
@@ -135,24 +135,17 @@ class RAGPipeline:
135
  logger.info(f"Query type: {query_type}")
136
  if query_type == "Sources_needed":
137
  docs = self.search(query, k=10)
138
- # disabling reranking for now because HFApi is too slow
139
- # reranked_docs = self.rerank(query, docs, keep_top_k=10)
140
- reranked_docs = docs
141
  else:
142
- reranked_docs = []
143
 
144
- answer = self.generate_answer(query, reranked_docs)
145
 
146
- if reranked_docs:
147
- annotated_answer = self.add_context(answer, reranked_docs)
148
  else:
149
  annotated_answer = answer
150
 
151
  return (
152
  annotated_answer,
153
- list(
154
- set(
155
- [doc.metadata["filename"].split(".pdf")[0] for doc in reranked_docs]
156
- )
157
- ),
158
  )
 
135
  logger.info(f"Query type: {query_type}")
136
  if query_type == "Sources_needed":
137
  docs = self.search(query, k=10)
 
 
 
138
  else:
139
+ docs = []
140
 
141
+ answer = self.generate_answer(query, docs)
142
 
143
+ if docs:
144
+ annotated_answer = self.add_context(answer, docs)
145
  else:
146
  annotated_answer = answer
147
 
148
  return (
149
  annotated_answer,
150
+ list(set([doc.metadata["filename"].split(".pdf")[0] for doc in docs])),
 
 
 
 
151
  )