基于 RAG 的小说智能问答系统实现

RAGLLMLlamaIndexpgvectorPython

其实以前看小说的时候就有过类似的想法，把小说输入到AI系统，可以和系统进行剧情讨论，这样岂不是很有意思？只不过原来没有思路，也不知道这就是 RAG，所以一直没有实现。经过一段时间的学习，有了一些理解，所以想自己实现一个基于 RAG 的小说智能问答系统，一方面是作为面试准备，另一方面是作为学习笔记。

引言

检索增强生成（Retrieval-Augmented Generation，RAG）是一种结合了信息检索和文本生成的 AI 技术。本文将以解析小说为例，详细介绍 RAG 的实现过程。 在文章的最后，我们将构建一个能够智能回答小说相关问题的系统。

技术栈

  • 后端：FastAPI、PGVector、Llamaindex
  • 向量嵌入模型：shibing624/text2vec-base-chinese
  • 查询模型：gpt-3.5-turbo
  • 数据库：PostgreSQL

项目地址：

系统工作流程

  1. 用户发送问题到/chat 接口。
  2. 系统使用查询引擎在向量存储中检索相关文本。
  3. 检索到的文本与用户问题一起发送给 LLM 模型。
  4. LLM 模型生成回答。
  5. 系统返回生成的回答给用户。

RAG 的核心概念

RAG 主要包含以下几个关键组件：

  1. 文档解析：将原始文本转换为结构化数据。
  2. 向量化：将文本转换为数值向量。
  3. 向量存储：高效存储和检索向量化后的文本。
  4. 检索：根据用户查询找到相关文本。
  5. 生成：基于检索结果生成回答。

实现步骤

1. 文档解析

首先，我们需要将 EPUB 格式的小说解析成文本文件。这一步骤在parse_epub.py中实现：

import os
from ebooklib import epub
from bs4 import BeautifulSoup
import re
import dotenv

dotenv.load_dotenv()


def clean_text(text):
    text = re.sub(r"<[^>]+>", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def parse_epub(epub_path, output_dir):
    book = epub.read_epub(epub_path)
    # ... 解析EPUB文件的代码 ...


if __name__ == "__main__":
    epub_path = os.getenv("NOVEL_FILE")
    output_dir = os.getenv("NOVEL_FILE_DIR")
    parse_epub(epub_path, output_dir)
    print("EPUB parsing completed!")

这个脚本将 EPUB 文件解析为单独的章节文本文件，便于后续处理。

2. 向量化和存储

接下来，我们使用shibing624/text2vec-base-chinese模型需要将解析后的文本向量化并存储。这一步骤在 embed_novels.py 中实现：

from llama_index.core import Document, VectorStoreIndex
from llama_index.vector_stores.postgres import PGVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 设置嵌入模型
Settings.embed_model = HuggingFaceEmbedding(
    model_name=os.getenv("EMBEDDING_MODEL_NAME", "shibing624/text2vec-base-chinese"),
    embed_batch_size=BATCH_SIZE,
    device=device,
)

# 初始化向量存储
vector_store = PGVectorStore.from_params(
    # ... 数据库配置 ...
)


def embed_novels():
    # ... 向量化和存储的代码 ...

这个脚本将解析后的章节文本向量化并存储到数据库中。

3. 检索和生成

最后，我们需要实现检索和生成功能。这一步骤在 novel.py 中实现：

from fastapi import FastAPI
from llama_index.core import VectorStoreIndex
from llama_index.llms.openai import OpenAI

# 初始化FastAPI应用
app = FastAPI()

# 设置LLM模型
llm = OpenAI(
    model="gpt-3.5-turbo",
    temperature=0.7,
    max_tokens=1024,
    system_prompt="你是一个小说阅读AI助手，请用简洁、准确的中文回答所有问题。",
)

# 创建查询引擎
query_engine = index.as_query_engine(llm=llm)

@app.post("/chat")
async def chat(request: ChatRequest):
    # 使用查询引擎获取响应
    response = query_engine.query(request.message)
    # ... 处理响应的代码 ...

4. 运行系统

在项目根目录下运行：

python app.py

使用httpie测试：

$ http POST http://localhost:8000/chat message="小说讲了一个什么故事"
HTTP/1.1 200 OK
content-length: 144
content-type: application/json
date: Mon, 14 Oct 2024 08:49:47 GMT
server: uvicorn

{
    "response": "这部小说讲述了关于球状闪电的目击者的叙述记录，以及一些关于宏物质和宏原子核的研究工作。"
}

$ http POST http://localhost:8000/chat message="小说里最后发现球状闪电的本质是什么"
HTTP/1.1 200 OK
content-length: 228
content-type: application/json
date: Mon, 14 Oct 2024 08:49:31 GMT
server: uvicorn

{
    "response": "球状闪电的本质是通过操纵宏电子释放的能量，具有精细的选择性，可以专门烧毁动物的骨骼或汽化动物的血液，而不伤及其肌肉组织，其攻击方式可怕而精密。"
}

这个脚本创建了一个 FastAPI 应用，使用 OpenAI 的 GPT-3.5-turbo 模型作为生成模型，并实现了一个聊天接口。

最后

通过实现这个基于 RAG 的小说智能问答系统，我们展示了如何将复杂的文本数据转化为可查询的知识库。这种方法不仅适用于小说，还可以扩展到其他领域，如技术文档、学术论文等。RAG 技术的优势在于它能够结合大规模预训练模型的生成能力和专门知识库的精确性，为用户提供更加准确和相关的回答。

实际上这只能算是一个 Demo，实际使用中，发现有不少问题系统回答的并不太好，需要进一步优化。

作者: Lerry
发表时间:2024-10-14