granite-speech-5.0-2b

# Adapted from the Granite Speech 5.0 2B model card. Audio must be 16 kHz mono.
# Split recordings longer than about 2 minutes, for example with VAD, before decoding.
import torch, torchaudio
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

model_name = "ibm-granite/granite-speech-5.0-2b"
device = "cuda" if torch.cuda.is_available() else "cpu"
processor = AutoProcessor.from_pretrained(model_name)
tokenizer = processor.tokenizer
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_name, device_map=device, torch_dtype=torch.bfloat16
)

wav, sr = torchaudio.load("segment.wav", normalize=True)
chat = [{"role": "user", "content": "<|audio|>transcribe the speech with proper punctuation and capitalization."}]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)

inputs = processor(prompt, wav, device=device, return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False, num_beams=1)
new_tokens = outputs[0, inputs["input_ids"].shape[-1]:].unsqueeze(0)
print(tokenizer.batch_decode(new_tokens, add_special_tokens=False, skip_special_tokens=True)[0])