文档:后端(backend.py)
该脚本将PDF文档处理为向量嵌入,并构建FAISS索引以进行语义搜索。
这是Indaba RAG聊天机器人的离线预处理管道。
主要职责
- 从文件夹中加载PDF。
- 使用PyPDF2提取原始文本。
- 将大型文档分块为较小的重叠文本段。
- 使用SentenceTransformers将块转换为嵌入。
- 构建并持久化FAISS索引以进行相似性搜索。
- 保存原始块以便后续检索。
逐步分解
导入和设置
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_FILE 和 CHUNKS_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 到当前目录