from datasets import Audio, load_dataset
from transformers import WhisperForConditionalGeneration, AutoProcessor
import torch
import numpy as np

model = WhisperForConditionalGeneration.from_pretrained(
    "openai/whisper-tiny", torch_dtype=torch.float16
)
processor = AutoProcessor.from_pretrained("openai/whisper-tiny")
model.to("mps")

ds = load_dataset("distil-whisper/meanwhile", "default")["test"]
ds = ds.cast_column("audio", Audio(sampling_rate=16000))
audio = ds[:8]["audio"]
audio = [x["array"] for x in audio]
audio[0][:] = np.random.normal(scale=0.05, size=audio[0].shape)
inputs = processor(
    audio,
    return_tensors="pt",
    truncation=False,
    padding="longest",
    return_attention_mask=True,
    sampling_rate=16_000,
)
inputs = inputs.to(model.device, torch.float16)

result = model.generate(
    **inputs,
    no_speech_threshold=0.2,
    logprob_threshold=0.0,
    temperature=(0.0,),
    task="transcribe",
    language="fr",
)
decoded = processor.batch_decode(
    result, skip_special_tokens=False, decode_with_timestamps=True
)
print(decoded)