Update indiclid_inference.py

indiclid_inference.py

@@ -0,0 +1,134 @@
+import os
+import re
+import pandas as pd
+import fasttext
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+# ------------------------------
+# Download required models
+# ------------------------------
+print("Downloading IndicLID models from Hugging Face...")
+
+FTN_PATH = hf_hub_download("ai4bharat/IndicLID-FTN", filename="model_baseline_roman.bin")
+FTR_PATH = hf_hub_download("ai4bharat/IndicLID-FTR", filename="model_baseline_roman.bin")
+BERT_PATH = hf_hub_download("ai4bharat/IndicLID-BERT", filename="basline_nn_simple.pt")
+
+print("Download complete.")
+
+# ------------------------------
+# Data helper for BERT batching
+# ------------------------------
+class IndicBERT_Data(Dataset):
+    def __init__(self, indices, X):
+        self.x = list(X)
+        self.i = list(indices)
+    def __len__(self):
+        return len(self.x)
+    def __getitem__(self, idx):
+        return self.i[idx], self.x[idx]
+
+# ------------------------------
+# Main IndicLID Class
+# ------------------------------
+class IndicLID:
+    def __init__(self, input_threshold=0.5, roman_lid_threshold=0.6):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.IndicLID_FTN = fasttext.load_model(FTN_PATH)
+        self.IndicLID_FTR = fasttext.load_model(FTR_PATH)
+        self.IndicLID_BERT = torch.load(BERT_PATH, map_location=self.device)
+        self.IndicLID_BERT.eval()
+        self.IndicLID_BERT_tokenizer = AutoTokenizer.from_pretrained("ai4bharat/IndicBERTv2-MLM-only")
+        self.input_threshold = input_threshold
+        self.model_threshold = roman_lid_threshold
+
+        # Official label map (index -> language code)
+        self.label_map_reverse = {
+            0: 'asm_Latn', 1: 'ben_Latn', 2: 'brx_Latn', 3: 'guj_Latn',
+            4: 'hin_Latn', 5: 'kan_Latn', 6: 'kas_Latn', 7: 'kok_Latn',
+            8: 'mai_Latn', 9: 'mal_Latn', 10: 'mni_Latn', 11: 'mar_Latn',
+            12: 'nep_Latn', 13: 'ori_Latn', 14: 'pan_Latn', 15: 'san_Latn',
+            16: 'snd_Latn', 17: 'tam_Latn', 18: 'tel_Latn', 19: 'urd_Latn',
+            20: 'eng_Latn', 21: 'other', 22: 'asm_Beng', 23: 'ben_Beng',
+            24: 'brx_Deva', 25: 'doi_Deva', 26: 'guj_Gujr', 27: 'hin_Deva',
+            28: 'kan_Knda', 29: 'kas_Arab', 30: 'kas_Deva', 31: 'kok_Deva',
+            32: 'mai_Deva', 33: 'mal_Mlym', 34: 'mni_Beng', 35: 'mni_Meti',
+            36: 'mar_Deva', 37: 'nep_Deva', 38: 'ori_Orya', 39: 'pan_Guru',
+            40: 'san_Deva', 41: 'sat_Olch', 42: 'snd_Arab', 43: 'tam_Tamil',
+            44: 'tel_Telu', 45: 'urd_Arab'
+        }
+
+    def char_percent_check(self, text):
+        total_chars = sum(c.isalpha() for c in text)
+        roman_chars = sum(bool(re.match(r"[A-Za-z]", c)) for c in text)
+        return roman_chars / total_chars if total_chars else 0
+
+    def native_inference(self, data, out_dict):
+        if not data: return out_dict
+        texts = [x[1] for x in data]
+        preds = self.IndicLID_FTN.predict(texts)
+        for (idx, txt), lbls, scrs in zip(data, preds[0], preds[1]):
+            out_dict[idx] = (txt, lbls[0][9:], float(scrs[0]), 'IndicLID-FTN')
+        return out_dict
+
+    def ftr_inference(self, data, out_dict, batch_size):
+        if not data: return out_dict
+        texts = [x[1] for x in data]
+        preds = self.IndicLID_FTR.predict(texts)
+        bert_inputs = []
+        for (idx, txt), lbls, scrs in zip(data, preds[0], preds[1]):
+            if float(scrs[0]) > self.model_threshold:
+                out_dict[idx] = (txt, lbls[0][9:], float(scrs[0]), 'IndicLID-FTR')
+            else:
+                bert_inputs.append((idx, txt))
+        return self.bert_inference(bert_inputs, out_dict, batch_size)
+
+    def bert_inference(self, data, out_dict, batch_size):
+        if not data: return out_dict
+        ds = IndicBERT_Data([x[0] for x in data], [x[1] for x in data])
+        dl = DataLoader(ds, batch_size=batch_size)
+        with torch.no_grad():
+            for idxs, texts in dl:
+                enc = self.IndicLID_BERT_tokenizer(
+                    list(texts), return_tensors="pt", padding=True,
+                    truncation=True, max_length=512
+                ).to(self.device)
+                outputs = self.IndicLID_BERT(**enc)
+                preds = torch.argmax(outputs.logits, dim=1)
+                probs = torch.softmax(outputs.logits, dim=1)
+                for batch_i, p in enumerate(preds):
+                    i = idxs[batch_i].item()
+                    label_idx = p.item()
+                    label = self.label_map_reverse[label_idx]
+                    score = probs[batch_i, label_idx].item()
+                    out_dict[i] = (texts[batch_i], label, score, 'IndicLID-BERT')
+        return out_dict
+
+    def batch_predict(self, texts, batch_size=8):
+        native, roman = [], []
+        for i, t in enumerate(texts):
+            if self.char_percent_check(t) > self.input_threshold:
+                roman.append((i, t))
+            else:
+                native.append((i, t))
+        out_dict = {}
+        out_dict = self.native_inference(native, out_dict)
+        out_dict = self.ftr_inference(roman, out_dict, batch_size)
+        return [out_dict[i] for i in sorted(out_dict.keys())]
+
+# ------------------------------
+# Quick test if run directly
+# ------------------------------
+if __name__ == "__main__":
+    detector = IndicLID()
+    samples = [
+        "यह एक हिंदी वाक्य है।",        # Hindi (native)
+        "ennai pudikkuma?",             # Tamil (romanized)
+        "ఇది ఒక తెలుగు వాక్యం",         # Telugu (native)
+        "Hello, how are you?"           # English
+    ]
+    results = detector.batch_predict(samples)
+    for text, label, score, model in results:
+        print(f"Text: {text}\nPredicted: {label} | Score: {score:.4f} | Model: {model}\n")