50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
|
from core import PreRetrieval
|
|||
|
|
|
|||
|
|
class Retriever:
|
|||
|
|
|
|||
|
|
def __init__(self, config):
|
|||
|
|
self.config = config
|
|||
|
|
self.pre_retrieval = PreRetrieval(config)
|
|||
|
|
|
|||
|
|
def retrieval(self, query):
|
|||
|
|
|
|||
|
|
refs = {}
|
|||
|
|
|
|||
|
|
# TODO: 查询分类、查询重写、查询分解、查询伪文档生成（HyDE）
|
|||
|
|
|
|||
|
|
if self.config.enable_knowledge_base:
|
|||
|
|
refs["knowledge_base"] = self.pre_retrieval.search(query)
|
|||
|
|
|
|||
|
|
return refs
|
|||
|
|
|
|||
|
|
def construct_query(self, query, refs):
|
|||
|
|
# TODO：Reranking
|
|||
|
|
|
|||
|
|
if len(refs) == 0:
|
|||
|
|
return query
|
|||
|
|
|
|||
|
|
external = ""
|
|||
|
|
|
|||
|
|
kb_res = refs.get("knowledge_base")
|
|||
|
|
if kb_res:
|
|||
|
|
kb_text = "\n".join([f"{r['id']}: {r['entity']['text']}" for r in kb_res])
|
|||
|
|
external += f"知识库信息: \n\n{kb_text}"
|
|||
|
|
|
|||
|
|
if len(external) > 0:
|
|||
|
|
query = f"以下是参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答：{query}"
|
|||
|
|
|
|||
|
|
return query
|
|||
|
|
|
|||
|
|
def query_classification(self, query):
|
|||
|
|
"""判断是否需要查询
|
|||
|
|
- 对于完全基于用户给定信息的任务，称之为“足够”“sufficient”，不需要检索；
|
|||
|
|
- 否则，称之为“不足”“insufficient”，可能需要检索，
|
|||
|
|
"""
|
|||
|
|
raise NotImplementedError
|
|||
|
|
|
|||
|
|
def rewrite_query(self, query):
|
|||
|
|
"""重写查询"""
|
|||
|
|
raise NotImplementedError
|
|||
|
|
|
|||
|
|
def __call__(self, query):
|
|||
|
|
return self.retrieval(query)
|