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)