使用Transformers Pipeline处理大音频内存不足，能否将输出写入文件？

当处理大型音频文件时，将整个识别结果存入变量会占用大量内存，导致Google Colab出现内存不足错误。可以通过回调函数让Hugging Face Pipeline每处理完一段音频分片就将结果写入文件，避免内存累积。

实现步骤： #

1. 定义写入文件的回调函数
2. 调用Pipeline时传入回调，实时输出结果

修改后的代码： #

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    chunk_length_s=20,
    batch_size=8,  # 降低batch_size进一步减少内存占用
    return_timestamps=True,
    torch_dtype=torch_dtype,
    device=device,
)

# 定义回调函数：处理每个分片结果并写入文件
def write_chunk_result(chunk_result, file):
    # 提取分片文本
    text = chunk_result["text"]
    # 写入纯文本
    file.write(text + "\n")
    
    # 如果需要保留时间戳，可替换为以下代码
    # for segment in chunk_result.get("chunks", []):
    #     start, end = segment["timestamp"]
    #     file.write(f"[{start:.2f}s - {end:.2f}s] {segment['text']}\n")

# 打开文件并执行识别，实时写入
with open("transcription_output.txt", "w", encoding="utf-8") as output_file:
    def callback(chunk_result):
        write_chunk_result(chunk_result, output_file)
        output_file.flush()  # 强制刷新缓冲区，确保内容实时写入
    
    pipe("filename3.mp3", callback=callback)

额外优化建议： #

• 降低batch_size（从16改为8或4），减少单次处理的内存占用
• 如果不需要逐段时间戳，可将return_timestamps设为False，进一步减少输出数据量

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