Update app.py

app.py

@@ -2,24 +2,6 @@ import re
 import gradio as gr
 from entity_dictionary import ENTITY_DICTIONARY
 
-# ============================================================
-# 🔢 CATEGORY PRIORITY
-# ============================================================
-CATEGORY_PRIORITY = {
-    "DISTRICT": 1,
-    "CITY": 2,
-    "VILLAGE": 3,
-    "LANDMARK": 4,
-    "PERSON": 5,
-    "PRODUCT": 6,
-    "TIME": 7,
-    "DATE": 8,
-    "PHONE": 9,
-    "MONEY": 10,
-    "NUMBER": 11,
-    "UNKNOWN": 50
-}
-
 # ============================================================
 # 🔹 TOKENIZATION
 # ============================================================
@@ -27,75 +9,106 @@ def tokenize(text):
     return re.findall(r"\w+|[.!?]", text)
 
 # ============================================================
-# 🔍 BUILD LOWERCASED DICTIONARY LOOKUP
+# 🔍 DICTIONARY LOOKUP (case-insensitive, multi-word)
+# RETURNS (category, length, text) OR None
 # ============================================================
-def build_dictionary_lookup():
-    lookup = {}
-    for category, words in ENTITY_DICTIONARY.items():
-        if category == "REGEX_PATTERNS":
-            continue
-        for w in words:
-            lookup[w.lower()] = category
-    return lookup
+def find_dictionary_match(tokens, start):
+    for category, phrases in ENTITY_DICTIONARY.items():
+        for phrase in phrases:
+            p_tokens = phrase.split()
+            span = tokens[start:start + len(p_tokens)]
+
+            if len(span) != len(p_tokens):
+                continue
 
-DICT_LOOKUP = build_dictionary_lookup()
+            if [t.lower() for t in span] == [p.lower() for p in p_tokens]:
+                return category, len(p_tokens), " ".join(span)
+
+    return None
 
 # ============================================================
-# 🧠 ENTITY DETECTION (FINAL RULE SET)
+# 🧠 ENTITY DETECTION — FINAL RULESET
 # ============================================================
 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
 
         # ====================================================
-        # 1️⃣ LOWERCASE WORD RULE
+        # 1️⃣ FIRST WORD
         # ====================================================
-        if token[0].islower():
-            key = token.lower()
-            if key in DICT_LOOKUP:
+        if sentence_start:
+            match = find_dictionary_match(tokens, i)
+            if match:
+                category, length, text_val = match
                 entities.append({
-                    "text": token,
-                    "type": DICT_LOOKUP[key]
+                    "text": text_val,
+                    "type": category
                 })
-            i += 1
+                i += length
+            else:
+                i += 1
+
+            sentence_start = False
             continue
 
         # ====================================================
-        # 2️⃣ UPPERCASE WORD RULE (GROUPING)
+        # 2️⃣ UPPERCASE WORDS
         # ====================================================
-        phrase = [token]
-        j = i + 1
-
-        while (
-            j < len(tokens)
-            and tokens[j].isalpha()
-            and tokens[j][0].isupper()
-        ):
-            phrase.append(tokens[j])
-            j += 1
+        if token[0].isupper():
 
-        phrase_text = " ".join(phrase)          # original casing
-        phrase_key = phrase_text.lower()        # ✅ LOWERCASED FOR LOOKUP
-
-        if phrase_key in DICT_LOOKUP:
-            etype = DICT_LOOKUP[phrase_key]
-        else:
-            etype = "UNKNOWN"
+            # Try dictionary match (single or multi-word)
+            match = find_dictionary_match(tokens, i)
+            if match:
+                category, length, text_val = match
+                entities.append({
+                    "text": text_val,
+                    "type": category
+                })
+                i += length
+                continue
+
+            # Group consecutive uppercase words
+            phrase = [token]
+            j = i + 1
+            while j < len(tokens) and tokens[j].isalpha() and tokens[j][0].isupper():
+                phrase.append(tokens[j])
+                j += 1
+
+            entities.append({
+                "text": " ".join(phrase),
+                "type": "UNKNOWN"
+            })
+            i = j
+            continue
 
-        entities.append({
-            "text": phrase_text,
-            "type": etype
-        })
+        # ====================================================
+        # 3️⃣ LOWERCASE WORDS
+        # ====================================================
+        match = find_dictionary_match(tokens, i)
+        if match:
+            category, _, text_val = match
+            entities.append({
+                "text": text_val,
+                "type": category
+            })
 
-        i = j  # move past grouped phrase
+        i += 1
 
     return entities
 
@@ -105,10 +118,6 @@ def detect_entities(text):
 def detect_and_mask(text):
     entities = detect_entities(text)
 
-    entities.sort(
-        key=lambda e: (CATEGORY_PRIORITY.get(e["type"], 999), -len(e["text"]))
-    )
-
     masked = text
     entity_map = {}
     counters = {}
@@ -136,16 +145,15 @@ def detect_and_mask(text):
 # 🔓 UNMASKING
 # ============================================================
 def unmask_entities(masked_text, entity_map):
-    restored = masked_text
-    for ph in sorted(entity_map, key=len, reverse=True):
-        restored = restored.replace(ph, entity_map[ph])
-    return restored
+    for k in sorted(entity_map, key=len, reverse=True):
+        masked_text = masked_text.replace(k, entity_map[k])
+    return masked_text
 
 # ============================================================
 # 🖥️ GRADIO UI
 # ============================================================
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Entity Detection & Masking API")
+    gr.Markdown("# 🧠 Entity Masking API (Dictionary-Driven)")
 
     with gr.Tab("Detect & Mask"):
         text_in = gr.Textbox(label="Input Text")
@@ -159,4 +167,4 @@ with gr.Blocks() as demo:
         gr.Button("Unmask").click(unmask_entities, [masked, entity_map], out)
 
 demo.queue()
-demo.launch()
+demo.launch()