Update app.py

app.py

@@ -1,29 +1,45 @@
-# app.py
-# Persian Zero-Shot NER (CPU) — Hugging Face Spaces (Gradio)
-# Uses a lightweight Seq2Seq model (mT5-small) and slow tokenizer (no GPU deps).
-
-import re
-import json
-import gradio as gr
+# app.py — Persian Zero-Shot NER (CPU) with few-shot prompting + beams
+import re, json, gradio as gr
 from typing import Dict, Any, List
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# ---- Config (CPU-friendly) ----
-MODEL_ID = "google/mt5-small"
-ALLOWED_LABELS: List[str] = [
-    "PERSON", "ORG", "LOC", "GPE", "DATE", "TIME", "PRODUCT", "EVENT"
-]
+MODEL_ID = "google/mt5-small"   # try "google/mt5-base" on CPU if still empty (slower, better)
+ALLOWED_LABELS: List[str] = ["PERSON","ORG","LOC","GPE","DATE","TIME","PRODUCT","EVENT"]
 DEFAULT_EXAMPLE = "من دیروز با علی در تهران در دفتر دیجی‌کالا جلسه داشتم."
 
-# ---- Prompt & Parsing ----
+# --- Few-shot examples (in Persian) to nudge the model ---
+FEW_SHOT = """
+نمونه ۱:
+متن: من با علی در تهران در شرکت دیجی‌کالا جلسه داشتم.
+خروجی:
+{"entities":[
+  {"text":"علی","label":"PERSON","start":7,"end":10},
+  {"text":"تهران","label":"LOC","start":14,"end":19},
+  {"text":"دیجی‌کالا","label":"ORG","start":29,"end":37}
+]}
+
+نمونه ۲:
+متن: سارا فردا ساعت ۱۰ در دانشگاه تهران سخنرانی دارد.
+خروجی:
+{"entities":[
+  {"text":"سارا","label":"PERSON","start":0,"end":4},
+  {"text":"فردا","label":"DATE","start":5,"end":9},
+  {"text":"۱۰","label":"TIME","start":15,"end":17},
+  {"text":"دانشگاه تهران","label":"ORG","start":21,"end":34}
+]}
+"""
+
 def build_prompt(text: str, labels: List[str]) -> str:
     return (
         "متن زیر را برای شناسایی موجودیت‌های نامدار (NER) تحلیل کن.\n"
         f"لیبل‌های مجاز: {', '.join(labels)}.\n"
-        "خروجی را فقط به صورت JSON معتبر با اسکیمای زیر بده:\n"
+        "فقط JSON معتبر با اسکیمای زیر را برگردان:\n"
         '{"entities":[{"text":"...", "label":"ORG|PERSON|...", "start":0, "end":0}]}\n'
-        "هیچ متن دیگری ننویس؛ فقط JSON.\n\n"
+        "هیچ متن دیگری ننویس؛ فقط JSON.\n"
+        + FEW_SHOT +
+        "\nاکنون متن زیر را پردازش کن و فقط JSON بده:\n"
         f"متن: {text}\n"
+        "خروجی:\n"
     )
 
 def extract_first_json(s: str) -> Dict[str, Any]:
@@ -31,11 +47,9 @@ def extract_first_json(s: str) -> Dict[str, Any]:
     if not m:
         return {"entities": []}
     raw = m.group(0)
-    # try direct parse
     try:
         return json.loads(raw)
     except Exception:
-        # quick repairs for trailing commas
         raw = re.sub(r",\s*}", "}", raw)
         raw = re.sub(r",\s*]", "]", raw)
         try:
@@ -44,72 +58,62 @@ def extract_first_json(s: str) -> Dict[str, Any]:
             return {"entities": []}
 
 def normalize_entities(data: Dict[str, Any], text: str, labels: List[str]) -> Dict[str, Any]:
+    text_norm = text or ""
     out = []
     for e in data.get("entities", []):
         try:
