Update app.py

app.py

@@ -1,8 +1,8 @@
 import re, json, gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
+# LIGHTWEIGHT, CPU-FRIENDLY MODEL
+MODEL_ID = "google/mt5-small"
 LABELS = ["PERSON","ORG","LOC","GPE","DATE","TIME","PRODUCT","EVENT"]
 
 def build_prompt(text, labels=LABELS):
@@ -29,7 +29,7 @@ def extract_json(s: str):
         except Exception:
             return {"entities": []}
 
-# lazy globals
+# Lazy load on CPU
 _tokenizer = None
 _model = None
 
@@ -37,51 +37,57 @@ def load_model():
     global _tokenizer, _model
     if _tokenizer is None or _model is None:
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
-        _model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else None,
-            device_map="auto"
-        )
+        _model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)  # CPU by default on Spaces
     return _tokenizer, _model
 
-def ner_infer(text, temperature=0.0, max_new_tokens=256):
-    if not text.strip():
+def ner_infer(text, max_new_tokens=192):
+    text = (text or "").strip()
+    if not text:
         return {"entities": []}
+
     tok, model = load_model()
     prompt = build_prompt(text)
-    inputs = tok(prompt, return_tensors="pt").to(model.device)
+    inputs = tok(prompt, return_tensors="pt")  # stays on CPU
 
     gen_ids = model.generate(
         **inputs,
         max_new_tokens=int(max_new_tokens),
-        do_sample=(float(temperature) > 0),
-        temperature=float(temperature),
-        pad_token_id=tok.eos_token_id or tok.pad_token_id,
+        do_sample=False,          # deterministic on CPU
+        temperature=0.0,
+        # pad_token_id must be set for some T5/mT5 variants:
+        pad_token_id=tok.pad_token_id,
+        eos_token_id=tok.eos_token_id
     )
-    out = tok.decode(gen_ids[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    out = tok.decode(gen_ids[0], skip_special_tokens=True)
     data = extract_json(out)
 
-    # normalize schema
+    # normalize; if model omits start/end, compute first occurrence
     ents = []
     for e in data.get("entities", []):
         try:
-            t = e["text"]; lab = e["label"]
-            st = int(e.get("start", 0)); en = int(e.get("end", st + len(t)))
-            ents.append({"text": t, "label": lab, "start": st, "end": en})
+            t = str(e.get("text","")).strip()
+            lab = str(e.get("label","")).strip()
+            if not t or not lab: 
+                continue
+            st = e.get("start"); en = e.get("end")
+            if not isinstance(st, int) or not isinstance(en, int):
+                idx = text.find(t)
+                if idx >= 0:
+                    st, en = idx, idx + len(t)
+                else:
+                    st, en = 0, 0
+            ents.append({"text": t, "label": lab, "start": int(st), "end": int(en)})
         except Exception:
             pass
     return {"entities": ents}
 
-with gr.Blocks(title="Persian Zero-Shot NER (LLM)") as demo:
-    gr.Markdown("## Persian Zero-Shot NER (LLM) — JSON output")
+with gr.Blocks(title="Persian Zero-Shot NER (CPU)") as demo:
+    gr.Markdown("## Persian Zero-Shot NER (LLM) — CPU version (mT5-small)")
     inp = gr.Textbox(label="متن فارسی", lines=4, value="من دیروز با علی در تهران در دفتر دیجی‌کالا جلسه داشتم.")
-    with gr.Row():
-        temp = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature")
-        max_tok = gr.Slider(64, 512, value=256, step=16, label="Max new tokens")
+    max_tok = gr.Slider(64, 512, value=192, step=16, label="Max new tokens (CPU)")
     btn = gr.Button("Extract Entities")
     out = gr.JSON(label="خروجی JSON")
-
-    btn.click(fn=ner_infer, inputs=[inp, temp, max_tok], outputs=out)
+    btn.click(fn=ner_infer, inputs=[inp, max_tok], outputs=out)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()