app.py

import spaces
import gradio as gr
import torch
import time
import gc
import re
import string
import threading
from collections import OrderedDict
from transformers import pipeline, AutoTokenizer

# ZeroGPU Configuration
MAX_MODELS_LOADED = 5 
MODEL_IDLE_TIMEOUT = 15 * 60   # 15 minutes
CLEANUP_INTERVAL = 15 * 60      # check every 15 minutes

# Complete list: 36 target languages
MODELS = {
    "Assamese": "prachuryyaIITG/CLASSER_Assamese_MuRIL",
    "Bengali": "prachuryyaIITG/MultiCoNER2_Bengali_XLM",
    "Bhojpuri": "prachuryyaIITG/FiNE-MiBBiC_Bhojpuri_MuRIL",
    "Bishnupriya": "prachuryyaIITG/FiNE-MiBBiC_Bishnupriya_MuRIL",
    "Bodo": "prachuryyaIITG/CLASSER_Bodo_MuRIL",
    "Chhattisgarhi": "prachuryyaIITG/FiNE-MiBBiC_Chhattisgarhi_MuRIL",
    "Chinese": "prachuryyaIITG/MultiCoNER2_Chinese_XLM",
    "Dogri": "prachuryyaIITG/SampurNER_Dogri_IndicBERTv2",
    "English": "prachuryyaIITG/MultiCoNER2_English_XLM",
    "Farsi": "prachuryyaIITG/MultiCoNER2_Farsi_XLM",
    "French": "prachuryyaIITG/MultiCoNER2_French_XLM",
    "German": "prachuryyaIITG/MultiCoNER2_German_XLM",
    "Gujarati": "prachuryyaIITG/SampurNER_Gujarati_IndicBERTv2",
    "Hindi": "prachuryyaIITG/MultiCoNER2_Hindi_XLM",
    "Italian": "prachuryyaIITG/MultiCoNER2_Italian_XLM",
    "Kannada": "prachuryyaIITG/SampurNER_Kannada_IndicBERTv2",
    "Kashmiri": "prachuryyaIITG/SampurNER_Kashmiri_IndicBERTv2",
    "Konkani": "prachuryyaIITG/SampurNER_Konkani_IndicBERTv2",
    "Maithili": "prachuryyaIITG/SampurNER_Maithili_IndicBERTv2",
    "Malayalam": "prachuryyaIITG/SampurNER_Malayalam_IndicBERTv2",
    "Manipuri": "prachuryyaIITG/FiNERVINER_Manipuri_IndicBERTv2",
    "Marathi": "prachuryyaIITG/CLASSER_Marathi_MuRIL",
    "Mizo": "prachuryyaIITG/FiNERVINER_Mizo_XLM",
    "Nepali": "prachuryyaIITG/CLASSER_Nepali_MuRIL",
    "Odia": "prachuryyaIITG/SampurNER_Odia_IndicBERTv2",
    "Portuguese": "prachuryyaIITG/MultiCoNER2_Portuguese_XLM",
    "Punjabi": "prachuryyaIITG/SampurNER_Punjabi_IndicBERTv2",
    "Sanskrit": "prachuryyaIITG/CLASSER_Sanskrit_MuRIL",
    "Santali": "prachuryyaIITG/SampurNER_Santali_IndicBERTv2",
    "Sindhi": "prachuryyaIITG/SampurNER_Sindhi_IndicBERTv2",
    "Spanish": "prachuryyaIITG/MultiCoNER2_Spanish_XLM",
    "Swedish": "prachuryyaIITG/MultiCoNER2_Swedish_XLM",
    "Tamil": "prachuryyaIITG/APTFiNER_Tamil_MuRIL",
    "Telugu": "prachuryyaIITG/APTFiNER_Telugu_MuRIL",
    "Ukrainian": "prachuryyaIITG/MultiCoNER2_Ukrainian_XLM",
    "Urdu": "prachuryyaIITG/Urdu_CLASSER_XLM",
}

