kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix reranker bug

Browse Source
This commit is contained in:
Wenjie Zhang 2024-09-14 02:45:13 +08:00
parent 5925e359e9
commit 5c03d3d7cc
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/core/retriever.py
View File

@ -17,6 +17,7 @@ class Retriever:
def retrieval(self, query, history, meta): def retrieval(self, query, history, meta):
refs = {"query": query, "history": history, "meta": meta} refs = {"query": query, "history": history, "meta": meta}
refs["model_name"] = self.config.model_name
refs["entities"] = self.reco_entities(query, history, refs) refs["entities"] = self.reco_entities(query, history, refs)
refs["knowledge_base"] = self.query_knowledgebase(query, history, refs) refs["knowledge_base"] = self.query_knowledgebase(query, history, refs)
refs["graph_base"] = self.query_graph(query, history, refs) refs["graph_base"] = self.query_graph(query, history, refs)
@ -96,7 +97,10 @@ class Retriever:
r["file"] = kb.id2file(r["entity"]["file_id"]) r["file"] = kb.id2file(r["entity"]["file_id"])
# use distance threshold to filter results # use distance threshold to filter results
kb_res = [r for r in all_kb_res if r["distance"] > distance_threshold] if refs["meta"].get("mode") == "search":
kb_res = all_kb_res
else:
kb_res = [r for r in all_kb_res if r["distance"] > distance_threshold]
if self.config.enable_reranker: if self.config.enable_reranker:
for r in kb_res: for r in kb_res: