gte-base-zh批量处理技巧高效完成海量文档向量化如果你手头有堆积如山的文档——可能是公司多年的报告、海量的用户评论或是整个图书馆的电子书——需要把它们全部转换成计算机能理解的向量你可能会觉得这是个令人望而生畏的工程。手动一个个处理那得处理到猴年马月。直接跑个脚本又怕内存爆炸或者程序中途崩溃一夜回到解放前。别担心这正是我们今天要解决的问题。gte-base-zh 是一个优秀的中文文本向量化模型但面对TB级别的数据简单的“读取-处理-写入”循环就不够看了。这篇文章我就以一个处理过数亿条文本的老兵身份跟你聊聊怎么搭建一套健壮、高效的批量向量化流水线。咱们不谈空洞的理论直接上能跑起来的代码和踩过坑的实战经验。1. 理解挑战海量文档向量化不是简单的循环在开始写代码之前我们得先搞清楚处理海量文档和处理几百个文档本质区别在哪里。这决定了我们的技术方案。首先数据量巨大意味着数据无法一次性装入内存。你不能简单地把所有文本读进一个列表里然后交给模型。其次处理耗时极长。即使模型推理很快乘以百万、千万的量级后时间成本也非常可观系统稳定性如网络波动、GPU内存溢出成为巨大风险。最后故障恢复成本高。程序运行到一半崩溃如果无法从中断点继续损失是难以接受的。因此我们的核心思路从“快速处理”转变为“稳定、可恢复、高效地处理”。这套流程通常包含几个关键环节稳定可靠的数据读取、充分利用硬件的并行计算、优雅的错误处理与日志记录以及高效持久化的结果存储。接下来我们就逐一拆解。2. 构建稳健的数据读取管道数据读取是流水线的源头必须稳固。我们的原则是流式读取、分批供给、记录进度。2.1 从多种来源流式读取你的文档可能存放在各种地方本地的一个个大文本文件、CSV、JSON文件或者远程数据库、对象存储中。关键是使用**生成器generator**来流式地yield数据避免内存爆掉。假设我们的文档是每行一个JSON对象的文本文件例如每行是{“id”: “doc1”, “text”: “...”}下面是一个基础的读取器import json import logging from pathlib import Path logging.basicConfig(levellogging.INFO, format‘%(asctime)s - %(levelname)s - %(message)s’) logger logging.getLogger(__name__) def document_reader(file_path, batch_size100, start_from0): 流式读取文档的生成器。 Args: file_path: 文档文件路径每行一个JSON。 batch_size: 每次yield的文档数量。 start_from: 从第几行开始读取用于断点续传。 Yields: 包含多个文档的列表。 count 0 current_batch [] with open(file_path, ‘r’, encoding‘utf-8’) as f: for line_num, line in enumerate(f): # 断点续传跳过已经处理过的行 if line_num start_from: continue try: data json.loads(line.strip()) # 假设JSON结构中有‘id’和‘text’字段 doc_id data.get(‘id’, f‘line_{line_num}’) text data.get(‘text’, ‘’).strip() if text: # 只处理非空文本 current_batch.append({‘id’: doc_id, ‘text’: text}) count 1 if len(current_batch) batch_size: yield current_batch current_batch [] logger.info(f“已读取 {count} 个文档”) except json.JSONDecodeError as e: logger.warning(f“跳过第 {line_num} 行JSON解析错误: {e}”) continue # 返回最后一批不足batch_size的数据 if current_batch: yield current_batch logger.info(f“读取完成总计 {count} 个文档”)这个document_reader函数是一个生成器它一次只读一行攒够一个批次比如100条才yield出去内存占用始终很小。start_from参数是关键它记录了上次处理到了哪一行程序重启后可以从这里继续实现断点续传。2.2 记录处理进度为了实现断点续传我们需要把处理进度持久化下来。一个简单有效的方法是使用一个独立的进度文件。import pickle import os PROGRESS_FILE “processing_progress.pkl” def save_progress(last_processed_line): 保存最后处理到的行号 with open(PROGRESS_FILE, ‘wb’) as f: pickle.dump({‘last_line’: last_processed_line}, f) def load_progress(): 加载处理进度如果文件不存在则从0开始 if os.path.exists(PROGRESS_FILE): with open(PROGRESS_FILE, ‘rb’) as f: progress pickle.load(f) return progress.get(‘last_line’, 0) return 0在主流程中每成功处理并存储完一个批次就调用save_progress更新行号。这样即使程序崩溃重启后load_progress也能告诉我们从哪里开始。3. 利用并行计算加速推理gte-base-zh模型在GPU上推理很快但单个GPU处理海量数据仍是瓶颈。我们需要并行化。3.1 单机多进程/多GPU并行transformers库的pipeline配合Python的multiprocessing可以轻松实现批处理但更高效的方式是结合torch的数据并行。不过对于向量化这种计算密集、数据独立的任务使用**多进程池Process Pool**处理数据批次是更直观且能充分利用多核CPU进行数据预加载的方法。下面是一个结合多进程进行批量推理的示例。我们假设已经有一个处理单个批次的函数process_batch。from transformers import AutoTokenizer, AutoModel import torch import torch.nn.functional as F from multiprocessing import Pool, cpu_count import numpy as np # 加载模型和分词器全局变量在每个进程中只加载一次 MODEL_NAME “thenlper/gte-base-zh” tokenizer None model None device None def init_worker(gpu_idNone): 初始化工作进程加载模型到指定设备 global tokenizer, model, device if tokenizer is None: tokenizer AutoTokenizer.from_pretrained(MODEL_NAME) if model is None: model AutoModel.from_pretrained(MODEL_NAME) if gpu_id is not None and torch.cuda.is_available(): device torch.device(f‘cuda:{gpu_id}’) model.to(device) else: device torch.device(‘cpu’) model.eval() def embed_texts(texts): 计算一个文本列表的向量在已初始化的进程中运行 global tokenizer, model, device inputs tokenizer(texts, paddingTrue, truncationTrue, max_length512, return_tensors“pt”) inputs {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs model(**inputs) # 使用mean pooling获取句子向量 embeddings mean_pooling(outputs, inputs[‘attention_mask’]) embeddings F.normalize(embeddings, p2, dim1) # 可选归一化 return embeddings.cpu().numpy() def mean_pooling(model_output, attention_mask): Mean Pooling - 取token embeddings的平均值并考虑attention mask token_embeddings model_output[0] input_mask_expanded attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min1e-9) def process_batch(batch): 处理一个批次的文档提取文本并向量化 texts [doc[‘text’] for doc in batch] ids [doc[‘id’] for doc in batch] try: vectors embed_texts(texts) # 形状: [batch_size, embedding_dim] return {‘ids’: ids, ‘vectors’: vectors, ‘status’: ‘success’} except Exception as e: logger.error(f“处理批次失败: {e}”) return {‘ids’: ids, ‘vectors’: None, ‘status’: ‘failed’, ‘error’: str(e)}但请注意在多进程环境中每个进程都需要加载一份模型非常耗内存。更推荐的方法是使用单个进程负责推理但利用多线程/异步进行数据加载和结果写入或者使用专门的分布式任务队列如Celery。不过对于单机多GPU我们可以采用更直接的方式将数据批次分配给不同的GPU进程。下面是一个简化的主流程框架它使用进程池并为每个进程分配一个GPU ID假设有多个GPUdef main_processing_loop(file_path, num_gpustorch.cuda.device_count()): 主处理循环使用多进程每个进程绑定一个GPU start_line load_progress() reader document_reader(file_path, batch_size64, start_fromstart_line) # 创建进程池进程数等于GPU数或CPU核心数取较小值 num_workers min(num_gpus, cpu_count()) if num_gpus 0 else cpu_count() # 为每个工作进程分配GPU ID如果没有GPU则为None gpu_ids list(range(num_workers)) if num_gpus 0 else [None] * num_workers with Pool(processesnum_workers, initializerinit_worker, initargs(gpu_ids[0],)) as pool: # 注意上面的initargs只能传一个固定值更复杂的分配需要其他机制。 # 一种更实际的做法是在进程内部根据进程索引决定使用哪个GPU。 # 这里为了简化假设每个进程都初始化到不同的GPU这需要更复杂的进程初始化逻辑。 # 更常见的模式是使用一个单独的推理服务进程或者使用torch的DataParallel。 # 鉴于篇幅我们转向另一种更易于实现的并行模式使用线程池进行数据准备和IO模型推理部分保持单进程批处理。3.2 更实用的并行架构生产者-消费者模式对于大多数实际场景我推荐使用生产者-消费者模式配合一个大型批处理推理。即一个或多个线程生产者负责快速从数据源读取数据并放入队列一个主线程消费者从队列中取出大批数据进行GPU批量推理batch_size可以设大些如256或512然后将结果交给另一个线程写入者进行持久化。import queue import threading import time class VectorizationPipeline: def __init__(self, model, tokenizer, device, batch_size256, max_queue_size20): self.model model self.tokenizer tokenizer self.device device self.batch_size batch_size self.input_queue queue.Queue(maxsizemax_queue_size) self.output_queue queue.Queue(maxsizemax_queue_size) self.stop_event threading.Event() def producer(self, file_path, start_from): 生产者读取文档并放入输入队列 reader document_reader(file_path, batch_size100, start_fromstart_from) # 读取批次可以小点 for batch in reader: self.input_queue.put(batch) # 可以在这里记录读取到的行号用于更细粒度的进度保存 self.input_queue.put(None) # 发送结束信号 logger.info(“生产者任务完成”) def consumer(self): 消费者从队列取数据批量推理结果放入输出队列 while not self.stop_event.is_set(): try: batches [] # 积累数据直到达到推理批次大小或超时 while len(batches) self.batch_size // 100: # 假设读取批次是100 try: batch self.input_queue.get(timeout1) if batch is None: # 遇到结束信号 self.input_queue.put(None) # 传递给其他可能的生产者/消费者 if batches: self._process_and_put(batches) self.output_queue.put(None) # 通知写入者结束 logger.info(“消费者任务完成”) return batches.append(batch) except queue.Empty: if batches: # 有数据就处理没数据就继续等 break continue if batches: self._process_and_put(batches) except Exception as e: logger.error(f“消费者发生错误: {e}”, exc_infoTrue) def _process_and_put(self, batches): 处理积累的批次并放入输出队列 # 合并批次 merged_batch [] for b in batches: merged_batch.extend(b) texts [doc[‘text’] for doc in merged_batch] ids [doc[‘id’] for doc in merged_batch] try: vectors embed_texts(texts) # 使用之前定义的embed_texts这里需要能访问到model等 self.output_queue.put({‘ids’: ids, ‘vectors’: vectors}) except Exception as e: logger.error(f“推理失败: {e}”) # 可以选择将失败批次重新放回队列或记录到失败列表 def writer(self, output_connector): 写入者从输出队列取结果并持久化 while not self.stop_event.is_set(): try: result self.output_queue.get(timeout5) if result is None: logger.info(“写入者任务完成”) break output_connector.save_batch(result[‘ids’], result[‘vectors’]) # 在这里调用 save_progress 更新进度需要知道处理到的行号 logger.info(f“已写入一批 {len(result[‘ids’])} 个向量”) except queue.Empty: continue except Exception as e: logger.error(f“写入结果失败: {e}”, exc_infoTrue) def run(self, file_path): start_from load_progress() # 创建并启动线程 producer_thread threading.Thread(targetself.producer, args(file_path, start_from)) consumer_thread threading.Thread(targetself.consumer) writer_thread threading.Thread(targetself.writer, args(your_output_connector,)) # 需要实现 producer_thread.start() consumer_thread.start() writer_thread.start() producer_thread.join() consumer_thread.join() writer_thread.join()这个架构将IO读取、写入与计算GPU推理分离避免了相互等待能显著提升吞吐量。你可以根据硬件情况调整队列大小和批次大小。4. 错误处理、日志与监控海量处理中错误是常态而非例外。健壮的程序必须能处理错误并继续运行。错误重试对于网络超时、临时性错误可以加入重试机制。隔离失败将处理失败的文档ID和错误原因记录到单独的“失败清单”文件中方便后续排查和重试而不是让整个程序停止。详细日志使用Python的logging模块记录不同级别INFO, WARNING, ERROR的日志并输出到文件便于事后分析。简易监控可以定期打印处理速度如 docs/sec估算剩余时间让运行过程更透明。import time class RetryDecorator: 简单的重试装饰器 def __init__(self, max_retries3, delay1): self.max_retries max_retries self.delay delay def __call__(self, func): def wrapper(*args, **kwargs): last_exception None for attempt in range(self.max_retries): try: return func(*args, **kwargs) except Exception as e: last_exception e logger.warning(f“{func.__name__} 第{attempt1}次尝试失败: {e}”) if attempt self.max_retries - 1: time.sleep(self.delay * (2 ** attempt)) # 指数退避 logger.error(f“{func.__name__} 重试{self.max_retries}次后仍失败”) raise last_exception return wrapper # 使用示例 RetryDecorator(max_retries3, delay2) def call_model_api(texts): # 模拟可能失败的调用 return embed_texts(texts)5. 结果的高效持久化生成的向量需要存起来。对于海量数据直接写单个大文件或频繁插入数据库效率很低。向量数据库这是最专业的方案如Milvus、Qdrant、Weaviate或Pinecone云服务。它们专为向量检索设计支持批量导入、自动创建索引。通常提供Python客户端使用起来类似这样from qdrant_client import QdrantClient from qdrant_client.models import PointStruct, VectorParams, Distance client QdrantClient(host“localhost”, port6333) client.create_collection( collection_name“my_docs”, vectors_configVectorParams(size768, distanceDistance.COSINE) # gte-base-zh向量维度是768 ) # 批量插入点 points [ PointStruct( iddoc_id, # 使用你的文档ID vectorvector.tolist(), # 向量 payload{“text”: text} # 可选的元数据 ) for doc_id, vector, text in zip(batch_ids, batch_vectors, batch_texts) ] client.upsert(collection_name“my_docs”, pointspoints)分布式文件系统索引文件如果暂时不想引入向量数据库可以将向量以NumPy数组或Parquet格式按分片shard保存到对象存储如S3、MinIO或HDFS同时维护一个元数据文件如JSONL或SQLite记录文档ID与向量存储路径/偏移量的映射关系。这种方法成本低但检索时需要自己加载索引。import numpy as np import pyarrow.parquet as pq import pyarrow as pa def save_vectors_to_parquet(ids, vectors, shard_path): 将一批向量保存为Parquet文件 table pa.table({ ‘id’: ids, ‘vector’: vectors.tolist() # Parquet支持列表列 }) pq.write_table(table, shard_path) logger.info(f“向量已保存至 {shard_path}”)6. 总结处理TB级文档的向量化核心思想是将它看作一个数据流水线工程而不仅仅是调用几次模型API。我们需要关注稳定性断点续传、错误处理、效率并行计算、异步IO和可维护性清晰日志、模块化设计。回顾一下关键点首先用生成器流式读取数据并记录进度以实现断点续传。其次通过生产者-消费者等多线程模式将数据读取、模型推理和结果写入解耦最大化硬件利用率。再者为所有可能失败的操作加上重试和降级处理并记录详尽的日志。最后根据你的应用场景选择向量数据库或分片文件系统来存储结果。这套方案听起来有点复杂但一旦搭建起来你就可以放心地把海量数据交给它去喝杯咖啡甚至睡一觉回来就能拿到全部向量的结果。它带来的稳定性和效率提升对于长期、大规模的数据处理任务来说是绝对值得的。你可以先从简单的单进程版本开始逐步引入并行和容错机制最终构建出适合自己业务场景的强力工具。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。