RAG灌库，帮助大模型背诵古诗-CSDN博客

本文链接：https://blog.csdn.net/disbubu/article/details/142578822

前言：

我尝试使用chat-gpt和Chatglm-6b去鉴赏古诗文，在他们能正确的知道古诗的时候，鉴赏能力还是不错的，但是他们很多时候会将古诗背诵错误，诗人或者诗句背诵错误，并且他们还非常自信。所以，我使用RAG给大模型增加一下知识库，这样就不需要微调了（我的最终目标是想部署一个可以精确鉴赏诗歌的模型，目前也在努力中）。

1.古诗资料的获取

自己一首一首的收集古诗显然是不太可能的，所以我在网上找到了古诗词的资料https://github.com/chinese-poetry/chinese-poetry?tab=readme-ov-file。这里面的诗歌内容非常丰富。

2.代码部分

我这里不是用的本地的模型，用的是chat-gpt的api，首先是配置一下环境，

_ = load_dotenv(find_dotenv())
api_key = os.getenv("api_key")
base_url = os.getenv("base_url")

llm = ChatOpenAI(
    api_key=api_key,
    base_url=base_url,
)

配置环境的目的是为了确保你的api不会被别人使用，可以隐藏一下。

query1 = "背诵一下《在岳咏蝉》"
# def search(query: str) -> None:
"""
:param query: 输入用户提出的问题
:return: 返回RAG搜索的结果
"""
loader = TextLoader('C:\\Users\\86139\\Desktop\\model\\test.json', encoding='utf-8')
pages = loader.load()
# 用于切分文章
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=200,
    chunk_overlap=100,
    length_function=len,
    add_start_index=True,
)

# 灌库
embeddings = OpenAIEmbeddings(model="text-embedding-ada-002",
                              api_key=api_key,
                              base_url=base_url,
                              )
db = FAISS.from_documents(pages, embeddings)

# 检索 top-2 结果
retriever = db.as_retriever(search_kwargs={"k": 1})

docs = retriever.invoke(query1)

我这里灌库的文件是json文件，pdf格式也是可以转化的，但是 TextLoader()需要换一下，换成pdf的。代码逻辑还是非常简单的，我写的也比较潦草。然后如果这样写的时候，每次提问都需要转化向量，速度慢，api也消耗快。所以可以pip一下chromadb这个库吧，可以满足基本要求了。另外这里也用到了prompt控制大模型的输出。

prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。

已知信息:
{context}

用户问：
{query}

如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。
请不要输出已知信息中不包含的信息或答案。
请用中文回答用户问题。
"""

这样的写法可以有效的避免模型的幻觉，然后也可以告诉模型什么是我需要的回答。最后就是结果输出部分了。

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 = v
        inputs[k] = val
    return prompt_template.format(**inputs)


prompt = build_prompt(prompt_template, context=docs, query=query1)
# print(docs)
# print(retriever)
response = llm.invoke(prompt)
# print(prompt)
print(response.content)