Update app.py

app.py

@@ -0,0 +1,137 @@
+import os
+import re
+import gradio as gr
+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...")
+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]
+
+# ----------------------------
+# 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.FTN = fasttext.load_model(FTN_PATH)
+        self.FTR = fasttext.load_model(FTR_PATH)
+        self.BERT = torch.load(BERT_PATH, map_location=self.device)
+        self.BERT.eval()
+        self.tokenizer = AutoTokenizer.from_pretrained("ai4bharat/IndicBERTv2-MLM-only")
+        self.input_threshold = input_threshold
+        self.model_threshold = roman_lid_threshold
+        
+        # Official label map
+        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.FTN.predict(texts)
+        for (idx, txt), lbls, scrs in zip(data, preds[0], preds[1]):
+            out_dict[idx] = {"text": txt, "label": lbls[0][9:], "score": float(scrs[0]), "model": "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.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] = {"text": txt, "label": lbls[0][9:], "score": float(scrs[0]), "model": "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.tokenizer(list(texts), return_tensors="pt", padding=True,
+                                     truncation=True, max_length=512).to(self.device)
+                outputs = self.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()
+                    out_dict[i] = {
+                        "text": texts[batch_i],
+                        "label": self.label_map_reverse[label_idx],
+                        "score": probs[batch_i, label_idx].item(),
+                        "model": "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())]
+
+# ----------------------------
+# Gradio UI
+# ----------------------------
+lid_model = IndicLID()
+
+def detect(text_block):
+    lines = [l.strip() for l in text_block.splitlines() if l.strip()]
+    if not lines:
+        return []
+    return lid_model.batch_predict(lines)
+
+with gr.Blocks(title="IndicLID by AI4Bharat") as demo:
+    gr.Markdown("## IndicLID (AI4Bharat) — Full Ensemble\nDetects Indian languages in native & roman scripts.")
+    inp = gr.Textbox(lines=8, label="Enter one sentence per line")
+    out = gr.JSON(label="Predictions")
+    btn = gr.Button("Detect Language")
+    btn.click(fn=detect, inputs=inp, outputs=out)
+
+if __name__ == "__main__":
+    demo.launch()