chatglm3实现RAG增强检索

下载模型

从huggingface或镜像网站下载生成式语言模型及ice_text.model文件，如chatglm3-6b-32k
https://huggingface.co/THUDM/chatglm3-6b-32k/tree/mainhttps://huggingface.co/THUDM/chatglm-6b/blob/main/ice_text.model

从github下载chatglm3的代码，解压后在项目的主目录建立model等文件夹放自己下载的模型，并将ice_text.model也放在模型文件夹

从huggingface或镜像网站下载生成文本嵌入模型，如bert-base-nli-mean-tokens ，其中1_Pooling、onnx和openvino可以不下载，否则加载模型时可能报错，下载后在项目的主目录建立其他文件夹放自己下载的模型，比如RAG(li)/embedding_model/bert-base-nli-mean-tokens

sentence-transformers/bert-base-nli-mean-tokens at main

从网上找一个文本或者自己编写，这里保存格式以docx为例，并保存在自己的文件夹。

代码实现

在RAG(li)建立rag.py，准备实现RAG增强检索。

首先，下载和导入必要的框架、包和模块。由于这里要下载的模块较多，基本是pip install 模块名可以完成，这里不做演示，注意像提示缺失docx模块时应该下载python-docx。另外，如果你不想用docx文件的话，是要导入对应的Loader，比如pdf对应PyPDFLoader。

from langchain_community.document_loaders import UnstructuredWordDocumentLoader#PyPDFLoader,Docx2txtLoader,TextLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_chroma import Chroma

from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from langchain_community.llms import HuggingFacePipeline

from transformers import AutoTokenizer, AutoModelForCausalLM,pipeline

加载文档，导入文本嵌入模型，其中embedding_model_path是你的文本嵌入模型路径。

# 加载文档
loader = UnstructuredWordDocumentLoader(r'E:\Project\ML\ChatGLM3\RAG(li)\file\中华人民共和国商标法.docx')
documents = loader.load()

embedding_model_path = r'E:/Project/ML/ChatGLM3/RAG(li)/embedding_model/bert-base-nli-mean-tokens'

# 定义文本嵌入模型
local_embeddings = HuggingFaceEmbeddings(model_name=embedding_model_path.replace('\\','/').split('/')[-1],cache_folder=embedding_model_path, model_kwargs = {'device': 'cuda'},encode_kwargs={'normalize_embeddings': True, 'batch_size': 16})  

创建向量存储，编写模板，并创建聊天提示对象，再将向量存储转换为检索器

template = """你是AI助手，请根据检索的上下文回答问题。如果不知道答案，请说我不知道，最多使用五句话，保持回答简结。
        {context}

        回答问题: {question}
    """

prompt = ChatPromptTemplate.from_template(template)#根据模板创建聊天提示对象，用于生成对话内容

retriever = vectorstore.as_retriever()#将向量存储转换为检索器，用于基于查询向量进行相关文档检索

加载本地生成模型和分词器，创建文本生成管道，并包装成Hugging Face语言模型对象，然后创建任务链对象，注意如果没有cuda改为cpu。

# 加载本地生成模型和分词器
model_path = r'E:\Project\ML\ChatGLM3\model\chatglm3-6b-32k'
tokenizer = AutoTokenizer.from_pretrained(model_path,trust_remote_code=True)#,model_max_length=512
model = AutoModelForCausalLM.from_pretrained(model_path,trust_remote_code=True)

#使用Hugging Face的pipeline函数创建文本生成管道，配置模型、tokenizer，以及生成参数。
pipe = pipeline(
    "text-generation", #chatglm传入这个参数不起作用，会有警告
    model=model, 
    tokenizer=tokenizer, 
    max_new_tokens=64, #64
    # do_sample=True,#是否使用采样，如果为False，则每次生成都是相同的，temperature参数无效
    # temperature=0.5,#0.0-1.0温度，值越大，生成结果越随机  
    # repetition_penalty=1.0,#重复惩罚，值越大，重复结果越少
    # top_k=50,#top_k采样，值越大，生成结果越随机
    # top_p=0.95,#top_p采样，值越大，生成结果越随机
    # max_length=1000,#最大长度
    # min_length=100,#最小长度
    # early_stopping=True,#是否提前停止
    # length_penalty=1.0,#长度惩罚，值越大，生成结果越短
    # num_beams=1,#束搜索，值越大，生成结果越稳定
    # num_beam_groups=1,#束搜索组数，值越大，生成结果越稳定
    device='cuda:0',  # 设置设备为GPU
)

llm = HuggingFacePipeline(pipeline=pipe)#将生成管道包装成Hugging Face语言模型对象，用于与任务链集成和调用。

chain = LLMChain(llm=llm, prompt=prompt)# 创建任务链对象，将LLM与提示模板结合，以生成基于输入上下文的自然语言输出。

提出问题，从向量存储检索与输入问题相似的文档集合，这里k也可设为1，将查询文档转化为上下文字符串。

question = "以三维标志申请注册商标什么形状不得注册？"#
docs = vectorstore.similarity_search(question,k=5)#从向量存储中检索与输入问题最相似的文档集合。
# 确保将文档内容格式化为字符串
#context = "\n".join(doc.page_content for doc in docs)
max_length = 512  #512 设置最大长度
context = "\n".join(doc.page_content for doc in docs).replace('\n\n', '\n')#可限制长度doc.page_content[:max_length]

运行问答链，打印响应和答案，注意提取答案与你的模板设置有关。

# 运行问答链
response = chain.invoke(input={"context": context, "question": question})
print('response：\n',response)
# print('context：\n',response['context'])
text=response['text'].replace('\n\n', '\n')
# print('text：\n',text)
#print("type(response['text'])",type(response['text']))#<class 'str'>

answer=text.split(f'回答问题: {question}')[-1].strip()
print('answer：\n', answer)

部分输出情况如下。

创作不易，禁止抄袭，转载请附上原文链接及标题