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-13 15:25:01 +08:00
parent 8c3d9e0f0a
commit 0f249759a7
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/core/retriever.py
View File

@ -84,6 +84,7 @@ class Retriever:
db_name = refs["meta"]["db_name"]
kb = self.dbm.metaname2db[db_name]
logger.debug(f"{refs['meta']=}")
max_query_count = refs["meta"].get("maxQueryCount", 10)
rerank_threshold = refs["meta"].get("rerankThreshold", 0.1)
@ -99,7 +100,7 @@ class Retriever:
if self.config.enable_reranker:
for r in kb_res:
r["rerank_score"] = self.reranker.compute_score([rw_query, r["entity"]["text"]], normalize=True)
r["rerank_score"] = self.reranker.compute_score([rw_query, r["entity"]["text"]], normalize=True)[0]
kb_res.sort(key=lambda x: x["rerank_score"], reverse=True)
kb_res = [_res for _res in kb_res if _res["rerank_score"] > rerank_threshold]