LangChain构建FAISS向量库内存溢出，求增量构建解决方案

核心思路 #

• 避免一次性加载所有文档到内存，改为分批加载pickle分片，每处理一批就将文档向量添加到FAISS索引
• 利用LangChain的FAISS类提供的add_documents方法实现增量更新
• 每处理完一批可临时保存索引，防止意外中断丢失进度

修改后的代码实现 #

import os
import pickle
from tqdm import tqdm
from dotenv import load_dotenv
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS

load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_KEY")

def build_faiss_incrementally(batch_size=100):
    # 获取排序后的分片文件列表，保持原处理顺序
    chunks_files = sorted(
        [f for f in os.listdir("chunks") if f.endswith(".pkl")],
        key=lambda x: int(x.split(".")[0])
    )
    total_chunks = len(chunks_files)
    embeddings = OpenAIEmbeddings()
    db = None

    # 按批次处理分片
    for i in tqdm(range(0, total_chunks, batch_size), desc="Processing Batches"):
        batch_files = chunks_files[i:i+batch_size]
        batch_docs = []
        
        # 加载当前批次的所有文档
        for chunk_file in batch_files:
            with open(f'chunks/{chunk_file}', 'rb') as file:
                batch_docs.extend(pickle.load(file))
        
        # 初始化或增量更新FAISS索引
        if db is None:
            # 首次处理，创建初始索引
            db = FAISS.from_documents(batch_docs, embeddings)
        else:
            # 后续批次，增量添加文档到已有索引
            db.add_documents(batch_docs)
        
        # 可选：每处理一批保存临时索引，防止中断丢失进度
        temp_save_path = f"faiss_openai_embeddings_temp_{i//batch_size}"
        db.save_local(temp_save_path)
        # 主动清理当前批次文档，释放内存
        del batch_docs
    
    # 保存最终完整索引
    final_save_path = "faiss_openai_embeddings"
    db.save_local(final_save_path)
    print(f"-- 最终索引已保存到 {final_save_path} --")

if __name__ == '__main__':
    build_faiss_incrementally(batch_size=100)

关键优化点说明 #

• 分批加载控制：通过batch_size参数调整每次加载的分片数量，可根据自身内存情况灵活增减（内存紧张就调小，充足则调大）
• 增量索引构建：用add_documents替代一次性from_documents，避免一次性生成所有向量占用大量内存
• 内存主动释放：每批处理完成后删除当前批次的文档列表，强制释放内存空间
• 临时进度保存：可选的临时索引保存机制，防止程序意外中断导致前序工作全部丢失
• 顺序一致性：保持分片文件的排序逻辑与原代码一致，确保索引构建顺序不变

注意事项 #

• OpenAI Embeddings接口存在速率限制，若遇到API请求报错，可添加重试逻辑或调整请求间隔
• 临时保存的索引文件可在最终构建完成后删除，节省存储空间
• 若后续需要继续添加新分片，可加载已保存的索引，重复执行add_documents逻辑

内容的提问来源于stack exchange，提问作者Asim