from transformers import WhisperProcessor, WhisperForConditionalGeneration
import torch
import torchaudio
import spaces

# Load ASR model and processor
processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2")
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v2")

# Move model to CUDA device
model.to("cuda")
model.eval()

@spaces.GPU(duration_s=30)  # ASR can take longer for complex audio
def transcribe(audio_path):
    waveform, sample_rate = torchaudio.load(audio_path)
    if sample_rate != 16000:
        waveform = torchaudio.functional.resample(waveform, sample_rate, 16000)
    input_features = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features
    input_features = input_features.to("cuda")
    with torch.no_grad():
        # Explicitly set language='en' to ensure translation to English
        # This addresses the breaking change in transformers mentioned in PR #28687
        predicted_ids = model.generate(
            input_features,
            language="en",  # Force English output
            task="translate"  # Ensure translation, not just transcription
        )
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    return transcription