Update asr.py

asr.py

@@ -1,6 +1,7 @@
 import librosa
 import torch
 import numpy as np
+import langid  # Language detection library
 from transformers import Wav2Vec2ForCTC, AutoProcessor
 
 ASR_SAMPLING_RATE = 16_000
@@ -11,6 +12,11 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
 model.eval()
 
+def detect_language(text):
+    """Detects language using langid (fast & lightweight)."""
+    lang, _ = langid.classify(text)
+    return lang if lang in ["en", "sw"] else "en"  # Default to English
+
 def transcribe_auto(audio_data=None):
     if not audio_data:
         return "<<ERROR: Empty Audio Input>>"
@@ -30,20 +36,24 @@ def transcribe_auto(audio_data=None):
 
     inputs = processor(audio_samples, sampling_rate=ASR_SAMPLING_RATE, return_tensors="pt")
 
-    # **Step 1: Detect Language**
+    # **Step 1: Transcribe without Language Detection**
     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]
+        outputs = model(**inputs).logits
+        ids = torch.argmax(outputs, dim=-1)[0]
+        raw_transcription = processor.decode(ids)
+
+    # **Step 2: Detect Language from Transcription**
+    detected_lang = detect_language(raw_transcription)
+    lang_code = "eng" if detected_lang == "en" else "swh"
 
-    # **Step 2: Load Detected Language Adapter**
-    processor.tokenizer.set_target_lang(detected_lang)
-    model.load_adapter(detected_lang)
+    # **Step 3: Reload Model with Correct Adapter**
+    processor.tokenizer.set_target_lang(lang_code)
+    model.load_adapter(lang_code)
 
-    # **Step 3: Transcribe Audio**
+    # **Step 4: Transcribe Again with Correct Adapter**
     with torch.no_grad():
         outputs = model(**inputs).logits
         ids = torch.argmax(outputs, dim=-1)[0]
-        transcription = processor.decode(ids)
+        final_transcription = processor.decode(ids)
 
-    return f"Detected Language: {detected_lang}\n\nTranscription:\n{transcription}"
-    
+    return f"Detected Language: {detected_lang.upper()}\n\nTranscription:\n{final_transcription}"