Update app.py

app.py

@@ -1,4 +1,4 @@
-# Import spaces & other imports
+# Import spaces and others
 import spaces
 import gradio as gr
 import torch
@@ -10,12 +10,12 @@ from transformers import pipeline
 
 # ZeroGPU Configuration
 MAX_MODELS_LOADED = 5
-MODEL_IDLE_TIMEOUT = 15 * 60   # 15 minutes
-CLEANUP_INTERVAL = 5 * 60      # check every 5 minutes
+MODEL_IDLE_TIMEOUT = 15 * 60
+CLEANUP_INTERVAL = 5 * 60
 
 DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-# Model mapping (unchanged)
+# Model mapping
 MODELS = {
     "Assamese": "prachuryyaIITG/CLASSER_Assamese_MuRIL",
     "Bengali": "prachuryyaIITG/MultiCoNER2_Bengali_XLM",
@@ -55,7 +55,7 @@ MODELS = {
     "Urdu": "prachuryyaIITG/CLASSER_Urdu_XLM",
 }
 
-# Cache and locking
+# Model cache with LRU
 pipelines = OrderedDict()
 last_used = {}
 lock = threading.Lock()
@@ -69,7 +69,7 @@ def clear_memory():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
 
-# Load pipeline (ZeroGPU safe)
+# Load pipeline
 def load_pipeline(model_id):
 
     ner = pipeline(
@@ -82,7 +82,7 @@ def load_pipeline(model_id):
 
     return ner
 
-# LRU manager
+# Pipeline manager
 def get_pipeline(model_id):
 
     with lock:
@@ -111,7 +111,7 @@ def get_pipeline(model_id):
 
         return ner
 
-# Background cleanup
+# Cleanup thread
 def cleanup_worker():
 
     while True:
@@ -122,26 +122,26 @@ def cleanup_worker():
 
             now = time.time()
 
-            to_remove = []
+            remove = []
 
             for model_id in last_used:
 
                 if now - last_used[model_id] > MODEL_IDLE_TIMEOUT:
-                    to_remove.append(model_id)
+                    remove.append(model_id)
 
-            for model_id in to_remove:
+            for model_id in remove:
 
                 del pipelines[model_id]
                 del last_used[model_id]
 
-            if to_remove:
+            if remove:
                 clear_memory()
 
 threading.Thread(target=cleanup_worker, daemon=True).start()
 
-# GPU inference (required)
+# GPU inference API (Agent-compatible)
 @spaces.GPU
-def gpu_infer(text, language):
+def predict(text, language):
 
     model_id = MODELS[language]
 
@@ -149,20 +149,9 @@ def gpu_infer(text, language):
 
     results = ner(text)
 
-    return {
-        "text": text,
-        "entities": results
-    }
+    return results
 
-# Public wrapper
-def get_ner_results(text, language):
-
-    if not text:
-        return {"text": "", "entities": []}
-
-    return gpu_infer(text, language)
-
-# Original CSS (unchanged)
+# CSS
 custom_css = """
 #blue-button {
     background-color: #00568b !important;
@@ -173,24 +162,23 @@ custom_css = """
     background-color: #00488b !important;
 }
 """
-# UI with ORIGINAL MARKDOWN (exact)
+
+# UI (Original Markdown preserved exactly)
 with gr.Blocks() as demo:
 
     gr.Markdown("# AWED-FiNER: Agents, Web applications, and Expert Detectors for Fine-grained Named Entity Recognition across 36 languages for 6.6 billion speakers")
 
     gr.Markdown("A tool for Fine-grained NER across 36 languages, ranging from four most spoken languages (English, Chinese, Hindi, & Spanish) to four extremely low-resource vulnerable languages (Bishnupriya, Bodo, Manipuri, & Mizo)")
 
-    gr.Markdown("The models used in this tool are based on the [SampurNER](https://huggingface.co/datasets/prachuryyaIITG/SampurNER), [MultiCoNERv2](https://huggingface.co/datasets/MultiCoNER/multiconer_v2), [CLASSER](https://huggingface.co/datasets/prachuryyaIITG/CLASSER), [FiNE-MiBBiC](https://github.com/PrachuryyaKaushik/FiNE-MiBBiC), [FiNERVINER](https://huggingface.co/datasets/prachuryyaIITG/FiNERVINER) and [APTFiNER](https://huggingface.co/datasets/prachuryyaIITG/APTFiNER) research works.")
+    gr.Markdown("The models used in this tool are based on the [SampurNER](https://huggingface.co/datasets/prachuryyaIITG/SampurNER), [MultiCoNERv2](https://huggingface.co/datasets/MultiCoNER/multiconer_v2), [CLASSER](https://huggingface.co/datasets/prachuryyaIITG/CLASSER), FiNE-MiBBiC, [FiNERVINER](https://huggingface.co/datasets/prachuryyaIITG/FiNERVINER) and [APTFiNER](https://huggingface.co/datasets/prachuryyaIITG/APTFiNER) research works.")
 
     gr.Markdown("To access the agentic tool, please check out: [AWED-FiNER tool](https://github.com/PrachuryyaKaushik/AWED-FiNER)")
 
-    with gr.Row():
-
-        lang_dropdown = gr.Dropdown(
-            choices=list(MODELS.keys()),
-            value="English",
-            label="1. Select Language"
-        )
+    lang_dropdown = gr.Dropdown(
+        choices=list(MODELS.keys()),
+        value="English",
+        label="1. Select Language"
+    )
 
     input_text = gr.Textbox(
         value="Jude Bellingham joined Real Madrid in 2023.",
@@ -199,11 +187,7 @@ with gr.Blocks() as demo:
         lines=3
     )
 
-    submit_btn = gr.Button(
-        "Analyze",
-        variant="primary",
-        elem_id="blue-button"
-    )
+    submit_btn = gr.Button("Analyze", elem_id="blue-button")
 
     output = gr.HighlightedText(
         label="3. Fine-grained NER Results",
@@ -211,7 +195,7 @@ with gr.Blocks() as demo:
     )
 
     submit_btn.click(
-        fn=get_ner_results,
+        fn=predict,
         inputs=[input_text, lang_dropdown],
         outputs=output,
         api_name="predict"
@@ -236,7 +220,7 @@ with gr.Blocks() as demo:
             ["சச்சின் டெண்டுல்கர் சென்னையில் கிரிக்கெட் விளையாடினார்.", "Tamil"],
             ["فردوسی در توس به دنیا آمد.", "Farsi"],
             ["Dante Alighieri è nato a Firenze.", "Italian"],
-            ["રવિન્દ્રનાથ ટાગોરે કલકત્તામાં ગીતાંજલિ લખી હતી.", "Gujarati"],
+            ["રવિન્દ્રનাথ ટાગોરે કલકત્તામાં ગીતાંજલિ લખી હતી.", "Gujarati"],
             ["ಕೆ. ಶಿವರಾಮ ಕಾರಂತರು ಪುತ್ತೂರಿನಲ್ಲಿ ಜನಿಸಿದರು.", "Kannada"],
             ["Тарас Шевченко народився в Моринцях.", "Ukrainian"],
             ["भानुभक्त आचार्य तनहुँमा जन्मिएका हुन्।", "Nepali"],
@@ -247,7 +231,7 @@ with gr.Blocks() as demo:
         ],
         inputs=[input_text, lang_dropdown],
         outputs=output,
-        fn=get_ner_results,
+        fn=predict,
         cache_examples=False
     )