如何使用OpenAI API批量生成Embeddings并提升处理速度？

并行发送请求确实能显著提升处理速度，但需要结合合理的批量划分和API速率限制来落地，以下是具体思路和实现方式：

核心逻辑 #

逐个请求完全串行等待，单次大列表请求如果包含数千条句子，单请求处理耗时会被拉长；而将句子划分为符合token限制的批量，并行发送这些批量请求，可以同时利用OpenAI API的处理能力，大幅压缩总耗时。

关键注意事项 #

• 遵守速率限制：OpenAI API对请求频率（每分钟请求数）和token消耗（每分钟token数）有明确限制，并行请求的总请求数和token数不能超过你的API配额，否则会触发429限流错误。
• 合理划分批量：text-embedding-ada-002模型单请求支持最多8191个token，需要将你的数千条句子拆分为多个批量，每个批量的总token数不超过上限（可通过tiktoken库计算单句token数来划分）。

代码示例（线程池并行） #

import requests
from concurrent.futures import ThreadPoolExecutor
import tiktoken

# 初始化token计算器
tokenizer = tiktoken.get_encoding("cl100k_base")
API_KEY = "your-api-key"
MODEL = "text-embedding-ada-002"
MAX_TOKENS_PER_REQUEST = 8191

def get_embeddings(batch):
    try:
        response = requests.post(
            "https://api.openai.com/v1/embeddings",
            json={
                "model": MODEL,
                "input": batch
            },
            headers={
                "Authorization": f"Bearer {API_KEY}"
            }
        )
        response.raise_for_status()
        return response.json()["data"]
    except Exception as e:
        print(f"请求失败: {e}")
        return None

def split_into_batches(sentences):
    batches = []
    current_batch = []
    current_token_count = 0

    for sentence in sentences:
        token_count = len(tokenizer.encode(sentence))
        # 加上当前句子后超过上限则新建批次
        if current_token_count + token_count > MAX_TOKENS_PER_REQUEST:
            batches.append(current_batch)
            current_batch = [sentence]
            current_token_count = token_count
        else:
            current_batch.append(sentence)
            current_token_count += token_count
    # 添加最后一个批次
    if current_batch:
        batches.append(current_batch)
    return batches

# 替换为你的实际数千条句子列表
sentences = ["text:This is a test", "text:This is another test", ...]
batches = split_into_batches(sentences)

# 调整max_workers控制并行数（不要超过API速率限制）
with ThreadPoolExecutor(max_workers=5) as executor:
    results = list(executor.map(get_embeddings, batches))

# 整理所有embeddings结果
all_embeddings = []
for batch_result in results:
    if batch_result:
        all_embeddings.extend([item["embedding"] for item in batch_result])

额外优化建议 #

• 加入重试机制：对429限流错误进行指数退避重试，避免请求失败。
• 异步请求：使用aiohttp库实现异步HTTP请求，相比线程池可能有更高的效率（适合超大规模请求场景）。

内容的提问来源于stack exchange，提问作者Constantly Groovin'