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

优化改写逻辑，默认启用改写查询

Browse Source
This commit is contained in:
Wenjie Zhang 2025-05-09 23:47:45 +08:00
parent 144632b703
commit ea40fa4d1b
3 changed files with 15 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/config/__init__.py
View File

@ -62,7 +62,7 @@ class Config(SimpleConfig):
self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys())) self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys()))
self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys())) self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys()))
self.add_item("model_local_paths", default={}, des="本地模型路径") self.add_item("model_local_paths", default={}, des="本地模型路径")
self.add_item("use_rewrite_query", default="off", des="重写查询", choices=["off", "on", "hyde"]) self.add_item("use_rewrite_query", default="on", des="重写查询", choices=["off", "on", "hyde"])
self.add_item("device", default="cuda", des="运行本地模型的设备", choices=["cpu", "cuda"]) self.add_item("device", default="cuda", des="运行本地模型的设备", choices=["cpu", "cuda"])
### <<< 默认配置结束 ### <<< 默认配置结束

11
src/core/retriever.py
View File

@ -145,13 +145,12 @@ class Retriever:
rewrite_query_span = config.use_rewrite_query rewrite_query_span = config.use_rewrite_query
if rewrite_query_span == "off": if rewrite_query_span == "off":
rewritten_query = query return query
else:
from src.utils.prompts import rewritten_query_prompt_template
history_query = [entry["content"] for entry in history if entry["role"] == "user"] if history else "" from src.utils.prompts import rewritten_query_prompt_template2 as rw_template
rewritten_query_prompt = rewritten_query_prompt_template.format(history=history_query, query=query) history_query = [entry["content"] for entry in history if entry["role"] == "user"] if history else ""
rewritten_query = model.predict(rewritten_query_prompt).content rewritten_query_prompt = rw_template.format(history=history_query, query=query)
rewritten_query = model.predict(rewritten_query_prompt).content
if rewrite_query_span == "hyde": if rewrite_query_span == "hyde":
res = HyDEOperator.call(model_callable=model.predict, query=query, context_str=history_query) res = HyDEOperator.call(model_callable=model.predict, query=query, context_str=history_query)

11
src/utils/prompts.py
View File

@ -34,8 +34,15 @@ rewritten_query_prompt_template = """
rewritten_query_prompt_template2 = """ rewritten_query_prompt_template2 = """
你是一个用来辅助查询的助手，请根据历史对话以及最新的问题，改写出多个与查询相关的查询问题，用于从知识库中匹配到参考资料； 你是一个用来辅助查询的助手，请根据历史对话以及最新的问题，改写出多个与查询相关的查询问题，用于从知识库中匹配到参考资料；
<历史信息>{history}</历史信息>
<问题>{query}</问题> <示例>
历史提问：无锡有哪些好吃的早点？
新的提问：火锅呢？
期望的改写：无锡有哪些好吃的火锅？
</示例>
<历史提问>{history}</历史提问>
<新的问题>{query}</新的问题>
""" """