DeepSeek + Faiss搭建本地知识库检索

# ChatPDF-Faiss智能文档问答


## 概述

 
ChatPDF-Faiss是一个基于大语言模型的PDF智能问答系统，能够对PDF文档内容进行语义理解和问答。系统使用FAISS向量数据库技术实现高效的相似度检索，并结合大语言模型提供准确的问答服务。

 
## 核心功能

 

### 1. PDF文本提取与页码关联

- 从PDF中提取文本内容

- 记录每段文本对应的原始页码信息

- 便于回溯答案的具体出处

 

### 2. 文本分块与向量化处理

- 使用RecursiveCharacterTextSplitter将长文本分割成适合向量化的小块

- 通过DashScope的文本嵌入模型将文本块转化为语义向量

- 构建FAISS向量索引，实现高效的相似度检索

 

### 3. 知识库持久化存储

- 将向量数据库保存到本地磁盘

- 存储文本块与对应页码的映射关系

- 支持后续快速加载，无需重复处理文档

 

### 4. 语义检索与问答

- 基于用户查询进行语义相似度匹配

- 检索与问题最相关的文本块

- 使用大语言模型(通义千问/DeepSeek)生成准确的答案

- 提供答案来源页码，增强可信度与可追溯性

 

## 技术栈

-**PDF处理**：PyPDF2

-**文本分割**：LangChain RecursiveCharacterTextSplitter

-**向量嵌入**：DashScopeEmbeddings (阿里云通义千亿级开源向量模型)

-**向量数据库**：FAISS (Facebook AI Similarity Search)

-**大语言模型**：通义/DeepSeek系列模型

-**问答链**：LangChain QA Chain

 

## 工作流程

 

1.**文档处理阶段**：

   - 提取PDF文本内容与页码信息

   - 文本分割为适当大小的块

   - 向量化处理并构建FAISS索引

   - 存储向量数据库和页码映射信息

 

2.**问答查询阶段**：

   - 接收用户问题

   - 问题向量化并在FAISS中检索相似文本块

   - 将相关文本块和问题一起传入大语言模型

   - 生成答案并提供来源页码

 

## 系统优势

 

-**高性能检索**：FAISS向量库支持亿级向量的快速相似度搜索

-**精准溯源**：回答问题时提供具体页码来源

-**灵活扩展**：支持不同的嵌入模型和大语言模型

-**本地存储**：向量数据库可持久化存储，支持增量更新

 

## 使用方法

 

1. 设置环境变量`DASHSCOPE_API_KEY`

2. 准备PDF文件，命名为"file.pdf"放在当前目录

3. 运行脚本，系统会自动处理文档并创建向量数据库

4. 设置查询问题，获取智能回答及来源页码 

 

 

from PyPDF2 import PdfReader
from langchain.chains.question_answering import load_qa_chain
from langchain_community.callbacks.manager import get_openai_callback
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
from typing import List, Tuple
import os
import pickle

DASHSCOPE_API_KEY = os.getenv('DASHSCOPE_API_KEY')
if not DASHSCOPE_API_KEY:
    raise ValueError("请设置环境变量 DASHSCOPE_API_KEY")

def extract_text_with_page_numbers(pdf) -> Tuple[str, List[int]]:
    """
    从PDF中提取文本并记录每行文本对应的页码
    
    参数:
        pdf: PDF文件对象
    
    返回:
        text: 提取的文本内容
        page_numbers: 每行文本对应的页码列表
    """
    text = ""
    page_numbers = []

    for page_number, page in enumerate(pdf.pages, start=1):
        extracted_text = page.extract_text()
        if extracted_text:
            text += extracted_text
            page_numbers.extend([page_number] * len(extracted_text.split("\n")))

    return text, page_numbers

