chatbot / modules /asr.py
amitshamsundar's picture
issue fix
f5a3837
Raw
History Blame Contribute Delete
1.29 kB
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