# Unified mapping for MultiCoNER2, CLASSER, and FewNERD / SampurNER taxonomies
TAG_TO_COARSE = {
    # --- PERSON ---
    # MultiCoNER2 / CLASSER / FiNERVINER / APTFiNER
    "Scientist": "PERSON", "Artist": "PERSON", "Athlete": "PERSON",
    "Politician": "PERSON", "Cleric": "PERSON", "SportsManager": "PERSON",
    "OtherPER": "PERSON", "PER": "PERSON", "Person": "PERSON",
    # SampurNER (person-*)
    "Actor": "PERSON", "Artist/Author": "PERSON", "Director": "PERSON",
    "Scholar": "PERSON", "Soldier": "PERSON", "person-actor": "PERSON",
    "person-artist/author": "PERSON", "person-athlete": "PERSON",
    "person-director": "PERSON", "person-other": "PERSON",
    "person-politician": "PERSON", "person-scholar": "PERSON",
    "person-soldier": "PERSON",

    # --- LOCATION & FACILITIES ---
    # MultiCoNER2 / CLASSER
    "Facility": "LOCATION", "OtherLOC": "LOCATION",
    "HumanSettlement": "LOCATION", "Station": "LOCATION",
    "LOC": "LOCATION", "Location": "LOCATION",
    # FewNERD / SampurNER (location-* & building-*)
    "GPE": "LOCATION", "Body of Water": "LOCATION", "Island": "LOCATION",
    "Mountain": "LOCATION", "Park": "LOCATION", "Road/Transit": "LOCATION",
    "Airport": "LOCATION", "Hospital": "LOCATION", "Hotel": "LOCATION",
    "Library": "LOCATION", "Restaurant": "LOCATION", "Sports Facility": "LOCATION",
    "Theater": "LOCATION", "location-GPE": "LOCATION", "location-bodiesofwater": "LOCATION",
    "location-island": "LOCATION", "location-mountain": "LOCATION", "location-other": "LOCATION",
    "location-park": "LOCATION", "location-road/railway/highway/transit": "LOCATION",
    "building-airport": "LOCATION", "building-hospital": "LOCATION", "building-hotel": "LOCATION",
    "building-library": "LOCATION", "building-other": "LOCATION", "building-restaurant": "LOCATION",
    "building-sportsfacility": "LOCATION", "building-theater": "LOCATION",

    # --- ORGANIZATION ---
    # MultiCoNER2 / CLASSER
    "MusicalGRP": "ORGANIZATION", "PublicCORP": "ORGANIZATION",
    "PrivateCORP": "ORGANIZATION", "AerospaceManufacturer": "ORGANIZATION",
    "SportsGRP": "ORGANIZATION", "CarManufacturer": "ORGANIZATION",
    "ORG": "ORGANIZATION", "GRP": "ORGANIZATION", "Organization": "ORGANIZATION",
    # FewNERD / SampurNER (organization-*)
    "Company": "ORGANIZATION", "Education": "ORGANIZATION", "Government": "ORGANIZATION",
    "Media": "ORGANIZATION", "Political Party": "ORGANIZATION", "Religion": "ORGANIZATION",
    "Sports League": "ORGANIZATION", "Show Organization": "ORGANIZATION",
    "organization-company": "ORGANIZATION", "organization-education": "ORGANIZATION",
    "organization-government/governmentagency": "ORGANIZATION", "organization-media/newspaper": "ORGANIZATION",
    "organization-other": "ORGANIZATION", "organization-politicalparty": "ORGANIZATION",
    "organization-religion": "ORGANIZATION", "organization-showorganization": "ORGANIZATION",
    "organization-sportsleague": "ORGANIZATION", "organization-sportsteam": "ORGANIZATION",

    # --- MEDICAL ---
    # MultiCoNER2 / CLASSER
    "Medication/Vaccine": "MEDICAL", "MedicalProcedure": "MEDICAL",
    "AnatomicalStructure": "MEDICAL", "Symptom": "MEDICAL",
    "Disease": "MEDICAL", "MED": "MEDICAL", "Medical": "MEDICAL",
    # FewNERD / SampurNER (misc-*)
    "misc-disease": "MEDICAL", "misc-medical": "MEDICAL", "misc-biology": "MEDICAL"
}

# Cache and locking
pipelines = OrderedDict()
last_used = {}
lock = threading.Lock()

def clear_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

def load_pipeline(model_id, language, use_gpu=True):
    strategy = "simple" if language == "Chinese" else "first"
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    
    if use_gpu:
        device = 0 
        current_dtype = torch.bfloat16
    else:
        device = -1
        current_dtype = torch.float32
    
    return pipeline(
        "ner",
        model=model_id,
        tokenizer=tokenizer,
        aggregation_strategy=strategy,
        device=device,
        torch_dtype=current_dtype
    )

