处理30万条记录的CSV时Embedding进程卡顿，求优化方案

核心误解纠正 #

OpenAI的Embedding.create接口完全支持批量输入，一次可传入多个文本（只要总token数不超过模型限制，text-embedding-ada-002单请求最大支持8191token），这是解决30万条数据处理缓慢的关键。

优化方案与代码示例 #

1. 批量调用API（最有效提速手段）

替换逐行apply为批量请求，大幅减少API调用次数：

import os
import pandas as pd
import openai
from dotenv import load_dotenv

load_dotenv('.env')
openai.api_key = os.getenv('OPENAI_API_KEY')

def get_batch_embeddings(texts, model="text-embedding-ada-002"):
    response = openai.Embedding.create(input=texts, model=model)
    return [item['embedding'] for item in response['data']]

# 加载数据
dataset = pd.read_csv('keywords.csv', encoding='ISO-8859-1')

# 分批次处理（根据文本长度调整batch_size，确保总token不超限）
batch_size = 1000
embeddings = []
for i in range(0, len(dataset), batch_size):
    batch_texts = dataset['keyword'].iloc[i:i+batch_size].tolist()
    batch_embeds = get_batch_embeddings(batch_texts)
    embeddings.extend(batch_embeds)
    print(f"已处理 {i+len(batch_embeds)} / {len(dataset)} 条数据")

# 赋值并保存（用parquet格式存储数组类型更高效）
dataset['embedding'] = embeddings
dataset.to_parquet('keywords_with_embeddings.parquet')

# 查询逻辑
keyword = input('Input:')
keyword_vector = get_batch_embeddings([keyword])[0]
print(keyword_vector)

2. 异步并发请求（进一步提升效率）

使用OpenAI异步客户端并发处理多个批次，充分利用带宽：

import asyncio
from openai import AsyncOpenAI

client = AsyncOpenAI(api_key=os.getenv('OPENAI_API_KEY'))

async def get_async_batch_embeddings(texts, model="text-embedding-ada-002"):
    response = await client.embeddings.create(input=texts, model=model)
    return [item.embedding for item in response.data]

async def process_batches_async(dataset, batch_size=1000):
    tasks = []
    for i in range(0, len(dataset), batch_size):
        batch_texts = dataset['keyword'].iloc[i:i+batch_size].tolist()
        tasks.append(get_async_batch_embeddings(batch_texts))
    all_embeddings = await asyncio.gather(*tasks)
    return [embed for batch in all_embeddings for embed in batch]

# 执行异步处理
embeddings = asyncio.run(process_batches_async(dataset))
dataset['embedding'] = embeddings
dataset.to_parquet('keywords_with_embeddings.parquet')

3. 增量处理与缓存（避免重复生成）

将已生成的embedding持久化，后续仅处理新增数据：

try:
    # 加载已有数据
    existing_dataset = pd.read_parquet('keywords_with_embeddings.parquet')
    # 筛选未处理的关键词
    new_dataset = dataset[~dataset['keyword'].isin(existing_dataset['keyword'])]
    print(f"已有 {len(existing_dataset)} 条数据，新增 {len(new_dataset)} 条待处理")
    
    if len(new_dataset) > 0:
        new_embeddings = get_batch_embeddings(new_dataset['keyword'].tolist())
        new_dataset['embedding'] = new_embeddings
        dataset = pd.concat([existing_dataset, new_dataset], ignore_index=True)
    else:
        dataset = existing_dataset
except FileNotFoundError:
    # 首次处理全量数据
    embeddings = get_batch_embeddings(dataset['keyword'].tolist())
    dataset['embedding'] = embeddings

dataset.to_parquet('keywords_with_embeddings.parquet')

4. 处理速率限制（避免限流中断）

加入重试机制，应对OpenAI的API速率限制：

import time

def get_batch_embeddings_with_retry(texts, model="text-embedding-ada-002"):
    while True:
        try:
            response = openai.Embedding.create(input=texts, model=model)
            return [item['embedding'] for item in response['data']]
        except openai.RateLimitError:
            print("触发速率限制，等待10秒后重试...")
            time.sleep(10)

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