Update app.py

app.py

@@ -4,6 +4,7 @@ import torch
 import time
 import gc
 import re
+import string
 import threading
 from collections import OrderedDict
 from transformers import pipeline, AutoTokenizer
@@ -38,7 +39,7 @@ MODELS = {
     "Urdu": "prachuryyaIITG/Urdu_CLASSER_XLM",
 }
 
-# Fine-grained tag mappings into coarse categories (Excludes Product & CW)
+# Fine-grained tag mappings into coarse categories
 TAG_TO_COARSE = {
     # Person
     "Scientist": "PERSON", "Artist": "PERSON", "Athlete": "PERSON",
@@ -140,12 +141,32 @@ def cpu_fallback_infer(text, language):
     ner = get_pipeline(model_id, language, use_gpu=False)
     return ner(text, stride=64)
 
-# --- MULTILINGUAL REGEX PII ENGINE ---
-def extract_regex_spans(text):
+# --- PUNCTUATION & SPAN CLEANUP HELPER ---
+PUNCT_PATTERN = r'^[\s\.,!?;:"\'\(\)\[\]\{\}।॥،؟’”…—]+|[\s\.,!?;:"\'\(\)\[\]\{\}।॥،؟’”…—]+$'
+
+def clean_span_boundaries(text, start, end):
     """
-    Extracts structured technical PII using script-aware regex rules.
-    Supports Western, Devanagari, Bengali/Assamese, Persian/Arabic, Tamil, Telugu, and CJK digits.
+    Trims leading and trailing punctuation/whitespace from span character offsets.
+    Prevents punctuation attached to words (e.g. 'Real Madrid.') from being included in the entity.
     """
+    val = text[start:end]
+    
+    # Trim leading punctuation
+    leading_match = re.search(r'^[\s\.,!?;:"\'\(\)\[\]\{\}।॥،؟’”…—]+', val)
+    if leading_match:
+        start += leading_match.end()
+        val = text[start:end]
+        
+    # Trim trailing punctuation
+    trailing_match = re.search(r'[\s\.,!?;:"\'\(\)\[\]\{\}।॥،؟’”…—]+$', val)
+    if trailing_match:
+        end -= (trailing_match.end() - trailing_match.start())
+        val = text[start:end]
+        
+    return start, end, val
+
+# --- MULTILINGUAL REGEX PII ENGINE ---
+def extract_regex_spans(text):
     spans = []
     
     # Universal Email
@@ -153,15 +174,14 @@ def extract_regex_spans(text):
     for m in re.finditer(email_pattern, text):
         spans.append({'start': m.start(), 'end': m.end(), 'category': 'EMAIL', 'text': m.group()})
         
-    # Script-Aware Phone Number Regex (Supports Western, Devanagari, Bengali, Arabic/Farsi, CJK digits)
-    # Digits range: 0-9, \u0966-\u096F (Devanagari), \u09E6-\u09EF (Bengali/Assamese), \u0660-\u0669/\u06F0-\u06F9 (Arabic/Farsi), \u0B66-\u0B6F (Odia), \u0BE6-\u0BEF (Tamil), \u0C66-\u0C6F (Telugu), \uFF10-\uFF19 (Fullwidth CJK)
+    # Script-Aware Phone Numbers
     digits = r'0-9\u0966-\u096F\u09E6-\u09EF\u0660-\u0669\u06F0-\u06F9\u0B66-\u0B6F\u0BE6-\u0BEF\u0C66-\u0C6F\uFF10-\uFF19'
     phone_pattern = rf'(?:\+?[' + digits + r']{1,3}[-.\s]?)?\(?[' + digits + r']{2,4}\)?[-.\s]?[' + digits + r']{3,4}[-.\s]?[' + digits + r']{3,4}\b'
     for m in re.finditer(phone_pattern, text):
-        if len(re.sub(rf'[^{digits}]', '', m.group())) >= 7:  # Ensure minimum numeric digits
+        if len(re.sub(rf'[^{digits}]', '', m.group())) >= 7:
             spans.append({'start': m.start(), 'end': m.end(), 'category': 'PHONE', 'text': m.group()})
             
-    # IP Addresses (v4)
+    # IP Addresses
     ip_pattern = r'\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b'
     for m in re.finditer(ip_pattern, text):
         spans.append({'start': m.start(), 'end': m.end(), 'category': 'IP_ADDRESS', 'text': m.group()})
@@ -173,10 +193,10 @@ def extract_regex_spans(text):
         
     return spans
 
-# --- HYBRID ALIGNMENT & PSEUDONYMIZATION PIPELINE ---
+# --- HYBRID ANONYMIZATION PIPELINE ---
 def process_pii_anonymization(text, language):
     if not text.strip():
-        return "", [], {}
+        return "", {}
         
     # Phase 1: Model Inference (FgNER)
     try:
@@ -185,21 +205,26 @@ def process_pii_anonymization(text, language):
         print(f"Switching to CPU Fallback due to: {e}")
         raw_ner_results = cpu_fallback_infer(text, language)
         
-    # Filter FgNER spans to only PERSON, LOCATION, ORGANIZATION, MEDICAL
     ner_spans = []
     for res in raw_ner_results:
         entity_type = res.get('entity_group', res.get('entity', ''))
-        # Normalize entity tag string
         entity_clean = entity_type.replace("B-", "").replace("I-", "").split("_")[0]
         
         if entity_clean in TAG_TO_COARSE:
             coarse_cat = TAG_TO_COARSE[entity_clean]
-            ner_spans.append({
-                'start': int(res['start']),
-                'end': int(res['end']),
-                'category': coarse_cat,
-                'text': res['word'] if 'word' in res else text[int(res['start']):int(res['end'])]
-            })
+            start_pos = int(res['start'])
+            end_pos = int(res['end'])
+            
+            # Clean span boundaries from trailing/leading punctuation
+            start_pos, end_pos, clean_val = clean_span_boundaries(text, start_pos, end_pos)
+            
+            if clean_val.strip():
+                ner_spans.append({
+                    'start': start_pos,
+                    'end': end_pos,
+                    'category': coarse_cat,
+                    'text': clean_val
+                })
 
     # Phase 2: Regex Scanning
     regex_spans = extract_regex_spans(text)
@@ -211,24 +236,17 @@ def process_pii_anonymization(text, language):
     for current in all_spans:
         overlap = False
         for kept in filtered_spans:
-            # Check for character boundary overlap
             if not (current['end'] <= kept['start'] or current['start'] >= kept['end']):
                 overlap = True
                 break
         if not overlap:
             filtered_spans.append(current)
 
-    # Re-sort spans chronologically by start position
-    filtered_spans = sorted(filtered_spans, key=lambda x: x['start'])
-
     # Phase 4: Pseudonymization Mapping & Reverse Offset Replacement
     category_counters = {}
-    entity_mapping = {}  # Original text -> Pseudonym placeholder
-    reverse_mapping = {} # Pseudonym placeholder -> Original text
-    
-    highlighted_entities = []
+    entity_mapping = {}  
+    reverse_mapping = {} 
     
-    # First pass: Assign pseudonyms consistently across occurrences
     for span in filtered_spans:
         original_val = text[span['start']:span['end']]
         cat = span['category']
@@ -242,9 +260,8 @@ def process_pii_anonymization(text, language):
             pseudonym = entity_mapping[original_val]
             
         span['pseudonym'] = pseudonym
-        highlighted_entities.append((span['start'], span['end'], cat))
 
-    # Reverse-offset slicing (Back-to-Front) to prevent index shift corruption
+    # Reverse-offset slicing (Back-to-Front)
     sanitized_text = text
     for span in reversed(filtered_spans):
         start = span['start']
@@ -252,18 +269,7 @@ def process_pii_anonymization(text, language):
         pseudonym = span['pseudonym']
         sanitized_text = sanitized_text[:start] + pseudonym + sanitized_text[end:]
 
-    # Format for Gradio HighlightedText visualization
-    gradio_highlights = []
-    last_idx = 0
-    for span in filtered_spans:
-        if span['start'] > last_idx:
-            gradio_highlights.append((text[last_idx:span['start']], None))
-        gradio_highlights.append((text[span['start']:span['end']], span['category']))
-        last_idx = span['end']
-    if last_idx < len(text):
-        gradio_highlights.append((text[last_idx:], None))
-
-    return sanitized_text, gradio_highlights, reverse_mapping
+    return sanitized_text, reverse_mapping
 
 # --- GRADIO UI CONFIGURATION ---
 
@@ -284,7 +290,7 @@ body, .gradio-container {
 with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
 
     gr.Markdown("# Multilingual PII Anonymizer & Synthetic Pseudonymizer")
-    gr.Markdown("Remove and anonymize sensitive PII (**PERSON**, **LOCATION**, **ORGANIZATION**, **MEDICAL**, Emails, Phone Numbers, IPs, Credit Cards) across **21 languages** using Fine-Grained NER + Hybrid Multilingual Regex Rules.")
+    gr.Markdown("Anonymize PII (**PERSON**, **LOCATION**, **ORGANIZATION**, **MEDICAL**, Emails, Phones, IPs, Credit Cards) across **21 languages** into synthetic placeholders.")
 
     with gr.Row():
         lang_dropdown = gr.Dropdown(
@@ -306,35 +312,24 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
         elem_id="action-button"
     )
 
-    with gr.Row():
-        sanitized_output = gr.Textbox(
-            label="3. Anonymized Text (Synthetic Pseudonyms)",
-            lines=4,
-            interactive=False
-        )
-        
-    with gr.Row():
-        visual_diff = gr.HighlightedText(
-            label="Detected PII Entities",
-            combine_adjacent=True
-        )
-        
-        mapping_json = gr.JSON(
-            label="De-Anonymization Dictionary Map"
-        )
-
-    def run_pii_app(text, language):
-        sanitized, highlights, mapping = process_pii_anonymization(text, language)
-        return sanitized, highlights, mapping
+    sanitized_output = gr.Textbox(
+        label="3. Anonymized Text (Synthetic Pseudonyms)",
+        lines=4,
+        interactive=False
+    )
+    
+    mapping_json = gr.JSON(
+        label="4. De-Anonymization Dictionary Map"
+    )
 
     submit_btn.click(
-        fn=run_pii_app,
+        fn=process_pii_anonymization,
         inputs=[input_text, lang_dropdown],
-        outputs=[sanitized_output, visual_diff, mapping_json],
+        outputs=[sanitized_output, mapping_json],
         api_name="anonymize"
     )
 
-    gr.Markdown("### Try Multilingual Examples:")
+    gr.Markdown("### Try Examples across Languages:")
 
     gr.Examples(
         examples=[
@@ -346,12 +341,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
             ["Albert Einstein wurde in Ulm geboren. Er litt an Diabetes.", "German"],
             ["مرزا غالب دہلی میں رہتے تھے۔", "Urdu"],
             ["Victor Hugo est né à Besançon. Appelez le +33-1-4268-5300.", "French"],
-            ["Dante Alighieri è nato a Firenze.", "Italian"],
-            ["ਰਬਿੰਦਰਨਾਥ ਟੈਗੋਰ ਦਾ ਜਨਮ ਕੋਲਕਾਤਾ ਵਿੱਚ ��ੋਇਆ ਸੀ।", "Punjabi"],
         ],
         inputs=[input_text, lang_dropdown],
-        outputs=[sanitized_output, visual_diff, mapping_json],
-        fn=run_pii_app,
+        outputs=[sanitized_output, mapping_json],
+        fn=process_pii_anonymization,
         cache_examples=False
     )