RAG的文档问答系统-PYTHON
引言
在当今信息爆炸的时代,如何从海量文档中快速准确地获取所需信息成为一个重要课题。本文将介绍如何实现一个基于RAG(Retrieval-Augmented Generation)的文档问答系统,该系统能够处理多种格式的文档,进行语义搜索,并结合大语言模型提供精准回答。
系统架构概述
我们的系统主要由以下几个模块组成:
-
文档处理模块:读取多种格式的文档
-
文本分片模块:将长文本分割为合适大小的片段
-
向量检索模块:实现语义相似度搜索
-
问答生成模块:结合检索结果生成回答
1. 文档处理模块实现
文档处理模块需要能够处理多种常见文档格式,包括TXT、PDF、DOC和DOCX等。我们使用以下技术栈:
python
复制
下载
import pdfplumber # 处理PDF文件 from docx import Document as DocxDocument # 处理DOCX文件 import win32com.client as win32 # 处理老版DOC文件
关键代码实现:
python
复制
下载
@staticmethod
def read_directory_files(directory_path: str) -> List[str]:
result_list = []
path = Path(directory_path)
for file_path in path.iterdir():
if file_path.is_file():
try:
file_content = ""
if file_path.suffix.lower() == '.txt':
file_content = ReadFiles.read_text_file(file_path)
elif file_path.suffix.lower() == '.docx':
file_content = ReadFiles.read_docx_file(file_path)
elif file_path.suffix.lower() == '.doc':
file_content = ReadFiles.read_doc_file(file_path)
elif file_path.suffix.lower() == '.pdf':
file_content = ReadFiles.read_pdf_file(file_path)
# ...省略后续处理...
技术要点:
-
使用Path对象进行跨平台路径操作
-
根据文件后缀选择不同的解析方法
-
对每种文件格式使用专门的解析库
2. 文本分片策略
为了便于后续的向量化处理,我们需要将长文档分割为适当大小的文本块。这里我们实现了基于句子的分片策略:
python
复制
下载
@staticmethod
def split_by_sentence_within_limit(input_str: str, max_number: int) -> List[str]:
result = []
if not input_str:
return result
current_position = 0
total_length = len(input_str)
while current_position < total_length:
end_position = min(current_position + max_number, total_length)
current_chunk = input_str[current_position:end_position]
# 查找最后一个句号作为分割点
last_dot_index = current_chunk.rfind('。')
# ...省略后续处理...
分片策略优势:
-
保持语义完整性:尽量在句子边界处分割
-
控制分片大小:确保每个分片不超过最大长度限制
-
处理中英文混合文本:支持中文句号作为分割点
3. 向量检索模块
向量检索是RAG系统的核心,我们实现了基于Ollama的嵌入函数和相似度计算:
python
复制
下载
class OllamaEmbeddingFunction(EmbeddingFunction):
def __init__(self, model: str, url: str):
self.model = model
self.url = url
self.session = requests.Session()
def generate_embedding(self, text: str) -> List[float]:
request_body = {
"model": self.model,
"prompt": text
}
# ...发送请求获取嵌入向量...
class Collection:
def query(self, query_text: str, top_k: int) -> List[Result]:
query_embedding = self.embedding_function.generate_embedding(query_text)
results = []
for doc in self.documents:
similarity = self.cosine_similarity(query_embedding, doc.get_embedding())
results.append(Result(doc, similarity))
# ...返回相似度最高的结果...
关键技术点:
-
使用Ollama本地模型生成嵌入向量
-
余弦相似度计算文档相关性
-
支持多线程处理文档嵌入
4. 问答生成流程
完整的问答流程整合了上述所有模块:
python
复制
下载
def rag_work(input_content_list: List[str], max_number: int, query_text: str, top_k: int,
similarity_threshold: float) -> str:
# 初始化嵌入函数和集合
embedding_func = OllamaEmbeddingFunction("bge-m3:latest", "http://localhost:11434/api/embeddings")
collection = Collection("knowledge-base", embedding_func)
# 处理并添加文档
total_document_list = []
for i, content in enumerate(input_content_list):
split_result = SubString.split_by_sentence_within_limit(content, max_number)
for j, chunk in enumerate(split_result):
doc = Document(f"{i}{j}", chunk)
total_document_list.append(doc)
# 执行查询
collection.add_documents(total_document_list)
result_list = collection.query(query_text, top_k)
# 组装相关知识
res = []
for result in result_list:
if result.get_similarity() > similarity_threshold:
res.append(result.get_document().get_content())
return '\n'.join(res)
系统集成与使用示例
最后,我们将所有模块集成起来,实现完整的问答功能:
python
复制
下载
if __name__ == "__main__":
# 1. 读取文档
input_content_list = ReadFiles.read_files_work("resources/")
# 2. 用户提问
query_text = "小马最终过河了吗?"
# 3. 检索相关知识
knowledge = Rag.rag_work(input_content_list, 500, query_text, 5, 0.35)
# 4. 构造提示词
content = f"""你的角色是:问题答复专家。
你的任务是:根据已知指示答复我的问题。
已知知识:
{knowledge}
根据以上要求,我输入的问题是:{query_text}"""
# 5. 调用语言模型生成回答
Ds_V3.chat_with_openai(content)
性能优化与扩展
-
并行处理:使用ThreadPoolExecutor并行处理文档嵌入
-
缓存机制:缓存已处理的文档嵌入结果
-
支持更多格式:可扩展支持PPT、Excel等格式
-
混合检索:结合关键词检索和向量检索
结语
本文介绍了一个完整的RAG系统实现,涵盖了从文档处理、文本分片到向量检索和问答生成的完整流程。该系统具有以下优势:
-
支持多种文档格式
-
本地化部署,保护数据隐私
-
语义级别的检索能力
-
可扩展性强
读者可以根据实际需求调整参数,如分片大小、相似度阈值等,以获得最佳效果。完整代码已在上文中给出,欢迎尝试实现和改进。
更多推荐
所有评论(0)