from transformers import AutoModel, AutoTokenizer
import torch

REPO_ID = "bharatgenai/Shrutam-2"

model = AutoModel.from_pretrained(REPO_ID, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model.to("cuda" if torch.cuda.is_available() else "cpu")
model.eval()

prompt = "Transcribe speech to Hindi text."

# Single file 
print(model.transcribe("audio.wav", prompts=prompt, tokenizer=tokenizer))

# Batch inference — one prompt per file
wavs = ["clip_hi.wav", "clip_mr.wav", "clip_ta.wav"]
prompts = [
    "Transcribe speech to Hindi text.",
    "Transcribe speech to Marathi text.",
    "Transcribe speech to Tamil text.",
]
print(model.transcribe(wavs, prompts=prompts, batch_size=2, tokenizer=tokenizer))