使用Python构建聊天机器人(后端)

文档:后端(backend.py)
该脚本将PDF文档处理为向量嵌入,并构建FAISS索引以进行语义搜索。
这是Indaba RAG聊天机器人的离线预处理管道。

主要职责

  1. 从文件夹中加载PDF。
  2. 使用PyPDF2提取原始文本。
  3. 将大型文档分块为较小的重叠文本段。
  4. 使用SentenceTransformers将块转换为嵌入。
  5. 构建并持久化FAISS索引以进行相似性搜索。
  6. 保存原始块以便后续检索。

逐步分解

导入和设置

import os
import pickle
import numpy as np
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
import faiss

os → 文件系统操作。
pickle → 保存预处理的块。
numpy → 数值数组处理。
PyPDF2 → 从PDF文件中提取文本。
SentenceTransformer → 嵌入模型(all-MiniLM-L6-v2)。
faiss → 高效的相似性搜索。

常量:

embedder = SentenceTransformer("all-MiniLM-L6-v2")
INDEX_FILE = "faiss_index.bin"
CHUNKS_FILE = "chunks.pkl"
    • embedder 是模型实例;加载此实例时,会下载模型权重(首次运行可能需要一些时间)。
    • INDEX_FILECHUNKS_FILE 定义了保存 FAISS 索引和数据块的位置。

 

加载 PDF 的函数

def load_pdf(file_path):
    pdf = PdfReader(file_path)
    text = ""
    for page in pdf.pages:
        text += page.extract_text() + "\n"
    return text
  • 使用 PyPDF2 读取 PDF 文件。
  • 逐页提取文本。
  • 将完整文档文本作为字符串返回。

文本分块功能

def chunk_text(text, chunk_size=500, overlap=100):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks
  • 将文本分割为每个大小为 chunk_size 字符的块,每次移动 chunk_size – overlap(因此连续的块之间重叠 overlap 字符)。
  • 表示方式:
    块 1 = 0–500
    块 2 = 400–900(重叠 100)

完整管道信息

函数的逐步讲解:

1. 收集所有 PDF 的块:

pdf_folder = "vault"   
 #这是存储pdf的文件夹/路径。

all_chunks = []
for filename in os.listdir(pdf_folder):
    if filename.endswith(".pdf"):
        text = load_pdf(os.path.join(pdf_folder, filename))
        chunks = chunk_text(text)
        all_chunks.extend(chunks)
  • 提取每个PDF的文本和块,并将所有块作为字符串保存在all_chunks列表中。

注意:顺序很重要(索引ID与顺序对齐)。

2. 嵌入块:

vectors = embedder.encode(all_chunks)
   vectors = np.array(vectors)
  • embedder.encode(list_of_texts) 返回一个向量的列表/数组。默认情况下,根据版本返回 float32 或 float64 — FAISS 期望使用 float32。在实际操作中,强制使用 dtype 为 float32 更安全。
  • 重要提示:一次性嵌入所有块可能会导致内存溢出(OOM),如果你有很多块。请使用批处理:
 vectors = embedder.encode(all_chunks, batch_size=32, show_progress_bar=True)
 vectors = np.array(vectors).astype('float32')

 

3. 创建 FAISS 索引:

dim = vectors.shape[1]
   index = faiss.IndexFlatL2(dim)
   index.add(vectors)

(基础)

  • 创建一个 FAISS 索引并将所有块向量添加到索引中。

(技术)
IndexFlatL2 = 使用 L2 距离进行精确(暴力)最近邻搜索。适用于小到中等规模的集合。

  • 优点:简单且精确。
  • 缺点:在大规模集合上速度较慢。

index.add(vectors) 以与 all_chunks 相同的顺序添加向量。FAISS 内部 ID = 0..N-1 以该顺序 — 这就是你如何映射回块。

4. 保存索引和块:

  faiss.write_index(index, INDEX_FILE)
   with open(CHUNKS_FILE, "wb") as f:
       pickle.dump(all_chunks, f)

将 FAISS 索引保存到 faiss_index.bin。
将文本块(原始文本)保存到 chunks.pkl。

这些文件稍后将在运行时由 Streamlit 前端加载。

如何运行此脚本

确保您的 PDF 文件在 docs/ 目录中

python -m backend

输出:faiss_index.bin 和 chunks.pkl 到当前目录

更多