Update ai4bharat/IndicLID.py

ai4bharat/IndicLID.py

@@ -43,16 +43,13 @@ class IndicLID():
         self.IndicLID_FTR_path = "models/indiclid-ftr.bin"  
         self.IndicLID_BERT_path = "models/indiclid-bert"
         
-        # Language mappings
+        # Updated language labels - Complete list as per your specification
         self.IndicLID_labels = [
-            'asm_Beng', 'asm_Latn', 'ben_Beng', 'ben_Latn', 'brx_Deva', 'brx_Latn',
-            'doi_Deva', 'doi_Latn', 'eng_Latn', 'guj_Gujr', 'guj_Latn', 'hin_Deva',
-            'hin_Latn', 'kan_Knda', 'kan_Latn', 'kas_Arab', 'kas_Deva', 'kas_Latn',
-            'kok_Deva', 'kok_Latn', 'mai_Deva', 'mai_Latn', 'mal_Mlym', 'mal_Latn',
-            'mni_Beng', 'mni_Meti', 'mni_Latn', 'mar_Deva', 'mar_Latn', 'nep_Deva',
-            'nep_Latn', 'ori_Orya', 'ori_Latn', 'pan_Guru', 'pan_Latn', 'san_Deva',
-            'san_Latn', 'sat_Olch', 'snd_Arab', 'snd_Latn', 'tam_Taml', 'tam_Latn',
-            'tel_Telu', 'tel_Latn', 'urd_Arab', 'urd_Latn', 'other'
+            'asm_Beng', 'ben_Beng', 'brx_Deva', 'doi_Deva', 'gom_Deva', 'guj_Gujr', 
+            'hin_Deva', 'kan_Knda', 'kas_Arab', 'kas_Deva', 'mai_Deva', 'mal_Mlym', 
+            'mni_Beng', 'mni_Mtei', 'mar_Deva', 'npi_Deva', 'ory_Orya', 'pan_Guru', 
+            'san_Deva', 'sat_Olck', 'snd_Arab', 'snd_Deva', 'tam_Taml', 'tel_Telu', 
+            'urd_Arab', 'eng_Latn', 'other'
         ]
         
         # Load models
@@ -107,7 +104,7 @@ class IndicLID():
                 output_dict[idx] = (text, label, score, 'IndicLID-FTN')
             else:
                 # Fallback - simple heuristic based on script
-                detected_lang = self.detect_script_simple(text)
+                detected_lang = self.detect_script_enhanced(text)
                 output_dict[idx] = (text, detected_lang, 0.8, 'Script-based')
         
         return output_dict
@@ -170,12 +167,16 @@ class IndicLID():
             print(f"BERT inference error: {e}")
             return 'eng_Latn', 0.5
     
-    def detect_script_simple(self, text):
-        """Simple script detection based on Unicode ranges"""
-        # Check for common Indian scripts
+    def detect_script_enhanced(self, text):
+        """Enhanced script detection based on Unicode ranges for all supported languages"""
+        # Check for various Indian scripts
         if any(ord(char) >= 0x0900 and ord(char) <= 0x097F for char in text):  # Devanagari
+            # Could be Hindi, Marathi, Sanskrit, Nepali, Dogri, Maithili, Sindhi
+            # Simple heuristic - default to Hindi for Devanagari
             return 'hin_Deva'
         elif any(ord(char) >= 0x0980 and ord(char) <= 0x09FF for char in text):  # Bengali
+            # Could be Bengali, Assamese, Manipuri
+            # Default to Bengali
             return 'ben_Beng'
         elif any(ord(char) >= 0x0B80 and ord(char) <= 0x0BFF for char in text):  # Tamil
             return 'tam_Taml'
@@ -189,6 +190,15 @@ class IndicLID():
             return 'guj_Gujr'
         elif any(ord(char) >= 0x0A00 and ord(char) <= 0x0A7F for char in text):  # Gurmukhi (Punjabi)
             return 'pan_Guru'
