如何使用变压器的管道处理大文件

Question

我使用管道从 Hugging Face 加载模型：

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=16,
    return_timestamps=True,
    torch_dtype=torch_dtype,
    device=device,
)

我想处理大音频文件

result = pipe("filename3.mp3")

但是谷歌协作说我的内存不足。是否可以将输出通过管道传输到文件而不是变量？

Answer 1

没有。

Whisper

large

模型的大小为几 Gb - 通常大于单个 GPU RAM。如果一个 GPU 上的 GPU RAM 不足，您可能需要使用

torch.nn

的函数

nn.DataParallel

将模型分散到多个 GPU（如果可用）。

或者，您可能需要使用较小的

medium

型号。