def get_pipeline(model_id, language, use_gpu=True):
    cache_key = f"{model_id}_{language}"
    with lock:
        now = time.time()
        if cache_key in pipelines:
            pipelines.move_to_end(cache_key)
            last_used[cache_key] = now
            return pipelines[cache_key]

        while len(pipelines) >= MAX_MODELS_LOADED:
            old_key, old_pipe = pipelines.popitem(last=False)
            del old_pipe
            last_used.pop(old_key, None)
            clear_memory()

        ner = load_pipeline(model_id, language, use_gpu=use_gpu)
        pipelines[cache_key] = ner
        last_used[cache_key] = now
        return ner

def cleanup_worker():
    while True:
        time.sleep(CLEANUP_INTERVAL)
        with lock:
            now = time.time()
            to_remove = [k for k, v in last_used.items() if now - v > MODEL_IDLE_TIMEOUT]
            for cache_key in to_remove:
                if cache_key in pipelines:
                    pipe = pipelines.pop(cache_key)
                    del pipe
                last_used.pop(cache_key, None)
            if to_remove:
                clear_memory()

threading.Thread(target=cleanup_worker, daemon=True).start()

@spaces.GPU
def try_gpu_infer(text, language):
    model_id = MODELS[language]
    ner = get_pipeline(model_id, language, use_gpu=True)
    return ner(text, stride=64)

def cpu_fallback_infer(text, language):
    model_id = MODELS[language]
    ner = get_pipeline(model_id, language, use_gpu=False)
    return ner(text, stride=64)

# --- PUNCTUATION & SPAN CLEANUP HELPER ---
def clean_span_boundaries(text, start, end):
    """
    Trims leading and trailing punctuation/whitespace from span character offsets.
    """
    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
    email_pattern = r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'
    for m in re.finditer(email_pattern, text):
        spans.append({'start': m.start(), 'end': m.end(), 'category': 'EMAIL', 'text': m.group()})
        
    # Script-Aware Phone Numbers
    digits = r'0-9\u0966-\u096F\u09E6-\u09EF\u0660-\u0669\u06F0-\u06F9\u0B66-\u0B6F\u0BE6-\u0BEF\u0C66-\u0C6F\u0A66-\u0A6F\u0AE6-\u0AEF\u0CDE-\u0CEF\u0D66-\u0D6F\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:
            spans.append({'start': m.start(), 'end': m.end(), 'category': 'PHONE', 'text': m.group()})
            
    # 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()})
        
    # Credit Card Numbers
    card_pattern = rf'\b(?:[' + digits + r']{4}[-\s]?){3}[' + digits + r']{4}\b'
    for m in re.finditer(card_pattern, text):
        spans.append({'start': m.start(), 'end': m.end(), 'category': 'CREDIT_CARD', 'text': m.group()})
        
    return spans

# --- HYBRID ANONYMIZATION PIPELINE ---
def process_pii_anonymization(text, language):
    if not text.strip():
        return "", {}
        
    # Phase 1: Model Inference (FgNER)
    try:
        raw_ner_results = try_gpu_infer(text, language)
    except Exception as e:
        print(f"Switching to CPU Fallback due to: {e}")
        raw_ner_results = cpu_fallback_infer(text, language)
        
    ner_spans = []
    for res in raw_ner_results:
        entity_type = res.get('entity_group', res.get('entity', ''))
        # Normalize entity tag string (handles B-, I-, sub-types)
        entity_clean = entity_type.replace("B-", "").replace("I-", "")
        
        # Check both full clean tag and prefix split
        matched_cat = None
        if entity_clean in TAG_TO_COARSE:
            matched_cat = TAG_TO_COARSE[entity_clean]
        elif entity_clean.split("_")[0] in TAG_TO_COARSE:
            matched_cat = TAG_TO_COARSE[entity_clean.split("_")[0]]
        elif entity_clean.split("-")[0] in TAG_TO_COARSE:
            matched_cat = TAG_TO_COARSE[entity_clean.split("-")[0]]
            
        if matched_cat:
            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': matched_cat,
                    'text': clean_val
                })

    # Phase 2: Regex Scanning
    regex_spans = extract_regex_spans(text)
    
    # Phase 3: Conflict Resolution (Prioritize Regex Spans over NER)
    # 1. First accept all valid non-overlapping regex spans
    filtered_spans = []
    for r_span in sorted(regex_spans, key=lambda x: (x['start'], -(x['end'] - x['start']))):
        if not any(not (r_span['end'] <= kept['start'] or r_span['start'] >= kept['end']) for kept in filtered_spans):
            filtered_spans.append(r_span)

    # 2. Add NER spans ONLY if they don't overlap with any accepted regex span
    for n_span in sorted(ner_spans, key=lambda x: (x['start'], -(x['end'] - x['start']))):
        overlap = any(not (n_span['end'] <= kept['start'] or n_span['start'] >= kept['end']) for kept in filtered_spans)
        if not overlap:
            filtered_spans.append(n_span)

    # Sort all resolved spans by start position
    filtered_spans = sorted(filtered_spans, key=lambda x: x['start'])

    # Phase 4: Pseudonymization Mapping & Reverse Offset Replacement
    category_counters = {}
    entity_mapping = {}  
    reverse_mapping = {} 
    
    for span in filtered_spans:
        original_val = text[span['start']:span['end']]
        cat = span['category']
        
        if original_val not in entity_mapping:
            category_counters[cat] = category_counters.get(cat, 0) + 1
            pseudonym = f"[{cat}_{category_counters[cat]}]"
            entity_mapping[original_val] = pseudonym
            reverse_mapping[pseudonym] = original_val
        else:
            pseudonym = entity_mapping[original_val]
            
        span['pseudonym'] = pseudonym

    # Reverse-offset slicing (Back-to-Front)
    sanitized_text = text
    for span in reversed(filtered_spans):
        start = span['start']
        end = span['end']
        pseudonym = span['pseudonym']
        sanitized_text = sanitized_text[:start] + pseudonym + sanitized_text[end:]

    return sanitized_text, reverse_mapping

