LangChain嵌入大文本块至Pinecone遇OpenAI限流，求逐块加载方案

使用LangChain读取PDF文件并转换为文本块，将数据嵌入为向量后加载至Pinecone向量库时出现MaxRetryError，推测是一次性加载所有文本块导致请求过载，询问是否存在类似add_document的函数用于逐块/分批加载数据。

PDF加载与文本分块（运行正常） #

def load_document(file):
    from langchain.document_loaders import PyPDFLoader
    print(f'Loading {file} ..')
    loader = PyPDFLoader(file)
    # 返回LangChain文档列表，每页对应一个文档
    data = loader.load() 
    return data


data=load_document("DATA/capacitance.pdf")
# 打印第二页内容
print(data[1].page_content)
print(data[2].metadata)

# 文本分块
def chunk_data(data,chunk_size=256):
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    text_splitter=RecursiveCharacterTextSplitter(chunk_size=chunk_size,chunk_overlap=0)
    chunks=text_splitter.split_documents(data)
    print(type(chunks))
    return chunks

chunks=chunk_data(data)
print(len(chunks))

Pinecone嵌入相关代码 #

已提前创建Pinecone索引electrostatics：

pinecone.create_index('electrostatics',dimension=1536,metric='cosine')

import os
from dotenv import load_dotenv,find_dotenv
load_dotenv("D:/test/.env")
print(os.environ.get("OPENAI_API_KEY"))

def insert_embeddings(index_name,chunks):
    import pinecone
    from langchain.vectorstores import Pinecone
    from langchain.embeddings.openai import OpenAIEmbeddings

    embeddings=OpenAIEmbeddings()
    pinecone.init(api_key=os.environ.get("PINECONE_API_KEY"),environment=os.environ.get("PINECONE_ENV"))
    vector_store=Pinecone.from_documents(chunks,embeddings,index_name=index_name)
    print("Ok")

调用方式：

index_name='electrostatics'
vector_store=insert_embeddings(index_name,chunks)

FAISS嵌入代码 #

from langchain.vectorstores import FAISS
from langchain.embeddings.openai import OpenAIEmbeddings
embeddings=OpenAIEmbeddings()
db = FAISS.from_documents(chunks, embeddings)

LangChain的Pinecone类提供了add_documents方法，可以实现逐块或分批加载文本块，避免一次性请求过多导致的重试错误。

修改后的嵌入函数示例 #

def insert_embeddings_batch(index_name, chunks, batch_size=10, delay=1):
    import pinecone
    import time
    from langchain.vectorstores import Pinecone
    from langchain.embeddings.openai import OpenAIEmbeddings

    embeddings = OpenAIEmbeddings()
    pinecone.init(api_key=os.environ.get("PINECONE_API_KEY"), environment=os.environ.get("PINECONE_ENV"))
    
    # 获取已创建的索引实例
    index = pinecone.Index(index_name)
    # 初始化Pinecone向量存储
    vector_store = Pinecone(index, embeddings.embed_query, "text")
    
    # 分批添加文档
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i+batch_size]
        vector_store.add_documents(batch)
        print(f"已添加第 {i//batch_size + 1} 批，共 {(len(chunks)-1)//batch_size + 1} 批")
        # 每批添加后延迟，避免请求过于频繁
        time.sleep(delay)
    print("所有文档添加完成")
    return vector_store

调用方式 #

index_name='electrostatics'
vector_store=insert_embeddings_batch(index_name, chunks, batch_size=10)

说明 #

• batch_size可根据文本块数量和请求限制调整，数量多可适当调大
• 添加delay是为了避免短时间内请求触发限流，可根据实际情况调整或移除
• 若需逐块加载，将batch_size设为1即可

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