Futuresony's picture
Create asr.py
40c9fce verified
raw
history blame
1.68 kB
import librosa
import torch
import numpy as np
from transformers import Wav2Vec2ForCTC, AutoProcessor
ASR_SAMPLING_RATE = 16_000
MODEL_ID = "facebook/mms-1b-all"
# Load MMS Model
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
model.eval()
def transcribe_auto(audio_data=None):
if not audio_data:
return "<<ERROR: Empty Audio Input>>"
# Process Microphone Input
if isinstance(audio_data, tuple):
sr, audio_samples = audio_data
audio_samples = (audio_samples / 32768.0).astype(np.float32)
if sr != ASR_SAMPLING_RATE:
audio_samples = librosa.resample(audio_samples, orig_sr=sr, target_sr=ASR_SAMPLING_RATE)
# Process File Upload Input
else:
if not isinstance(audio_data, str):
return "<<ERROR: Invalid Audio Input>>"
audio_samples = librosa.load(audio_data, sr=ASR_SAMPLING_RATE, mono=True)[0]
inputs = processor(audio_samples, sampling_rate=ASR_SAMPLING_RATE, return_tensors="pt")
# **Step 1: Detect Language**
with torch.no_grad():
lang_id = model.generate(**inputs, task="lang-id")
detected_lang = processor.tokenizer.batch_decode(lang_id, skip_special_tokens=True)[0]
# **Step 2: Load Detected Language Adapter**
processor.tokenizer.set_target_lang(detected_lang)
model.load_adapter(detected_lang)
# **Step 3: Transcribe Audio**
with torch.no_grad():
outputs = model(**inputs).logits
ids = torch.argmax(outputs, dim=-1)[0]
transcription = processor.decode(ids)
return f"Detected Language: {detected_lang}\n\nTranscription:\n{transcription}"