def process_text_with_splitter(text: str, page_numbers: List[int], save_path: str = None) -> FAISS:
    """
    处理文本并创建向量存储
    
    参数:
        text: 提取的文本内容
        page_numbers: 每行文本对应的页码列表
        save_path: 可选，保存向量数据库的路径
    
    返回:
        knowledgeBase: 基于FAISS的向量存储对象
    """
    # 创建文本分割器，用于将长文本分割成小块
    text_splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", ".", " ", ""],
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
    )

    # 分割文本
    chunks = text_splitter.split_text(text)
    print(f"文本被分割成 {len(chunks)} 个块。")
        
    # 创建嵌入模型
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v1",
        dashscope_api_key=DASHSCOPE_API_KEY,
    )
    
    # 从文本块创建知识库
    knowledgeBase = FAISS.from_texts(chunks, embeddings)
    print("已从文本块创建知识库。")
    
    # 改进：存储每个文本块对应的页码信息
    # 创建原始文本的行列表和对应的页码列表
    lines = text.split("\n")
    
    # 为每个chunk找到最匹配的页码
    page_info = {}
    for chunk in chunks:
        # 查找chunk在原始文本中的开始位置
        start_idx = text.find(chunk[:100])  # 使用chunk的前100个字符作为定位点
        if start_idx == -1:
            # 如果找不到精确匹配，则使用模糊匹配
            for i, line in enumerate(lines):
                if chunk.startswith(line[:min(50, len(line))]):
                    start_idx = i
                    break
            
            # 如果仍然找不到，尝试另一种匹配方式
            if start_idx == -1:
                for i, line in enumerate(lines):
                    if line and line in chunk:
                        start_idx = text.find(line)
                        break
        
        # 如果找到了起始位置，确定对应的页码
        if start_idx != -1:
            # 计算这个位置对应原文中的哪一行
            line_count = text[:start_idx].count("\n")
            # 确保不超出页码列表长度
            if line_count < len(page_numbers):
                page_info[chunk] = page_numbers[line_count]
            else:
                # 如果超出范围，使用最后一个页码
                page_info[chunk] = page_numbers[-1] if page_numbers else 1
        else:
            # 如果无法匹配，使用默认页码-1（这里应该根据实际情况设置一个合理的默认值）
            page_info[chunk] = -1
    
    knowledgeBase.page_info = page_info
    
    # 如果提供了保存路径，则保存向量数据库和页码信息
    if save_path:
        # 确保目录存在
        os.makedirs(save_path, exist_ok=True)
        
        # 保存FAISS向量数据库
        knowledgeBase.save_local(save_path)
        print(f"向量数据库已保存到: {save_path}")
        
        # 保存页码信息到同一目录
        with open(os.path.join(save_path, "page_info.pkl"), "wb") as f:
            pickle.dump(page_info, f)
        print(f"页码信息已保存到: {os.path.join(save_path, 'page_info.pkl')}")

    return knowledgeBase

def load_knowledge_base(load_path: str, embeddings = None) -> FAISS:
    """
    从磁盘加载向量数据库和页码信息
    
    参数:
        load_path: 向量数据库的保存路径
        embeddings: 可选，嵌入模型。如果为None，将创建一个新的DashScopeEmbeddings实例
    
    返回:
        knowledgeBase: 加载的FAISS向量数据库对象
    """
    # 如果没有提供嵌入模型，则创建一个新的
    if embeddings is None:
        embeddings = DashScopeEmbeddings(
            model="text-embedding-v1",
            dashscope_api_key=DASHSCOPE_API_KEY,
        )
    
    # 加载FAISS向量数据库，添加allow_dangerous_deserialization=True参数以允许反序列化
    knowledgeBase = FAISS.load_local(load_path, embeddings, allow_dangerous_deserialization=True)
    print(f"向量数据库已从 {load_path} 加载。")
    
    # 加载页码信息
    page_info_path = os.path.join(load_path, "page_info.pkl")
    if os.path.exists(page_info_path):
        with open(page_info_path, "rb") as f:
            page_info = pickle.load(f)
        knowledgeBase.page_info = page_info
        print("页码信息已加载。")
    else:
        print("警告: 未找到页码信息文件。")
    
    return knowledgeBase