-            t = str(e.get("text", "")).strip()
-            lab = str(e.get("label", "")).strip().upper()
-            if not t or not lab:
+            t = str(e.get("text","")).strip()
+            lab = str(e.get("label","")).strip().upper()
+            if not t or lab not in labels:
                 continue
-            # keep only allowed labels
-            if lab not in labels:
-                continue
-            st = e.get("start"); en = e.get("end")
+            st, en = e.get("start"), e.get("end")
             if not isinstance(st, int) or not isinstance(en, int) or st < 0 or en < 0:
-                # fallback: first occurrence
-                idx = text.find(t)
-                if idx >= 0:
-                    st, en = idx, idx + len(t)
-                else:
-                    st, en = 0, 0
+                idx = text_norm.find(t)
+                st, en = (idx, idx+len(t)) if idx >= 0 else (0, 0)
             out.append({"text": t, "label": lab, "start": int(st), "end": int(en)})
         except Exception:
-            # ignore malformed entries
             pass
     return {"entities": out}
 
-# ---- Lazy model load (CPU) ----
+# lazy CPU load
 _tokenizer = None
 _model = None
-
 def load_model():
     global _tokenizer, _model
     if _tokenizer is None or _model is None:
-        # IMPORTANT: use_fast=False to avoid SentencePiece fast-conversion issues on CPU
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
         _model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
     return _tokenizer, _model
 
-# ---- Inference ----
-def ner_infer(text: str, max_new_tokens: int = 192) -> Dict[str, Any]:
+def ner_infer(text: str, max_new_tokens: int = 256, num_beams: int = 4) -> Dict[str, Any]:
     text = (text or "").strip()
     if not text:
         return {"entities": []}
     tok, model = load_model()
     prompt = build_prompt(text, ALLOWED_LABELS)
-    inputs = tok(prompt, return_tensors="pt")  # stays on CPU
+    inputs = tok(prompt, return_tensors="pt")   # CPU
     gen_ids = model.generate(
         **inputs,
         max_new_tokens=int(max_new_tokens),
-        do_sample=False,    # deterministic for stable outputs on CPU
-        temperature=0.0,
+        do_sample=False,          # deterministic
+        num_beams=int(num_beams), # stronger decoding than greedy on CPU
+        length_penalty=1.05,
         pad_token_id=tok.pad_token_id,
-        eos_token_id=tok.eos_token_id
+        eos_token_id=tok.eos_token_id,
     )
     out_text = tok.decode(gen_ids[0], skip_special_tokens=True)
     raw = extract_first_json(out_text)
     return normalize_entities(raw, text, ALLOWED_LABELS)
 
-# ---- UI ----
 with gr.Blocks(title="Persian Zero-Shot NER (CPU)") as demo:
-    gr.Markdown("## Persian Zero-Shot NER (LLM) — **CPU version (mT5-small)**")
-    with gr.Row():
-        inp = gr.Textbox(label="متن فارسی", lines=4, value=DEFAULT_EXAMPLE)
+    gr.Markdown("## Persian Zero-Shot NER — CPU (mT5) + Few-Shot Prompting")
+    inp = gr.Textbox(label="متن فارسی", lines=4, value=DEFAULT_EXAMPLE)
     with gr.Row():
-        max_tok = gr.Slider(64, 512, value=192, step=16, label="حداکثر توکن خروجی (CPU)")
+        max_tok = gr.Slider(96, 512, value=256, step=16, label="حداکثر توکن خروجی")
+        beams = gr.Slider(1, 8, value=4, step=1, label="Beam size")
     btn = gr.Button("استخراج موجودیت‌ها")
-    out = gr.JSON(label="خروجی JSON (entities)")
-    btn.click(fn=ner_infer, inputs=[inp, max_tok], outputs=out)
+    out = gr.JSON(label="خروجی JSON")
+    btn.click(fn=ner_infer, inputs=[inp, max_tok, beams], outputs=out)
 
 if __name__ == "__main__":
     demo.launch()