【学习记录】从 PDF 到 FAISS 向量索引:构建本地 RAG 数据预处理流水线
在搭建本地 RAG(检索增强生成)系统时,第一步往往是将 PDF 文档转换为可检索的向量索引。然而,PDF 文件既有可直接提取文本的电子文档,也有扫描图片型的文档。本文基于 Python 实现了一套完整的流水线:自动判断文档类型,必要时调用 OCR,然后进行文本分块、向量化(使用中文嵌入模型)并构建 FAISS 索引。代码完全开源,可直接用于本地知识库的预处理。
📌 目录
- 代码功能概述
- 环境配置与依赖
- 输入与输出
- 核心逻辑解析
- 4.1 PDF 文本提取与 OCR 回退
- 4.2 文本分块
- 4.3 向量化与 FAISS 索引构建
- 4.4 索引持久化
- 注意事项与优化建议
- 使用示例:加载索引并查询
- 总结
1️⃣ 代码功能概述
本代码实现了一个端到端的数据预处理流水线:
- 读取 PDF 文件(支持文本型 PDF 和扫描图片型 PDF)。
- 自动判断文本量,若少于 5000 字符则触发 OCR(光学字符识别)。
- 将提取的文本切分成语义连贯的节点(chunk)。
- 使用 HuggingFace 的中文嵌入模型(
BAAI/bge-small-zh-v1.5)将每个节点向量化。 - 构建 FAISS 索引(L2 距离),并持久化到本地磁盘。
最终产出一个可被直接加载的向量索引,供下游的语义搜索、问答系统使用。
2️⃣ 环境配置与依赖
Python 库依赖
bash
pip install pdfplumber pytesseract pdf2image faiss-cpu llama_index transformers huggingface_hub
| 库 | 用途 |
|---|---|
pdfplumber |
提取文本型 PDF 的文字内容 |
pytesseract |
OCR 识别图片文字 |
pdf2image |
将 PDF 页面转换为图像 |
faiss-cpu |
高效的向量索引库 |
llama_index |
文档管理、文本分块、向量索引封装 |
transformers |
HuggingFace 嵌入模型的后端 |
系统依赖(必须安装)
| 组件 | 安装命令 (Ubuntu/Debian) | 作用 |
|---|---|---|
| Tesseract OCR | sudo apt install tesseract-ocr tesseract-ocr-chi-sim |
OCR 引擎 |
| Poppler | sudo apt install poppler-utils |
提供 pdftoppm 命令,供 pdf2image 使用 |
对于 Windows 或 macOS,请参考官方文档安装对应版本。
3️⃣ 输入与输出
输入
PDF_PATH:待处理的 PDF 文件路径(例如./data/y0664.pdf)。- 配置参数:
CHUNK_SIZE = 512:每个文本块的最大字符数。CHUNK_OVERLAP = 50:相邻块之间的重叠字符数。EMBED_DIM = 512:嵌入向量维度(与模型bge-small-zh一致)。INDEX_DIR = "./faiss_index":索引保存目录。
输出
- 文本节点列表 (
nodes):每个节点包含一段文本和元数据,长度约 512 字符。 - FAISS 索引文件 :
vector_store.faiss - LlamaIndex 元数据 :保存在
INDEX_DIR目录下(包括文档存储、索引结构等)。
生成的文件可用于后续的向量检索,例如:
python
query = "公司最新政策"
results = index.query(query, top_k=5)
4️⃣ 核心逻辑解析
4.1 PDF 文本提取与 OCR 回退
python
def extract_pdf(path):
with pdfplumber.open(path) as pdf:
return "\n".join(page.extract_text() or "" for page in pdf.pages)
raw_text = extract_pdf(PDF_PATH)
if len(raw_text) < 5000: # 文本太少,判定为扫描件
raw_text = ocr_pdf(PDF_PATH)
extract_pdf使用pdfplumber尝试提取文字层。- 如果总字符数不足 5000(可根据文档长度调整阈值),则认为该 PDF 是扫描件或图片型,触发 OCR。
- OCR 流程:
pdf2image.convert_from_path将每页转为 PIL 图片,再调用pytesseract.image_to_string识别文字(语言包chi_sim+eng)。
注意:OCR 非常耗时(每页约 1~3 秒),建议在后台处理,或对页数多的文档使用多进程。
4.2 文本分块
python
from llama_index.core.node_parser import SentenceSplitter
splitter = SentenceSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
documents = [Document(text=raw_text)]
nodes = splitter.get_nodes_from_documents(documents)
SentenceSplitter会尽量在句子边界切分,避免切断语义。- 重叠(
chunk_overlap)可以防止关键信息正好落在两个块的边缘而被遗漏。
4.3 向量化与 FAISS 索引构建
python
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings, VectorStoreIndex, StorageContext
from llama_index.vector_stores.faiss import FaissVectorStore
import faiss
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5", device="cpu")
faiss_index = faiss.IndexFlatL2(EMBED_DIM)
vector_store = FaissVectorStore(faiss_index=faiss_index)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex(nodes, storage_context=storage_context)
IndexFlatL2使用欧氏距离进行精确检索,适合百万级以下的向量。- 如果数据量巨大(>100 万),可改用
IndexIVFFlat加速。
4.4 索引持久化
python
storage_context.persist(persist_dir=INDEX_DIR)
faiss.write_index(faiss_index, os.path.join(INDEX_DIR, "vector_store.faiss"))
persist保存 LlamaIndex 的元数据(节点文本、文档存储等)。faiss.write_index保存原生 FAISS 索引,方便直接加载到内存。
5️⃣ 注意事项与优化建议
| 问题 | 说明 | 解决方案 |
|---|---|---|
| OCR 速度慢 | 每页需转为图片再识别,30 页文档约需 1 分钟 | 使用多进程(concurrent.futures)或仅对无文字页进行 OCR;升级到 GPU 版 Tesseract(较复杂) |
| 内存消耗大 | 转图片 + OCR 可能占用数百 MB 内存 | 分批处理(convert_from_path(first_page, last_page)) |
| 中文 OCR 准确率 | 默认 chi_sim 对印刷体较好,但对模糊图片有误差 |
提高图片 DPI(如 300)或使用更专业的 OCR(如 PaddleOCR) |
| FAISS 索引限制 | IndexFlatL2 精确但速度随数据量线性增长 |
使用 IndexIVFFlat + IndexFlatL2 量化器;设置 nprobe 参数平衡精度/速度 |
| 模型推理速度 | CPU 上 bge-small-zh 处理 1 万段文本约需 30 秒 |
使用 GPU 嵌入(device="cuda");或换用更轻量的 paraphrase-multilingual-MiniLM-L12-v2 |
| 目录覆盖风险 | persist 会覆盖已有目录 |
每次运行前备份或使用时间戳命名目录 |
6️⃣ 使用示例:加载索引并查询
python
from llama_index.core import StorageContext, load_index_from_storage
# 加载已保存的索引
storage_context = StorageContext.from_defaults(persist_dir=INDEX_DIR)
index = load_index_from_storage(storage_context)
# 创建查询引擎
query_engine = index.as_query_engine(similarity_top_k=5)
response = query_engine.query("文档中提到了哪些技术规范?")
print(response)
此外,你也可以直接使用 FAISS 索引进行向量检索(不通过 LlamaIndex):
python
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
embed_model = SentenceTransformer("BAAI/bge-small-zh-v1.5")
faiss_index = faiss.read_index(os.path.join(INDEX_DIR, "vector_store.faiss"))
query = "项目验收标准是什么?"
query_vec = embed_model.encode([query])
distances, indices = faiss_index.search(query_vec, k=5)
代码
build_index.py
python
import os
import pdfplumber
import pytesseract
from pdf2image import convert_from_path
from llama_index.core import (
Document,
VectorStoreIndex,
StorageContext,
Settings
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.huggingface import (
HuggingFaceEmbedding
)
from llama_index.vector_stores.faiss import (
FaissVectorStore
)
import faiss
# ==================================================
# 配置
# ==================================================
PDF_PATH = "./data/y0664.pdf"
INDEX_DIR = "./storage/faiss_index"
EMBED_MODEL = "BAAI/bge-small-zh-v1.5"
EMBED_DIM = 512
CHUNK_SIZE = 512
CHUNK_OVERLAP = 50
os.makedirs(INDEX_DIR, exist_ok=True)
# ==================================================
# OCR
# ==================================================
def ocr_pdf(pdf_path):
print("开始OCR...")
images = convert_from_path(
pdf_path,
dpi=300
)
full_text = ""
for i, img in enumerate(images):
text = pytesseract.image_to_string(
img,
lang="chi_sim+eng"
)
print(
f"第{i+1}页 OCR字符数:",
len(text)
)
full_text += text + "\n"
return full_text
# ==================================================
# PDF读取
# ==================================================
def extract_pdf(pdf_path):
text = ""
with pdfplumber.open(pdf_path) as pdf:
for page in pdf.pages:
page_text = page.extract_text()
if page_text:
text += page_text + "\n"
return text
# ==================================================
# 主程序
# ==================================================
print("加载Embedding模型...")
Settings.embed_model = HuggingFaceEmbedding(
model_name=EMBED_MODEL,
device="cpu"
)
print("读取PDF...")
raw_text = extract_pdf(PDF_PATH)
print("PDF字符数:", len(raw_text))
if len(raw_text) < 5000:
print("文本过少,自动OCR...")
raw_text = ocr_pdf(PDF_PATH)
print("最终字符数:", len(raw_text))
documents = [
Document(
text=raw_text,
metadata={
"source": "YY/T0664-2020"
}
)
]
splitter = SentenceSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP
)
nodes = splitter.get_nodes_from_documents(
documents
)
print("节点数:", len(nodes))
# ==================================================
# 创建FAISS
# ==================================================
faiss_index = faiss.IndexFlatL2(
EMBED_DIM
)
vector_store = FaissVectorStore(
faiss_index=faiss_index
)
storage_context = StorageContext.from_defaults(
vector_store=vector_store
)
print("构建索引...")
index = VectorStoreIndex(
nodes,
storage_context=storage_context
)
print("保存索引...")
storage_context.persist(
persist_dir=INDEX_DIR
)
# 强制写出FAISS文件
faiss.write_index(
faiss_index,
os.path.join(
INDEX_DIR,
"vector_store.faiss"
)
)
print("\n文件检查")
for f in os.listdir(INDEX_DIR):
fp = os.path.join(INDEX_DIR, f)
print(
f,
os.path.getsize(fp)
)
print("\n完成")
print(os.path.abspath(INDEX_DIR))
7️⃣ 总结
本文提供了一个生产可用的 PDF → 向量索引流水线,具备以下亮点:
- ✅ 自动区分文本型 PDF 和扫描件,回退到 OCR。
- ✅ 支持中文嵌入(
bge-small-zh),兼顾效果与速度。 - ✅ 基于 LlamaIndex 和 FAISS,代码简洁且易于扩展。
- ✅ 索引持久化,可重复加载使用。
你可以将此流水线集成到本地 RAG 系统中,作为知识库的预处理环节。下一步可以:
- 增加文件格式支持(如 DOCX、HTML)。
- 将 FAISS 索引迁移到更强大的向量数据库(如 Milvus、Qdrant)。
- 添加文档元数据(作者、时间)以支持过滤检索。
希望本文对你的项目有所帮助。欢迎在评论区分享你的实践经验!