

Let's see how to use a user-managed pool for batch decoding multiple audios
from multiprocessing import get_context
from transformers import AutoTokenizer, AutoProcessor, AutoModelForCTC
from datasets import load_dataset
import datasets
import torch
import model, feature extractor, tokenizer
model = AutoModelForCTC.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm").to("cuda")
processor = AutoProcessor.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm")
load example dataset
dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
dataset = dataset.cast_column("audio", datasets.Audio(sampling_rate=16_000))
def map_to_array(batch):
     batch["speech"] = batch["audio"]["array"]
     return batch
prepare speech data for batch inference
dataset = dataset.map(map_to_array, remove_columns=["audio"])
def map_to_pred(batch, pool):
     inputs = processor(batch["speech"], sampling_rate=16_000, padding=True, return_tensors="pt")
     inputs = {k: v.to("cuda") for k, v in inputs.items()}