引言

在当今信息爆炸的时代,如何从海量文档中快速准确地获取所需信息成为一个重要课题。本文将介绍如何实现一个基于RAG(Retrieval-Augmented Generation)的文档问答系统,该系统能够处理多种格式的文档,进行语义搜索,并结合大语言模型提供精准回答。

系统架构概述

我们的系统主要由以下几个模块组成:

  1. 文档处理模块:读取多种格式的文档

  2. 文本分片模块:将长文本分割为合适大小的片段

  3. 向量检索模块:实现语义相似度搜索

  4. 问答生成模块:结合检索结果生成回答

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)

性能优化与扩展

  1. 并行处理:使用ThreadPoolExecutor并行处理文档嵌入

  2. 缓存机制:缓存已处理的文档嵌入结果

  3. 支持更多格式:可扩展支持PPT、Excel等格式

  4. 混合检索:结合关键词检索和向量检索

结语

本文介绍了一个完整的RAG系统实现,涵盖了从文档处理、文本分片到向量检索和问答生成的完整流程。该系统具有以下优势:

  1. 支持多种文档格式

  2. 本地化部署,保护数据隐私

  3. 语义级别的检索能力

  4. 可扩展性强

读者可以根据实际需求调整参数,如分片大小、相似度阈值等,以获得最佳效果。完整代码已在上文中给出,欢迎尝试实现和改进。

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