如何基于HuggingFace批量处理超大规模CSV生成Embedding（不加载全量数据）

核心优化方向 #

• 模型仅初始化一次：避免循环内重复加载模型，大幅节省时间与硬件资源
• 分块读取CSV：逐批次加载数据，完全无需将全量文件放入内存
• 批量编码：利用GPU并行计算能力，通过合理设置batch_size最大化推理效率

具体实现方案 #

方案一：原生CSV模块（轻量低内存） #

适合内存紧张的场景，依赖Python原生模块，无额外安装成本：

import csv
from sentence_transformers import SentenceTransformer

# 仅初始化一次模型，指定目标GPU设备
model = SentenceTransformer('flax-sentence-embeddings/some_model_here', device='cuda:0')
batch_size = 512  # 根据GPU显存调整，12G显存可尝试512-1024，显存不足则降低
batch_ids = []
batch_sentences = []

# 逐行读取源CSV，攒够批次就执行编码并写入结果
with open('large_data.csv', 'r', encoding='utf-8') as infile, open('embeddings_result.csv', 'w', encoding='utf-8', newline='') as outfile:
    reader = csv.DictReader(infile)
    writer = csv.writer(outfile)
    writer.writerow(['id', 'embedding'])  # 写入结果文件表头
    
    for row in reader:
        batch_ids.append(row['id'])
        batch_sentences.append(row['sentence'])
        
        # 批次容量达标时处理
        if len(batch_sentences) >= batch_size:
            embeddings = model.encode(batch_sentences, normalize_embeddings=True, batch_size=batch_size)
            # 将编码结果写入文件
            for idx, emb in zip(batch_ids, embeddings):
                writer.writerow([idx, ','.join(map(str, emb))])
            # 清空批次缓存
            batch_ids = []
            batch_sentences = []
    
    # 处理最后一批不足batch_size的数据
    if batch_sentences:
        embeddings = model.encode(batch_sentences, normalize_embeddings=True, batch_size=len(batch_sentences))
        for idx, emb in zip(batch_ids, embeddings):
            writer.writerow([idx, ','.join(map(str, emb))])

方案二：Pandas分块读取（便捷易扩展） #

如果需要对数据做简单预处理（如空值过滤），Pandas的分块读取更高效：

import pandas as pd
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('flax-sentence-embeddings/some_model_here', device='cuda:0')
batch_size = 512
chunk_size = 10000  # 每次读取的行数，可根据内存情况调整

# 分块读取源CSV，逐块处理并追加写入结果
with open('embeddings_result.csv', 'w', encoding='utf-8', newline='') as outfile:
    # 先写入表头
    pd.DataFrame(columns=['id', 'embedding']).to_csv(outfile, index=False, header=True)
    
    for chunk in pd.read_csv('large_data.csv', chunksize=chunk_size):
        # 批量编码当前块的sentence数据
        embeddings = model.encode(chunk['sentence'].tolist(), normalize_embeddings=True, batch_size=batch_size)
        # 组装结果列
        chunk['embedding'] = [','.join(map(str, emb)) for emb in embeddings]
        # 追加写入结果文件（不重复写表头）
        chunk[['id', 'embedding']].to_csv(outfile, index=False, header=False, mode='a')

关键注意事项 #

• batch_size调试：根据GPU显存大小测试最优值，出现显存溢出时及时减小
• 设备验证：确保模型加载到GPU上，多GPU环境可指定对应编号（如cuda:1）
• 内存控制：每批次处理完成后立即写入磁盘，不要缓存所有编码结果到内存
• 进度可视化：可在model.encode中添加show_progress_bar=True查看实时进度

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