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)