def answer_question(knowledgeBase, llm, query, k=2):
    """
    回答用户问题
    
    参数:
        knowledgeBase: FAISS向量数据库
        llm: 大语言模型
        query: 用户问题
        k: 检索的文档数量
    
    返回:
        answer: 回答内容
        sources: 来源页码列表
    """
    # 执行相似度搜索，找到与查询相关的文档
    docs = knowledgeBase.similarity_search(query, k=k)
    
    # 加载问答链
    chain = load_qa_chain(llm, chain_type="stuff")
    
    # 准备输入数据
    input_data = {"input_documents": docs, "question": query}
    
    # 使用回调函数跟踪API调用成本
    with get_openai_callback() as cost:
        # 执行问答链
        response = chain.invoke(input=input_data)
        print(f"查询已处理。成本: {cost}")
    
    # 记录唯一的页码
    unique_pages = set()
    
    # 收集每个文档块的来源页码
    for doc in docs:
        text_content = getattr(doc, "page_content", "")
        source_page = knowledgeBase.page_info.get(
            text_content.strip(), "未知"
        )
        
        if source_page not in unique_pages:
            unique_pages.add(source_page)
    
    return response["output_text"], list(unique_pages)

def initialize_system():
    """初始化系统，加载或创建知识库"""
    save_dir = "./vector_db"
    
    # 检查是否已经存在向量数据库
    if os.path.exists(save_dir) and os.path.isdir(save_dir):
        print("检测到已有向量数据库，正在加载...")
        try:
            # 创建嵌入模型
            embeddings = DashScopeEmbeddings(
                model="text-embedding-v1",
                dashscope_api_key=DASHSCOPE_API_KEY,
            )
            # 从磁盘加载向量数据库
            knowledgeBase = load_knowledge_base(save_dir, embeddings)
            print("成功加载已有知识库！")
        except Exception as e:
            print(f"加载失败: {e}")
            print("将创建新的知识库...")
            knowledgeBase = create_new_knowledge_base(save_dir)
    else:
        print("未检测到向量数据库，将创建新的知识库...")
        knowledgeBase = create_new_knowledge_base(save_dir)
    
    # 初始化大语言模型
    from langchain_community.llms import Tongyi
    llm = Tongyi(model_name="deepseek-v3", dashscope_api_key=DASHSCOPE_API_KEY)
    
    return knowledgeBase, llm

def create_new_knowledge_base(save_dir):
    """创建新的知识库"""
    # 检查PDF文件是否存在
    pdf_path = './file.pdf'
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"未找到PDF文件: {pdf_path}，请确保文件存在！")
    
    # 读取PDF文件
    pdf_reader = PdfReader(pdf_path)
    
    # 提取文本和页码信息
    text, page_numbers = extract_text_with_page_numbers(pdf_reader)
    print(f"提取的文本长度: {len(text)} 个字符。")
    
    # 处理文本并创建知识库，同时保存到磁盘
    knowledgeBase = process_text_with_splitter(text, page_numbers, save_path=save_dir)
    return knowledgeBase

def chat_interface():
    """交互式聊天界面"""
    print("\n" + "="*50)
    print("欢迎使用 ChatPDF-Faiss 智能文档问答系统")
    print("="*50)
    
    # 初始化系统
    try:
        knowledgeBase, llm = initialize_system()
    except Exception as e:
        print(f"系统初始化失败: {e}")
        return
    
    print("\n系统已准备就绪！您可以开始提问了。")
    print("输入 'exit' 或 'quit' 退出对话。\n")
    
    # 开始对话循环
    while True:
        # 获取用户输入
        query = input("\n请输入您的问题: ")
        
        # 检查是否退出
        if query.lower() in ['exit', 'quit', '退出', '结束']:
            print("感谢使用，再见！")
            break
        
        # 如果输入为空，继续循环
        if not query.strip():
            continue
        
        try:
            # 回答问题
            answer, sources = answer_question(knowledgeBase, llm, query)
            
            # 打印回答
            print("\n回答:")
            print(answer)
            
            # 打印来源
            if sources:
                print("\n来源页码:", end=" ")
                print(", ".join(str(page) for page in sources if page != -1))
        except Exception as e:
            print(f"\n处理问题时出错: {e}")
            print("请尝试重新提问或检查系统配置。")

if __name__ == "__main__":
    chat_interface()