从 PDF 到 FAISS 向量索引:构建本地 RAG 数据预处理流水线
【学习记录】从 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 库依赖pipinstallpdfplumber pytesseract pdf2image faiss-cpu llama_index transformers huggingface_hub库用途pdfplumber提取文本型 PDF 的文字内容pytesseractOCR 识别图片文字pdf2image将 PDF 页面转换为图像faiss-cpu高效的向量索引库llama_index文档管理、文本分块、向量索引封装transformersHuggingFace 嵌入模型的后端系统依赖必须安装组件安装命令 (Ubuntu/Debian)作用Tesseract OCRsudo apt install tesseract-ocr tesseract-ocr-chi-simOCR 引擎Popplersudo 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.faissLlamaIndex 元数据保存在INDEX_DIR目录下包括文档存储、索引结构等。生成的文件可用于后续的向量检索例如query公司最新政策resultsindex.query(query,top_k5)4️⃣ 核心逻辑解析4.1 PDF 文本提取与 OCR 回退defextract_pdf(path):withpdfplumber.open(path)aspdf:return\n.join(page.extract_text()orforpageinpdf.pages)raw_textextract_pdf(PDF_PATH)iflen(raw_text)5000:# 文本太少判定为扫描件raw_textocr_pdf(PDF_PATH)extract_pdf使用pdfplumber尝试提取文字层。如果总字符数不足 5000可根据文档长度调整阈值则认为该 PDF 是扫描件或图片型触发 OCR。OCR 流程pdf2image.convert_from_path将每页转为 PIL 图片再调用pytesseract.image_to_string识别文字语言包chi_simeng。注意OCR 非常耗时每页约 1~3 秒建议在后台处理或对页数多的文档使用多进程。4.2 文本分块fromllama_index.core.node_parserimportSentenceSplitter splitterSentenceSplitter(chunk_sizeCHUNK_SIZE,chunk_overlapCHUNK_OVERLAP)documents[Document(textraw_text)]nodessplitter.get_nodes_from_documents(documents)SentenceSplitter会尽量在句子边界切分避免切断语义。重叠chunk_overlap可以防止关键信息正好落在两个块的边缘而被遗漏。4.3 向量化与 FAISS 索引构建fromllama_index.embeddings.huggingfaceimportHuggingFaceEmbeddingfromllama_index.coreimportSettings,VectorStoreIndex,StorageContextfromllama_index.vector_stores.faissimportFaissVectorStoreimportfaiss Settings.embed_modelHuggingFaceEmbedding(model_nameBAAI/bge-small-zh-v1.5,devicecpu)faiss_indexfaiss.IndexFlatL2(EMBED_DIM)vector_storeFaissVectorStore(faiss_indexfaiss_index)storage_contextStorageContext.from_defaults(vector_storevector_store)indexVectorStoreIndex(nodes,storage_contextstorage_context)IndexFlatL2使用欧氏距离进行精确检索适合百万级以下的向量。如果数据量巨大100 万可改用IndexIVFFlat加速。4.4 索引持久化storage_context.persist(persist_dirINDEX_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如 PaddleOCRFAISS 索引限制IndexFlatL2精确但速度随数据量线性增长使用IndexIVFFlatIndexFlatL2量化器设置nprobe参数平衡精度/速度模型推理速度CPU 上bge-small-zh处理 1 万段文本约需 30 秒使用 GPU 嵌入devicecuda或换用更轻量的paraphrase-multilingual-MiniLM-L12-v2目录覆盖风险persist会覆盖已有目录每次运行前备份或使用时间戳命名目录6️⃣ 使用示例加载索引并查询fromllama_index.coreimportStorageContext,load_index_from_storage# 加载已保存的索引storage_contextStorageContext.from_defaults(persist_dirINDEX_DIR)indexload_index_from_storage(storage_context)# 创建查询引擎query_engineindex.as_query_engine(similarity_top_k5)responsequery_engine.query(文档中提到了哪些技术规范)print(response)此外你也可以直接使用 FAISS 索引进行向量检索不通过 LlamaIndeximportfaissimportnumpyasnpfromsentence_transformersimportSentenceTransformer embed_modelSentenceTransformer(BAAI/bge-small-zh-v1.5)faiss_indexfaiss.read_index(os.path.join(INDEX_DIR,vector_store.faiss))query项目验收标准是什么query_vecembed_model.encode([query])distances,indicesfaiss_index.search(query_vec,k5)代码build_index.pyimportosimportpdfplumberimportpytesseractfrompdf2imageimportconvert_from_pathfromllama_index.coreimport(Document,VectorStoreIndex,StorageContext,Settings)fromllama_index.core.node_parserimportSentenceSplitterfromllama_index.embeddings.huggingfaceimport(HuggingFaceEmbedding)fromllama_index.vector_stores.faissimport(FaissVectorStore)importfaiss# # 配置# PDF_PATH./data/y0664.pdfINDEX_DIR./storage/faiss_indexEMBED_MODELBAAI/bge-small-zh-v1.5EMBED_DIM512CHUNK_SIZE512CHUNK_OVERLAP50os.makedirs(INDEX_DIR,exist_okTrue)# # OCR# defocr_pdf(pdf_path):print(开始OCR...)imagesconvert_from_path(pdf_path,dpi300)full_textfori,imginenumerate(images):textpytesseract.image_to_string(img,langchi_simeng)print(f第{i1}页 OCR字符数:,len(text))full_texttext\nreturnfull_text# # PDF读取# defextract_pdf(pdf_path):textwithpdfplumber.open(pdf_path)aspdf:forpageinpdf.pages:page_textpage.extract_text()ifpage_text:textpage_text\nreturntext# # 主程序# print(加载Embedding模型...)Settings.embed_modelHuggingFaceEmbedding(model_nameEMBED_MODEL,devicecpu)print(读取PDF...)raw_textextract_pdf(PDF_PATH)print(PDF字符数:,len(raw_text))iflen(raw_text)5000:print(文本过少自动OCR...)raw_textocr_pdf(PDF_PATH)print(最终字符数:,len(raw_text))documents[Document(textraw_text,metadata{source:YY/T0664-2020})]splitterSentenceSplitter(chunk_sizeCHUNK_SIZE,chunk_overlapCHUNK_OVERLAP)nodessplitter.get_nodes_from_documents(documents)print(节点数:,len(nodes))# # 创建FAISS# faiss_indexfaiss.IndexFlatL2(EMBED_DIM)vector_storeFaissVectorStore(faiss_indexfaiss_index)storage_contextStorageContext.from_defaults(vector_storevector_store)print(构建索引...)indexVectorStoreIndex(nodes,storage_contextstorage_context)print(保存索引...)storage_context.persist(persist_dirINDEX_DIR)# 强制写出FAISS文件faiss.write_index(faiss_index,os.path.join(INDEX_DIR,vector_store.faiss))print(\n文件检查)forfinos.listdir(INDEX_DIR):fpos.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。添加文档元数据作者、时间以支持过滤检索。希望本文对你的项目有所帮助。欢迎在评论区分享你的实践经验