Adrien
commited on
Commit
·
3823e3e
1
Parent(s):
0090181
fix: reranking was done two times
Browse files- 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 |
-
|
143 |
|
144 |
-
answer = self.generate_answer(query,
|
145 |
|
146 |
-
if
|
147 |
-
annotated_answer = self.add_context(answer,
|
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 |
)
|