RAG系统(一)系统介绍与向量检索
RAG系统(二)文档的加载与分段
RAG系统(三)向量数据库
完整代码需要依赖前边三章中的代码,本节主要展示整合后的RAG系统,及运行示例。
from MyVectorDBConnector import MyVectorDBConnector
from chatgptDemo import get_completion, prompt_template, build_prompt
from embedding import get_embeddings
from pdf_handle import PdfHandlerclass RAG_Bot:def __init__(self, vector_db, llm_api, n_results=2):self.vector_db = vector_dbself.llm_api = llm_apiself.n_results = n_resultsdef chat(self, user_query):# 1. 检索search_results = self.vector_db.search(user_query, self.n_results)# 2. 构建 Promptprompt = build_prompt(prompt_template, context=search_results['documents'][0], query=user_query)# 3. 调用 LLMresponse = self.llm_api(prompt)return response
运行示例:
if __name__ == '__main__':# 只取两页pdf = PdfHandler("llama2.pdf")paragraphs = pdf.extract_text_from_pdf(page_numbers=[2, 3],min_line_length=10)# 创建一个向量数据库对象vector_db = MyVectorDBConnector("demo", get_embeddings)# 向向量数据库中添加文档vector_db.add_documents(paragraphs)# 创建一个RAG机器人bot = RAG_Bot(vector_db,llm_api=get_completion)user_query = "llama 2有多少参数?"response = bot.chat(user_query)print(response)
其中”get_completion“是chatgpt的模型封装,实际使用可根据自身选用的模型API或SDK封装一个调用方法传递进来。
prompt
prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。已知信息:
{context}用户问:
{query}如果已知信息不包含用户问题的答案,或者已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。
请不要输出已知信息中不包含的信息或答案。
请用中文回答用户问题。
"""
prompt模板赋值
def build_prompt(prompt_template, **kwargs):'''将 Prompt 模板赋值'''inputs = {}for k, v in kwargs.items():if isinstance(v, list) and all(isinstance(elem, str) for elem in v):val = '\n\n'.join(v)else:val = vinputs[k] = valreturn prompt_template.format(**inputs)