Update app.py

app.py

@@ -1,16 +1,18 @@
 import gradio as gr
 import re
-from entity_dictionary import ENTITY_DICTIONARY, WEEKDAY_MAP, MONTH_MAP
+from entity_dictionary import ENTITY_DICTIONARY
 
 # ============================================================
-# LOAD VOCABULARY (Bhojpuri common words)
+# 📘 LOAD VOCABULARY
 # ============================================================
 VOCAB_PATH = "vocab.txt"
+
 with open(VOCAB_PATH, "r", encoding="utf-8") as f:
     VOCAB = {w.strip().lower() for w in f if w.strip()}
 
+
 # ============================================================
-# PRIORITY ORDER — Dictionary > Regex > UNKNOWN
+# 🔢 CATEGORY PRIORITY (lower = higher priority)
 # ============================================================
 CATEGORY_PRIORITY = {
     "DISTRICT": 1,
@@ -19,48 +21,37 @@ CATEGORY_PRIORITY = {
     "LANDMARK": 4,
     "PERSON": 5,
     "PRODUCT": 6,
-    "WEEKDAY": 7,   # ✅ added
-    "MONTH": 8,     # ✅ added
-    "TIME": 9,
-    "DATE": 10,
-    "PHONE": 11,
-    "MONEY": 12,
-    "NUMBER": 13,
+    "TIME": 7,
+    "DATE": 8,
+    "PHONE": 9,
+    "MONEY": 10,
+    "NUMBER": 11,
     "UNKNOWN": 50
 }
 
+
 # ============================================================
-# NORMALIZE TEXT SPACING (removes double spaces + fixes . Tu)
+# ✨ TEXT NORMALISATION
 # ============================================================
 def normalize_sentence_spacing(text):
-    # 1️⃣ Remove multiple spaces
     text = re.sub(r"\s+", " ", text)
-
-    # 2️⃣ Remove spaces directly after punctuation
     text = re.sub(r"([.!?])\s+", r"\1", text)
-
-    # 3️⃣ Trim leading/trailing spaces
     return text.strip()
 
-# ============================================================
-# HELPER — Split text into tokens, keeping punctuation
-# ============================================================
+
 def split_with_punctuation(text):
     return re.findall(r"\w+|[.!?]", text)
 
+
 # ============================================================
-# CAPITAL PHRASE DETECTION (UNKNOWN)
+# 🔍 CAPITAL PHRASE DETECTION (UNKNOWN)
 # ============================================================
 def detect_capital_phrases(text, dictionary_entities):
     tokens = text.split()
     capital_entities = []
 
     dict_full = {e["text"].lower() for e in dictionary_entities}
-    dict_words = {
-        part
-        for e in dictionary_entities
-        for part in e["text"].lower().split()
-    }
+    dict_words = {p for e in dictionary_entities for p in e["text"].lower().split()}
 
     i = 1  # skip first word
     while i < len(tokens):
@@ -73,21 +64,15 @@ def detect_capital_phrases(text, dictionary_entities):
         if len(w) > 1 and w[0].isupper():
             phrase = [w]
             j = i + 1
-            while j < len(tokens) and len(tokens[j]) > 1 and tokens[j][0].isupper():
+            while j < len(tokens) and tokens[j][0].isupper():
                 phrase.append(tokens[j])
                 j += 1
 
-            full_phrase = " ".join(phrase)
-            full_lower = full_phrase.lower()
+            full_phrase = " ".join(phrase).lower()
 
-            if (
-                full_lower in dict_full
-                or all(tok.lower() in dict_words for tok in phrase)
-            ):
-                i = j
-                continue
+            if full_phrase not in dict_full and not all(p.lower() in dict_words for p in phrase):
+                capital_entities.append({"text": " ".join(phrase), "type": "UNKNOWN"})
 
-            capital_entities.append({"text": full_phrase, "type": "UNKNOWN"})
             i = j
             continue
 
@@ -95,121 +80,87 @@ def detect_capital_phrases(text, dictionary_entities):
 
     return capital_entities
 
+
 # ============================================================
-# ENTITY DETECTION
+# 🧠 ENTITY DETECTION PIPELINE
 # ============================================================
 def detect_entities(text):
     entities = []
-
-    # ---------- Dictionary detection ----------
     dictionary_entities = []
+
+    # ---------- Dictionary ----------
     for category, words in ENTITY_DICTIONARY.items():
         if category == "REGEX_PATTERNS":
             continue
 
         for w in words:
-            pattern = re.compile(
-                rf"(?<![A-Za-z]){re.escape(w)}(?![A-Za-z])",
-                re.IGNORECASE
-            )
+            pattern = re.compile(rf"(?<![A-Za-z]){re.escape(w)}(?![A-Za-z])", re.IGNORECASE)
             for m in pattern.finditer(text):
-                dictionary_entities.append(
-                    {"text": m.group(), "type": category}
-                )
+                dictionary_entities.append({"text": m.group(), "type": category})
 
     entities.extend(dictionary_entities)
 
-    # ---------- Regex detection ----------
+    # ---------- Regex ----------
     for etype, pattern in ENTITY_DICTIONARY["REGEX_PATTERNS"].items():
         for m in re.findall(pattern, text):
             entities.append({"text": m, "type": etype})
 
-    # ---------- Sentence-aware first-word UNKNOWN logic ----------
+    # ---------- Sentence-start UNKNOWN ----------
     text = normalize_sentence_spacing(text)
     tokens = split_with_punctuation(text)
-    sentence_start = True
 
+    sentence_start = True
     for w in tokens:
-        w = w.strip()
-        if not w:
-            continue
-
         if w in [".", "!", "?"]:
             sentence_start = True
             continue
 
-        if not re.match(r"^[A-Za-z]+$", w):
+        if not w.isalpha():
             continue
 
         if sentence_start:
             lw = w.lower()
-
-            if any(ent["text"].lower() == lw for ent in dictionary_entities):
-                sentence_start = False
-                continue
-
             if lw not in VOCAB:
                 entities.append({"text": w, "type": "UNKNOWN"})
-
             sentence_start = False
 
-    # ---------- Capital UNKNOWN detection ----------
+    # ---------- Capital UNKNOWN ----------
     entities.extend(detect_capital_phrases(text, dictionary_entities))
 
     # ---------- Deduplicate ----------
-    uniq = []
     seen = set()
-    for ent in entities:
-        key = (ent["text"].lower(), ent["type"])
+    uniq = []
+    for e in entities:
+        key = (e["text"].lower(), e["type"])
         if key not in seen:
             seen.add(key)
-            uniq.append(ent)
+            uniq.append(e)
 
-    # ---------- Remove NUMBER inside TIME ----------
-    cleaned = []
-    for e in uniq:
-        if e["type"] == "NUMBER":
-            if any(
-                t["type"] == "TIME" and e["text"] in t["text"]
-                for t in uniq
-            ):
-                continue
-        cleaned.append(e)
+    return uniq
 
-    return cleaned
 
 # ============================================================
-# MASK ENTITIES
+# 🎭 MASK / UNMASK
 # ============================================================
 def mask_entities(text):
     detected = detect_entities(text)
 
-    sorted_entities = sorted(
-        detected,
-        key=lambda e: (
-            CATEGORY_PRIORITY.get(e["type"], 999),
-            -len(e["text"])
-        )
+    detected.sort(
+        key=lambda e: (CATEGORY_PRIORITY.get(e["type"], 999), -len(e["text"]))
     )
 
     masked = text
     entity_map = {}
     counters = {}
 
-    for ent in sorted_entities:
-        cat = ent["type"]
+    for e in detected:
+        cat = e["type"]
         counters[cat] = counters.get(cat, 1)
         placeholder = f"<{cat}_{counters[cat]}>"
         counters[cat] += 1
 
-        masked = re.sub(
-            re.escape(ent["text"]),
-            placeholder,
-            masked,
-            flags=re.IGNORECASE
-        )
-
-        entity_map[placeholder] = ent["text"]
+        masked = re.sub(re.escape(e["text"]), placeholder, masked, flags=re.IGNORECASE)
+        entity_map[placeholder] = e["text"]
 
     return {
         "masked_text": masked,
@@ -217,68 +168,34 @@ def mask_entities(text):
         "entities": detected
     }
 
-# ============================================================
-# UNMASK ENTITIES (WEEKDAY + MONTH SAFE)
-# ============================================================
-def unmask_entities(masked_text, entity_map):
-    unmasked = masked_text
-
-    for ph in sorted(entity_map.keys(), key=len, reverse=True):
-        original = entity_map[ph]
-        lw = original.lower()
-
-        if ph.startswith("<WEEKDAY_"):
-            replacement = WEEKDAY_MAP.get(lw, original)
 
-        elif ph.startswith("<MONTH_"):
-            replacement = MONTH_MAP.get(lw, original)
-
-        else:
-            replacement = original
-
-        unmasked = unmasked.replace(ph, replacement)
+def unmask_entities(masked_text, entity_map):
+    for ph in sorted(entity_map, key=len, reverse=True):
+        masked_text = masked_text.replace(ph, entity_map[ph])
+    return masked_text
 
-    return unmasked
 
 # ============================================================
-# GRADIO UI
+# 🖥️ GRADIO UI
 # ============================================================
 with gr.Blocks() as demo:
-    gr.Markdown(
-        "# 🧠 Entity Detection & Masking API\n"
-        "### Bhojpuri ⇄ Kreol Morisien (Weekdays & Months Enabled)"
-    )
+    gr.Markdown("# 🧠 Entity Detection & Masking API")
 
     with gr.Tab("Detect Entities"):
-        detect_in = gr.Textbox(label="Input Text")
-        detect_out = gr.JSON(label="Detected Entities")
-        gr.Button("Detect").click(
-            detect_entities,
-            detect_in,
-            detect_out,
-            api_name="detect_entities"
-        )
+        text_in = gr.Textbox(label="Input Text")
+        out = gr.JSON()
+        gr.Button("Detect").click(detect_entities, text_in, out)
 
     with gr.Tab("Mask Entities"):
-        mask_in = gr.Textbox(label="Input Text")
-        mask_out = gr.JSON(label="Masking Output")
-        gr.Button("Mask").click(
-            mask_entities,
-            mask_in,
-            mask_out,
-            api_name="mask_entities"
-        )
+        text_in = gr.Textbox(label="Input Text")
+        out = gr.JSON()
+        gr.Button("Mask").click(mask_entities, text_in, out)
 
     with gr.Tab("Unmask Entities"):
-        unmask_text = gr.Textbox(label="Masked Text")
-        unmask_map = gr.JSON(label="Entity Map")
-        unmask_out = gr.Textbox(label="Unmasked Output")
-        gr.Button("Unmask").click(
-            unmask_entities,
-            [unmask_text, unmask_map],
-            unmask_out,
-            api_name="unmask_entities"
-        )
+        masked = gr.Textbox(label="Masked Text")
+        entity_map = gr.JSON(label="Entity Map")
+        out = gr.Textbox(label="Unmasked Output")
+        gr.Button("Unmask").click(unmask_entities, [masked, entity_map], out)
 
 demo.queue()
 demo.launch()