Update app.py

app.py

@@ -11,7 +11,7 @@ with open(VOCAB_PATH, "r", encoding="utf-8") as f:
     VOCAB = {w.strip().lower() for w in f if w.strip()}
 
 # ============================================================
-# 🔢 CATEGORY PRIORITY
+# 🔢 CATEGORY PRIORITY (lower = higher priority)
 # ============================================================
 CATEGORY_PRIORITY = {
     "DISTRICT": 1,
@@ -35,19 +35,26 @@ def tokenize(text):
     return re.findall(r"\w+|[.!?]", text)
 
 # ============================================================
-# 🔍 DICTIONARY MATCHER
+# 🔍 DICTIONARY MATCHER (multi-word, case-insensitive)
 # ============================================================
 def find_dictionary_match(tokens, start_idx):
     for category, words in ENTITY_DICTIONARY.items():
         if category == "REGEX_PATTERNS":
             continue
+
         for w in words:
             w_tokens = w.split()
             span = tokens[start_idx:start_idx + len(w_tokens)]
+
             if len(span) != len(w_tokens):
                 continue
+
             if [t.lower() for t in span] == [x.lower() for x in w_tokens]:
-                return {"text": " ".join(span), "type": category, "length": len(w_tokens)}
+                return {
+                    "text": " ".join(span),
+                    "type": category,
+                    "length": len(w_tokens)
+                }
     return None
 
 # ============================================================
@@ -56,86 +63,149 @@ def find_dictionary_match(tokens, start_idx):
 def detect_entities(text):
     tokens = tokenize(text)
     entities = []
+
     i = 0
     sentence_start = True
 
     while i < len(tokens):
         token = tokens[i]
+
+        # Sentence boundary
         if token in [".", "!", "?"]:
             sentence_start = True
             i += 1
             continue
+
         if not token.isalpha():
             i += 1
             continue
 
-        # First word logic
+        # ====================================================
+        # RULE 1 — FIRST WORD OF SENTENCE
+        # ====================================================
         if sentence_start:
             match = find_dictionary_match(tokens, i)
+
             if match:
-                entities.append({"text": match["text"], "type": match["type"]})
+                entities.append({
+                    "text": match["text"],
+                    "type": match["type"]
+                })
                 i += match["length"]
             else:
                 if token.lower() not in VOCAB:
-                    entities.append({"text": token, "type": "UNKNOWN"})
+                    entities.append({
+                        "text": token,
+                        "type": "UNKNOWN"
+                    })
                 i += 1
+
             sentence_start = False
             continue
 
-        # Dictionary priority
+        # ====================================================
+        # RULE 2 — SECOND WORD ONWARDS
+        # ====================================================
+        # Dictionary always has priority
         match = find_dictionary_match(tokens, i)
         if match:
-            entities.append({"text": match["text"], "type": match["type"]})
+            entities.append({
+                "text": match["text"],
+                "type": match["type"]
+            })
             i += match["length"]
             continue
 
-        # Capital phrase grouping
+        # Capital phrase grouping (e.g. Shopping Mall)
         if token[0].isupper():
             phrase = [token]
             j = i + 1
-            while j < len(tokens) and tokens[j].isalpha() and tokens[j][0].isupper():
+
+            while (
+                j < len(tokens)
+                and tokens[j].isalpha()
+                and tokens[j][0].isupper()
+            ):
                 phrase.append(tokens[j])
                 j += 1
+
             phrase_text = " ".join(phrase)
-            entities.append({"text": phrase_text, "type": "UNKNOWN"})
+
+            entities.append({
+                "text": phrase_text,
+                "type": "UNKNOWN"
+            })
+
             i = j
             continue
 
+        # Normal word
         i += 1
 
-    # Regex patterns
+    # ====================================================
+    # REGEX RULES (UNCHANGED)
+    # ====================================================
     for etype, pattern in ENTITY_DICTIONARY.get("REGEX_PATTERNS", {}).items():
         for m in re.findall(pattern, text):
-            entities.append({"text": m, "type": etype})
-
-    # Deduplicate
-    seen, uniq = set(), []
+            entities.append({
+                "text": m,
+                "type": etype
+            })
+
+    # ====================================================
+    # DEDUPLICATION
+    # ====================================================
+    seen = set()
+    uniq = []
     for e in entities:
         key = (e["text"].lower(), e["type"])
         if key not in seen:
             seen.add(key)
             uniq.append(e)
+
     return uniq
 
 # ============================================================
-# 🎭 MASKING
+# 🎭 MASKING + AUTO-UNMASK (SAFE, BACKWARD-COMPATIBLE)
 # ============================================================
 def detect_and_mask(text):
     detected = detect_entities(text)
-    detected.sort(key=lambda e: (CATEGORY_PRIORITY.get(e["type"], 999), -len(e["text"])))
-    masked, entity_map, counters = text, {}, {}
+
+    detected.sort(
+        key=lambda e: (CATEGORY_PRIORITY.get(e["type"], 999), -len(e["text"]))
+    )
+
+    masked = text
+    entity_map = {}
+    counters = {}
 
     for e in detected:
         cat = e["type"]
         counters[cat] = counters.get(cat, 0) + 1
         placeholder = f"<<{cat}_{counters[cat]}>>"
-        masked = re.sub(rf"(?<!\w){re.escape(e['text'])}(?!\w)", placeholder, masked)
+
+        masked = re.sub(
+            rf"(?<!\w){re.escape(e['text'])}(?!\w)",
+            placeholder,
+            masked
+        )
+
         entity_map[placeholder] = e["text"]
 
-    return masked, entity_map, detected
+    # 🔑 AUTO-UNMASK (NEW, SAFE EXTENSION)
+    unmasked = masked
+    for ph in sorted(entity_map, key=len, reverse=True):
+        unmasked = unmasked.replace(ph, entity_map[ph])
+
+    return {
+        "masked_text": masked,
+        "entity_map": entity_map,
+        "entities": detected,
+        "unmasked_text": unmasked   # ✅ NEW FIELD (non-breaking)
+    }
 
 # ============================================================
-# 🔓 UNMASKING
+# 🔓 UNMASKING ENDPOINT (UNCHANGED)
 # ============================================================
 def unmask_entities(masked_text, entity_map):
     restored = masked_text
@@ -144,43 +214,21 @@ def unmask_entities(masked_text, entity_map):
     return restored
 
 # ============================================================
-# 🖥️ GRADIO UI (FIXED INTERFACE)
+# 🖥️ GRADIO UI
 # ============================================================
-def ui_mask_and_unmask(text):
-    masked_text, entity_map, entities = detect_and_mask(text)
-    unmasked_text = unmask_entities(masked_text, entity_map)
-    return {
-        "Masked Text": masked_text,
-        "Entity Map": entity_map,
-        "Entities": entities,
-        "Unmasked Result": unmasked_text
-    }
-
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Entity Detection & Masking API (Fixed Interface)")
-
-    with gr.Row():
-        text_in = gr.Textbox(label="Input Text", lines=3, placeholder="Type sentence here...")
-        run_btn = gr.Button("Run Masking + Unmasking")
-
-    with gr.Row():
-        masked_out = gr.Textbox(label="Masked Text", interactive=False)
-        unmasked_out = gr.Textbox(label="Unmasked Output", interactive=False)
-
-    with gr.Row():
-        map_out = gr.JSON(label="Entity Map")
-        ents_out = gr.JSON(label="Detected Entities")
-
-    def process_pipeline(text):
-        masked_text, entity_map, entities = detect_and_mask(text)
-        unmasked_text = unmask_entities(masked_text, entity_map)
-        return masked_text, unmasked_text, entity_map, entities
-
-    run_btn.click(
-        fn=process_pipeline,
-        inputs=text_in,
-        outputs=[masked_out, unmasked_out, map_out, ents_out]
-    )
+    gr.Markdown("# 🧠 Entity Detection & Masking API")
+
+    with gr.Tab("Detect & Mask"):
+        text_in = gr.Textbox(label="Input Text")
+        out = gr.JSON()
+        gr.Button("Run").click(detect_and_mask, text_in, out)
+
+    with gr.Tab("Unmask"):
+        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()