# --- GRADIO UI CONFIGURATION ---

custom_css = """
body, .gradio-container {
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", "Noto Sans", Helvetica, Arial, sans-serif !important;
}
#action-button {
    background-color: #00568b !important;
    color: white !important;
    border: none !important;
}
#action-button:hover {
    background-color: #00488b !important;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:

    gr.Markdown("# Multilingual PII Anonymizer & Synthetic Pseudonymizer")
    gr.Markdown("Anonymize sensitive PII (**PERSON**, **LOCATION**, **ORGANIZATION**, **MEDICAL**, Emails, Phones, IPs, Credit Cards) across **36 languages** into synthetic placeholders.")

    with gr.Row():
        lang_dropdown = gr.Dropdown(
            choices=list(MODELS.keys()),
            value="English",
            label="1. Select Language"
        )

    input_text = gr.Textbox(
        value="Jude Bellingham joined Real Madrid in 2023. You can reach him at jude@realmadrid.es or +15550199.",
        placeholder="Type or paste multilingual text here...",
        label="2. Input Text",
        lines=4
    )

    submit_btn = gr.Button(
        "Anonymize PII",
        variant="primary",
        elem_id="action-button"
    )

    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=process_pii_anonymization,
        inputs=[input_text, lang_dropdown],
        outputs=[sanitized_output, mapping_json],
        api_name="anonymize"
    )

    gr.Markdown("### Try Examples across Languages:")

    gr.Examples(
        examples=[
            ["Jude Bellingham joined Real Madrid in 2023. You can reach him at jude@realmadrid.es or +15550199.", "English"],
            ["姚明出生于上海。联系电话是 +8613800138000。", "Chinese"],
            ["अमिताभ बच्चन मुंबई में रहते हैं। उनका ईमेल contact@bachchan.com है।", "Hindi"],
            ["Madrid es la capital de España. Contactar con Dr. Garcia al +34912345678.", "Spanish"],
            ["সকলোৱে ভাল পায় জুবিন গাৰ্গক। গুৱাহাটীত তেওঁৰ ঘৰ।", "Assamese"],
            ["Albert Einstein wurde in Ulm geboren. Er litt an Diabetes.", "German"],
            ["مرزا غالب دہلی میں رہتے تھے۔", "Urdu"],
            ["Victor Hugo est né à Besançon. Appelez le +33142685300.", "French"],
            ["ਰਵਿੰਦਰਨਾਥ ਟੈਗੋਰ ਕਲਕੱਤੇ ਵਿੱਚ ਰਹਿੰਦੇ ਸਨ।", "Punjabi"],
            ["ಶಿವರಾಮ ಕಾರಂತರು ಪುತ್ತೂರಿನಲ್ಲಿ ಜನಿಸಿದರು.", "Kannada"],
        ],
        inputs=[input_text, lang_dropdown],
        outputs=[sanitized_output, mapping_json],
        fn=process_pii_anonymization,
        cache_examples=False
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_error=True)