+        elif any(ord(char) >= 0x0B00 and ord(char) <= 0x0B7F for char in text):  # Odia
+            return 'ory_Orya'
+        elif any(ord(char) >= 0x0600 and ord(char) <= 0x06FF for char in text):  # Arabic script
+            # Could be Urdu, Kashmiri, Sindhi
+            return 'urd_Arab'
+        elif any(ord(char) >= 0x1C00 and ord(char) <= 0x1C4F for char in text):  # Ol Chiki (Santali)
+            return 'sat_Olck'
+        elif any(ord(char) >= 0xAAE0 and ord(char) <= 0xAAFF for char in text):  # Meetei Mayek (Manipuri)
+            return 'mni_Mtei'
         else:
             return 'eng_Latn'  # Default to English
     
@@ -232,7 +242,8 @@ class IndicLID():
                     'text_lang': lang_name,
                     'text_lang_score': confidence,
                     'script': 'native' if not self.is_roman_text(sentence) else 'roman',
-                    'model_used': model_used
+                    'model_used': model_used,
+                    'lang_code': lang_code
                 }
             }
             results.append(result_dict)
@@ -240,25 +251,35 @@ class IndicLID():
         return results
     
     def code_to_language(self, lang_code):
-        """Convert language code to language name"""
+        """Convert language code to language name - Updated with all 26 languages"""
         code_to_lang = {
-            'hin_Deva': 'hindi', 'hin_Latn': 'hindi',
-            'ben_Beng': 'bengali', 'ben_Latn': 'bengali', 
-            'tam_Taml': 'tamil', 'tam_Latn': 'tamil',
-            'tel_Telu': 'telugu', 'tel_Latn': 'telugu',
-            'kan_Knda': 'kannada', 'kan_Latn': 'kannada',
-            'mal_Mlym': 'malayalam', 'mal_Latn': 'malayalam',
-            'guj_Gujr': 'gujarati', 'guj_Latn': 'gujarati',
-            'pan_Guru': 'punjabi', 'pan_Latn': 'punjabi',
-            'mar_Deva': 'marathi', 'mar_Latn': 'marathi',
-            'urd_Arab': 'urdu', 'urd_Latn': 'urdu',
-            'ori_Orya': 'odia', 'ori_Latn': 'odia',
-            'asm_Beng': 'assamese', 'asm_Latn': 'assamese',
-            'nep_Deva': 'nepali', 'nep_Latn': 'nepali',
-            'kas_Arab': 'kashmiri', 'kas_Deva': 'kashmiri', 'kas_Latn': 'kashmiri',
-            'snd_Arab': 'sindhi', 'snd_Latn': 'sindhi',
-            'san_Deva': 'sanskrit', 'san_Latn': 'sanskrit',
+            # Complete mapping for all supported languages
+            'asm_Beng': 'assamese',
+            'ben_Beng': 'bengali',
+            'brx_Deva': 'bodo',
+            'doi_Deva': 'dogri',
+            'gom_Deva': 'konkani',  # Goan Konkani
+            'guj_Gujr': 'gujarati',
+            'hin_Deva': 'hindi',
+            'kan_Knda': 'kannada',
+            'kas_Arab': 'kashmiri',
+            'kas_Deva': 'kashmiri',
+            'mai_Deva': 'maithili',
+            'mal_Mlym': 'malayalam',
+            'mni_Beng': 'manipuri',
+            'mni_Mtei': 'manipuri',
+            'mar_Deva': 'marathi',
+            'npi_Deva': 'nepali',
+            'ory_Orya': 'odia',
+            'pan_Guru': 'punjabi',
+            'san_Deva': 'sanskrit',
+            'sat_Olck': 'santali',
+            'snd_Arab': 'sindhi',
+            'snd_Deva': 'sindhi',
+            'tam_Taml': 'tamil',
+            'tel_Telu': 'telugu',
+            'urd_Arab': 'urdu',
             'eng_Latn': 'english',
             'other': 'other'
         }
-        return code_to_lang.get(lang_code, lang_code.split('_')[0])
+        return code_to_lang.get(lang_code, lang_code.split('_')[0] if '_' in lang_code else lang_code)