Update app.py

app.py

@@ -9,21 +9,15 @@ DESCRIPTION = "IndicConformer ASR with Automatic Language Identification"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # --- Model Loading ---
-# NOTE: If running on a Space with a HF_TOKEN secret,
-# transformers will automatically use it for gated models.
-
-# ASR Model (IndicConformer)
 print("Loading ASR model (IndicConformer)...")
 asr_model_id = "ai4bharat/indic-conformer-600m-multilingual"
 asr_model = AutoModel.from_pretrained(asr_model_id, trust_remote_code=True).to(device)
 asr_model.eval()
 print("✅ ASR Model loaded.")
 
-# Language Identification (LID) Model - Using your specified model
 print("\nLoading Language ID model (MMS-LID-1024)...")
 lid_model_id = "facebook/mms-lid-1024"
 lid_processor = Wav2Vec2FeatureExtractor.from_pretrained(lid_model_id)
-# Load the model with its audio classification head
 lid_model = AutoModelForAudioClassification.from_pretrained(lid_model_id).to(device)
 lid_model.eval()
 print("✅ Language ID Model loaded.")
@@ -36,7 +30,8 @@ LID_TO_ASR_LANG_MAP = {
     "kas_Deva": "ks", "gom_Deva": "kok", "mai_Deva": "mai", "mal_Mlym": "ml",
     "mni_Beng": "mni", "mar_Deva": "mr", "nep_Deva": "ne", "ory_Orya": "or",
     "pan_Guru": "pa", "san_Deva": "sa", "sat_Olck": "sat", "snd_Arab": "sd",
-    "tam_Taml": "ta", "tel_Telu": "te", "urd_Arab": "ur"
+    "tam_Taml": "ta", "tel_Telu": "te", "urd_Arab": "ur",
+    "pan": "pa"  # <-- ADDED THIS FIX FOR PUNJABI
 }
 
 ASR_CODE_TO_NAME = { "as": "Assamese", "bn": "Bengali", "br": "Bodo", "doi": "Dogri", "gu": "Gujarati", "hi": "Hindi", "kn": "Kannada", "ks": "Kashmiri", "kok": "Konkani", "mai": "Maithili", "ml": "Malayalam", "mni": "Manipuri", "mr": "Marathi", "ne": "Nepali", "or": "Odia", "pa": "Punjabi", "sa": "Sanskrit", "sat": "Santali", "sd": "Sindhi", "ta": "Tamil", "te": "Telugu", "ur": "Urdu"}
@@ -54,26 +49,26 @@ def transcribe_audio_with_lid(audio_path):
         return f"Error loading audio: {e}", "", ""
 
     try:
-        # 1. --- Language Identification ---
         inputs = lid_processor(waveform_16k.squeeze(), sampling_rate=16000, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = lid_model(**inputs)
-        
-        # CORRECTED: Access logits as the first element of the output tuple
+
         logits = outputs[0]
         predicted_lid_id = logits.argmax(-1).item()
         detected_lid_code = lid_model.config.id2label[predicted_lid_id]
 
-        # 2. --- Map to ASR Language Code ---
         asr_lang_code = LID_TO_ASR_LANG_MAP.get(detected_lid_code)
         
         if not asr_lang_code:
-            detected_lang_str = f"Detected '{detected_lid_code}', which is not supported by the ASR model."
-            return detected_lang_str, "N/A", "N/A"
+            # Fallback for simple codes like 'pan' from other LID models
+            if detected_lid_code in LID_TO_ASR_LANG_MAP:
+                 asr_lang_code = LID_TO_ASR_LANG_MAP[detected_lid_code]
+            else:
+                detected_lang_str = f"Detected '{detected_lid_code}', which is not supported by the ASR model."
+                return detected_lang_str, "N/A", "N/A"
 
         detected_lang_str = f"Detected Language: {ASR_CODE_TO_NAME.get(asr_lang_code, 'Unknown')}"
 
-        # 3. --- Transcription using the detected language ---
         with torch.no_grad():
             transcription_ctc = asr_model(waveform_16k.to(device), asr_lang_code, "ctc")
             transcription_rnnt = asr_model(waveform_16k.to(device), asr_lang_code, "rnnt")
@@ -83,6 +78,7 @@ def transcribe_audio_with_lid(audio_path):
 
     return detected_lang_str, transcription_ctc.strip(), transcription_rnnt.strip()
 
+
 # --- Gradio UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"## {DESCRIPTION}")