Rag实战案例代码

今天学了Rag技术，感觉挺有意思，所以就自己亲自玩了玩，下面是通过爬取知乎中关于选购机型机型指南作为大模型的外挂知识库，然后使用Rag技术实现对这篇文章中的内容进行回答。
如果您不太清楚百度api_key和secret_key的申请，可以参考这篇文章

import requests
from langchain_wenxin.llms import Wenxin
from langchain.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser

'''--------------------------------------------step1: 准备大模型未见过的知识--------------------------------------------'''
url = "https://www.zhihu.com/tardis/bd/art/511508925" # 爬取链接，里面包含了大模型没学过的知识
res = requests.get(url)
with open("93853.html", "w") as f:
    f.write(res.text)

'''--------------------------------------------ste2: 将知识文本划分成若干个chunk------------------------------------------'''
loader = TextLoader('./93853.html')
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = text_splitter.split_documents(documents) # 将文本划分成若干个chunk

'''--------------------------------------------step3: 使用m3e模型对这些chunk进行编码----------------------------------------'''
# embedding model: m3e-base
model_name = "moka-ai/m3e-base"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
embedding = HuggingFaceBgeEmbeddings(
                model_name=model_name,
                model_kwargs=model_kwargs,
                encode_kwargs=encode_kwargs,
                query_instruction="为文本生成向量表示用于文本检索"
            )

'''--------------------------------------------step4: 把编码后的特征向量存入chroma数据库中-----------------------------------'''
# load data to Chroma db
db = Chroma.from_documents(chunks, embedding)
retriever = db.as_retriever()

'''--------------------------------------------step5: 准备提升模板-------------------------------------------------------'''
template = """你是问答任务助手。使用以下检索到的上下文片段来回答问题。如果你不知道答案，就说你不知道。最多使用三个句子，保持答案简洁。
Question: {question} 
Context: {context} 
Answer:
"""
prompt = ChatPromptTemplate.from_template(template)

'''--------------------------------------------step6: 将检索出的知识与query组合通过大模型生成回答-----------------------------'''
llm = Wenxin(model="wenxin", baidu_api_key="Your api—key", baidu_secret_key="Your secret-key") #注意需要填下自己的key
rag_chain = (
    {"context": retriever,  "question": RunnablePassthrough()} 
    | prompt 
    | llm
    | StrOutputParser() 
)
query = "小米15ultra 升级了哪些地方？" # 用户提出的问题
print(rag_chain.invoke(query)) # 打印回答内容