diff --git a/src/core/operators.py b/src/core/operators.py new file mode 100644 index 00000000..af43f26c --- /dev/null +++ b/src/core/operators.py @@ -0,0 +1,48 @@ +""" +这里面存放是 RAG 相关的一些组件 +""" + +from src.utils import prompts + +class BaseOperator: + """ + 基类 + """ + template = None + + def __init__(self): + pass + + def call(self, **kwargs): + pass + + def __call__(self, **kwargs): + """ + 所有 RAG 相关组件的调用接口 + """ + return self.call(**kwargs) + + + +class HyDEOperator(BaseOperator): + """ + HyDE 重写查询 + """ + template = prompts.HYDE_PROMPT_TEMPLATE + + def __init__(self): + super().__init__() + + @classmethod + def call(cls, model_callable, query, context_str, **kwargs): + """ + 重写查询 + + Args: + model_callable: 模型调用函数 + query: 查询 + context_str: 上下文 + """ + prompt = cls.template.format(query=query, context_str=context_str) + response = model_callable(prompt) + return response diff --git a/src/core/retriever.py b/src/core/retriever.py index 92477b04..bc053afc 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -2,6 +2,8 @@ from src import config, knowledge_base, graph_base from src.models.rerank_model import get_reranker from src.utils.logging_config import logger from src.models import select_model +from src.core.operators import HyDEOperator + class Retriever: @@ -151,8 +153,8 @@ class Retriever: rewritten_query = model.predict(rewritten_query_prompt).content if rewrite_query_span == "hyde": - hy_doc = model.predict(rewritten_query).content - rewritten_query = f"{rewritten_query} {hy_doc}" + res = HyDEOperator.call(model_callable=model.predict, query=query, context_str=history_query) + rewritten_query = res.content return rewritten_query diff --git a/src/utils/prompts.py b/src/utils/prompts.py index 489fc671..09c41eaf 100644 --- a/src/utils/prompts.py +++ b/src/utils/prompts.py @@ -56,4 +56,16 @@ keywords_prompt_template = """ 返回的实体使用<->隔开。如:关键词1<->关键词<->关键词3 不要改变关键词的语言 <文本>{text} -""" \ No newline at end of file +""" + +HYDE_PROMPT_TEMPLATE = ( + "Please write a passage to answer the question\n" + "Try to include as many key details as possible.\n" + "\n" + "\n" + "{context_str}\n" + "\n" + "{query}\n" + "\n" + 'Passage:\n' +)