import streamlit as st import torch import torch.nn as nn import os import re import json from pathlib import Path from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModel from collections import defaultdict from huggingface_hub import hf_hub_download # Model paths for NER and Coreference Resolution on HuggingFace Hub. MODEL_NER_PATH = "inesctec/CitiLink-XLMR-Anonymization-pt" MODEL_COREF_PATH = "inesctec/CitiLink-SpanBERT-Coreference-pt" LABEL_NAMES = { "Name": ("Name", "Nome"), "AdministrativeInformation": ("Admin", "Admin"), "PositionDepartment": ("Position", "Cargo"), "Address": ("Address", "Morada"), "Date": ("Date", "Data"), "Location": ("Location", "Local"), "PersonalInformation": ("Info", "Info"), "Company": ("Company", "Empresa"), "ArtisticActivity": ("Artistic", "Artístico"), "Degree": ("Degree", "Curso"), "Time": ("Time", "Hora"), "LicensePlate": ("License", "Matrícula"), "Job": ("Job", "Profissão"), "Vehicle": ("Vehicle", "Veículo"), "Faculty": ("Faculty", "Faculdade"), "Family": ("Family", "Família"), "Other": ("Other", "Outro"), } UI = { "page_title": ("Text Pseudonymization Demo", "Demo de Pseudonimização de Texto"), "page_subtitle": ("Automatic text pseudonymization using NER + Coreference Resolution for municipal minutes", "Pseudonimização automática de texto utilizando NER + Resolução de Correferência para atas municipais"), "sidebar_header": ("⚙️ Configuration", "⚙️ Configuração"), "choose_example": ("Choose an example:", "Escolha um exemplo:"), "lang_toggle_label": ("Language", "Idioma"), "about_header": ("### 📊 About", "### 📊 Sobre"), "about_text": ( "- **Anonymization (NER)** uses Token Classification to identify sensitive entities.\n" "- **Linking (Coreference Resolution)** links mentions of the same entity across the document.\n" "- **Model (NER)**: XLM-RoBERTa fine-tuned for Portuguese NER.\n" "- **Model (Coreference Resolution)**: SpanBERT fine-tuned on Portuguese administrative texts.\n" "- **Language**: Portuguese (pt-pt).", "- **Anonimização (NER)** usa Classificação de Tokens para identificar entidades sensíveis.\n" "- **Ligação (Resolução de Correferência)** liga menções da mesma entidade ao longo do documento.\n" "- **Modelo (NER)**: XLM-RoBERTa ajustado para NER em Português.\n" "- **Modelo (Resolução de Correferência)**: SpanBERT ajustado em textos administrativos portugueses.\n" "- **Língua**: Português (PT-PT)." ), "resources_header": ("### 🔗 Resources", "### 🔗 Recursos"), "how_it_works": ("🎯 How it works", "🎯 Como funciona"), "how_text": ( "The pseudonymization pipeline uses two specialized models:\n\n" "1. **Detection (NER):** Identifies all sensitive entity spans across 17 categories.\n" "2. **Linking (Coreference Resolution):** Detects when different mentions refer to the same real-world entity, assigning consistent IDs.", "O pipeline de pseudonimização utiliza dois modelos especializados:\n\n" "1. **Deteção (NER):** Identifica todos os spans de entidades sensíveis em 17 categorias.\n" "2. **Ligação (Resolução de Correferência):** Deteta quando diferentes menções se referem à mesma entidade real, atribuindo IDs consistentes." ), "entities_title": ("Supported Entities:", "Entidades Suportadas:"), "input_label": ("INPUT:", "ENTRADA:"), "output_label": ("OUTPUT:", "SAÍDA:"), "col1_header": ("📝 Input Document", "📝 Documento de Entrada"), "col2_header": ("🔒 Pseudonymization Results", "🔒 Resultados da Pseudonimização"), "textarea_label": ("Enter your text here:", "Introduza o texto aqui:"), "textarea_placeholder": ("Paste your document text here...", "Cole aqui o texto do documento..."), "anonymize_btn": ("🔍 Pseudonymize", "🔍 Pseudonimizar"), "spinner_text": ("Processing...", "A processar..."), "tab_anon": ("📄 Pseudonymized Text", "📄 Texto Pseudonimizado"), "tab_entities": ("🔍 Extracted Entities", "🔍 Entidades Extraídas"), "tab_groups": ("🔗 Coreference Groups", "🔗 Grupos de Correferência"), "no_groups": ("No coreference groups found (no entity appears more than once).", "Nenhum grupo de correferência encontrado (nenhuma entidade aparece mais do que uma vez)."), "group_label": ("Group", "Grupo"), "mentions_label": ("mentions", "menções"), "no_entities": ("No entities found.", "Nenhuma entidade encontrada."), "json_expander": ("📋 Full Entity Report (JSON)", "📋 Relatório Completo (JSON)"), "download_btn": ("📥 Download JSON Report", "📥 Transferir Relatório JSON"), "no_results_info": ("Please process a document on the left to view the results.", "Por favor, processe um documento à esquerda para ver os resultados."), "metric_entities": ("Entities", "Entidades"), "metric_categories": ("Categories", "Categorias"), "metric_clusters": ("Groups", "Grupos"), "example_names": ( { "Custom Text": "Custom Text", "1º Portuguese Meeting Minute": "1º Portuguese Meeting Minute", "2º Portuguese Meeting Minute": "2º Portuguese Meeting Minute", "3º Portuguese Meeting Minute": "3º Portuguese Meeting Minute", }, { "Custom Text": "Texto Livre", "1º Portuguese Meeting Minute": "1ª Ata de Reunião Municipal", "2º Portuguese Meeting Minute": "2ª Ata de Reunião Municipal", "3º Portuguese Meeting Minute": "3ª Ata de Reunião Municipal", } ), } st.set_page_config( page_title="Text Pseudonymization Demo", page_icon="🛡️", layout="wide", initial_sidebar_state="expanded" ) st.markdown(""" """, unsafe_allow_html=True) # Returns the active language index: 1 for Portuguese, 0 for English. def lang_idx(): return 1 if st.session_state.get("lang_pt", True) else 0 # Returns the UI string for the given key in the active language. def t(key): return UI[key][lang_idx()] # Returns the display name of an entity label in the active language. def label_name(internal_tag): entry = LABEL_NAMES.get(internal_tag) if entry is None: return internal_tag return entry[lang_idx()] # Maps an integer distance to the nearest distance bucket index. def bucket_distance(d): buckets = [1, 2, 3, 4, 5, 8, 16, 32, 64] for i, b in enumerate(buckets): if d <= b: return i return len(buckets) # Feed-forward neural network with two hidden layers and dropout, used as mention and antecedent scorer. class FFNN(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.3): super().__init__() self.net = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, output_dim) ) def forward(self, x): return self.net(x) # Coreference resolution model based on mention-ranking, with a SpanBERT encoder and FFNN scorers. class CorefModel(nn.Module): def __init__(self, model_dir, ffnn_dim=1024, max_span_width=52, max_antecedents=50, max_cluster_size=12, dropout=0.3): super().__init__() self.encoder = AutoModel.from_pretrained(model_dir) self.hidden = self.encoder.config.hidden_size self.max_span_width = max_span_width self.max_antecedents = max_antecedents self.max_cluster_size = max_cluster_size self.width_emb = nn.Embedding(max_span_width + 1, 20) self.dist_emb = nn.Embedding(10, 20) span_dim = self.hidden * 3 + 20 self.mention_scorer = FFNN(span_dim, ffnn_dim, 1, dropout) ant_dim = span_dim * 3 + 20 self.antecedent_scorer = FFNN(ant_dim, ffnn_dim, 1, dropout) self.span_attn = nn.Linear(self.hidden, 1) # Builds the vector representation of a span from its token embeddings. def span_repr(self, token_embs, start, end): h_start = token_embs[start] h_end = token_embs[min(end, len(token_embs) - 1)] span_tokens = token_embs[start: min(end + 1, len(token_embs))] attn_w = torch.softmax(self.span_attn(span_tokens), dim=0) head = (attn_w * span_tokens).sum(0) width = torch.tensor(min(end - start, self.max_span_width), dtype=torch.long) w_emb = self.width_emb(width) return torch.cat([h_start, h_end, head, w_emb], dim=-1) # Runs the forward pass: encodes tokens, computes span vectors and returns (i, j, score) pairs. def forward(self, input_ids, attention_mask, spans): out = self.encoder(input_ids=input_ids, attention_mask=attention_mask) token_embs = out.last_hidden_state[0] span_vecs = [] for (s, e) in spans: span_vecs.append(self.span_repr(token_embs, s, e)) if not span_vecs: return [] span_vecs = torch.stack(span_vecs) m_scores = self.mention_scorer(span_vecs).squeeze(-1) n = len(spans) pairs = [] for i in range(n): for j in range(max(0, i - self.max_antecedents), i): d = i - j d_idx = torch.tensor(bucket_distance(d), dtype=torch.long) d_emb = self.dist_emb(d_idx) v = torch.cat([span_vecs[i], span_vecs[j], span_vecs[i] * span_vecs[j], d_emb], dim=-1) score = (m_scores[i] + m_scores[j] + self.antecedent_scorer(v).squeeze()) pairs.append((i, j, score.item())) return pairs # Find operation of the Union-Find algorithm with path compression. def find(parent, x): if parent[x] != x: parent[x] = find(parent, parent[x]) return parent[x] # Union operation of the Union-Find algorithm by rank. def union(parent, rank, a, b): ra, rb = find(parent, a), find(parent, b) if ra == rb: return if rank[ra] < rank[rb]: ra, rb = rb, ra parent[rb] = ra if rank[ra] == rank[rb]: rank[ra] += 1 @st.cache_resource # Loads and caches NER templates and HuggingFace matching. def load_models(): try: tokenizer_ner = AutoTokenizer.from_pretrained(MODEL_NER_PATH, add_prefix_space=True) model_ner = AutoModelForTokenClassification.from_pretrained(MODEL_NER_PATH) model_ner.eval() coref_cfg = { "ffnn_dim": 1024, "max_span_width": 52, "max_antecedents": 50, "max_cluster_size": 12, "dropout": 0.3, "max_tokens": 512, "stride": 128, "epsilon": 0.0001, "min_cluster_score": 0.3, "ner_union_margin": 0.5, } model_coref = CorefModel( model_dir=MODEL_COREF_PATH, ffnn_dim=coref_cfg["ffnn_dim"], max_span_width=coref_cfg["max_span_width"], max_antecedents=coref_cfg["max_antecedents"], max_cluster_size=coref_cfg["max_cluster_size"], dropout=coref_cfg["dropout"], ) model_coref = CorefModel( model_dir=MODEL_COREF_PATH, ffnn_dim=coref_cfg["ffnn_dim"], max_span_width=coref_cfg["max_span_width"], max_antecedents=coref_cfg["max_antecedents"], max_cluster_size=coref_cfg["max_cluster_size"], dropout=coref_cfg["dropout"], ) cabecas_path = hf_hub_download( repo_id=MODEL_COREF_PATH, filename="cabecas_coref.pt" ) cabecas = torch.load(cabecas_path, map_location="cpu", weights_only=False) model_coref.load_state_dict(cabecas, strict=False) model_coref.eval() tokenizer_coref = AutoTokenizer.from_pretrained(MODEL_COREF_PATH) return tokenizer_ner, model_ner, model_coref, tokenizer_coref, coref_cfg, None except Exception as e: return None, None, None, None, None, str(e) # Aggregates contextual token embeddings for the full document using a sliding window with overlap. def aggregate_token_embeddings(model_coref, tokenizer_coref, tokens, max_tokens=512, stride=128, device="cpu"): emb_sum = torch.zeros(len(tokens), model_coref.hidden) emb_count = torch.zeros(len(tokens)) encoded = tokenizer_coref( tokens, is_split_into_words=True, return_offsets_mapping=False, add_special_tokens=False, return_tensors=None ) word_ids_full = encoded.word_ids() subword_ids = encoded["input_ids"] total_subwords = len(subword_ids) start = 0 while start < total_subwords: end = min(start + max_tokens - 2, total_subwords) chunk_ids = [tokenizer_coref.cls_token_id] + subword_ids[start:end] + [tokenizer_coref.sep_token_id] chunk_tensor = torch.tensor([chunk_ids], dtype=torch.long).to(device) attn_mask = torch.ones_like(chunk_tensor) with torch.no_grad(): out = model_coref.encoder(input_ids=chunk_tensor, attention_mask=attn_mask) hidden = out.last_hidden_state[0] for sub_pos, (sub_id, h) in enumerate(zip(range(start, end), hidden[1:-1])): word_idx = word_ids_full[sub_id] if word_idx is not None: emb_sum[word_idx] += h.cpu() emb_count[word_idx] += 1 if end >= total_subwords: break start += max_tokens - 2 - stride count_safe = emb_count.unsqueeze(-1).clamp(min=1) return emb_sum / count_safe # Runs coreference resolution over NER spans and returns clusters of coreferential mentions. def predict_coref_clusters(model_coref, tokenizer_coref, tokens, ner_spans, coref_cfg, device="cpu"): if not ner_spans: return [] max_tokens = coref_cfg.get("max_tokens", 512) stride = coref_cfg.get("stride", 128) epsilon = coref_cfg.get("epsilon", 0.0001) min_cluster_score = coref_cfg.get("min_cluster_score", 0.3) ner_union_margin = coref_cfg.get("ner_union_margin", 0.5) threshold = max(epsilon, min_cluster_score) + ner_union_margin token_embs = aggregate_token_embeddings(model_coref, tokenizer_coref, tokens, max_tokens, stride, device) spans = [(s, e) for (s, e, _) in ner_spans] span_labels = [lbl for (_, _, lbl) in ner_spans] model_coref.eval() with torch.no_grad(): span_vecs = [] for (s, e) in spans: sv = model_coref.span_repr(token_embs, s, e) span_vecs.append(sv) if not span_vecs: return [] span_vecs_t = torch.stack(span_vecs) m_scores = model_coref.mention_scorer(span_vecs_t).squeeze(-1) n = len(spans) parent = list(range(n)) rank = [0] * n for i in range(n): for j in range(max(0, i - model_coref.max_antecedents), i): if span_labels[i] != span_labels[j]: continue d = i - j d_idx = torch.tensor(bucket_distance(d), dtype=torch.long) d_emb = model_coref.dist_emb(d_idx) v = torch.cat([span_vecs_t[i], span_vecs_t[j], span_vecs_t[i] * span_vecs_t[j], d_emb], dim=-1) score = (m_scores[i] + m_scores[j] + model_coref.antecedent_scorer(v).squeeze()).item() if score > threshold: union(parent, rank, i, j) cluster_map = defaultdict(list) for i, (s, e) in enumerate(spans): root = find(parent, i) cluster_map[root].append((s, e)) return [v for v in cluster_map.values() if len(v) > 1] # Runs the full pseudonymization pipeline: NER, coreference, ID assignment and anonymized text generation. def process_anonymization(text, tokenizer_ner, model_ner, model_coref, tokenizer_coref, coref_cfg): if not text.strip(): return "Por favor, insira um texto.", [] # ── PASSO 1: NER ────────────────────────────────────────────────────── id2label = model_ner.config.id2label inputs = tokenizer_ner( text, truncation=True, max_length=512, stride=164, return_overflowing_tokens=True, return_offsets_mapping=True, padding=True, return_tensors="pt" ) offset_mapping_all = inputs.pop("offset_mapping") inputs.pop("overflow_to_sample_mapping") input_ids_all = inputs["input_ids"] with torch.no_grad(): logits = model_ner(**inputs).logits all_predictions = torch.argmax(logits, dim=2).tolist() SPACE_PREFIXES = [" ", "▁", "Ġ"] entidades_brutas = [] for window_idx, (predictions, offsets) in enumerate(zip(all_predictions, offset_mapping_all)): temp_entity = {"tokens": [], "start": None, "end": None, "label": None} for idx, (pred_id, offset) in enumerate(zip(predictions, offsets)): label_name_str = id2label[pred_id] start_char, end_char = int(offset[0]), int(offset[1]) if start_char == end_char: continue tag_base = label_name_str.split('-', 1)[1] if '-' in label_name_str else None token_id = input_ids_all[window_idx][idx].item() token = tokenizer_ner.convert_ids_to_tokens(token_id) comeca_nova_palavra = any(token.startswith(p) for p in SPACE_PREFIXES) if label_name_str.startswith("B-"): if temp_entity["label"]: chars_to_remove = ".,;:!?" if temp_entity["label"] != "PERSONAL-PositionDepartment" else ",;:!?" while (temp_entity["end"] - temp_entity["start"]) > 0 and text[temp_entity["end"] - 1] in chars_to_remove: temp_entity["end"] -= 1 entidades_brutas.append(temp_entity) temp_entity = {"start": start_char, "end": end_char, "label": tag_base, "tokens": [token]} elif label_name_str.startswith("I-") and temp_entity["label"] == tag_base: temp_entity["tokens"].append(token) temp_entity["end"] = end_char elif not comeca_nova_palavra and temp_entity["label"] is not None: temp_entity["tokens"].append(token) temp_entity["end"] = end_char else: if temp_entity["label"]: chars_to_remove = ".,;:!?" while (temp_entity["end"] - temp_entity["start"]) > 0 and text[temp_entity["end"] - 1] in chars_to_remove: temp_entity["end"] -= 1 entidades_brutas.append(temp_entity) temp_entity = {"tokens": [], "start": None, "end": None, "label": None} if temp_entity["label"]: entidades_brutas.append(temp_entity) entidades_brutas.sort(key=lambda x: x["start"]) entidades_finais = [] for atual in entidades_brutas: if not entidades_finais: entidades_finais.append(atual) continue ultima = entidades_finais[-1] distancia = atual["start"] - ultima["end"] if distancia <= 1 and atual["label"] == ultima["label"]: ultima["end"] = atual["end"] ultima["tokens"].extend(atual["tokens"]) else: adicionar = True for i, selecionada in enumerate(entidades_finais): interseccao = max(0, min(atual["end"], selecionada["end"]) - max(atual["start"], selecionada["start"])) if interseccao > 0: if (atual["end"] - atual["start"]) > (selecionada["end"] - selecionada["start"]): entidades_finais[i] = atual adicionar = False break if adicionar: entidades_finais.append(atual) SINAIS_REMOVER = ".-,;:_^~´`/*+!?\"'" for ent in entidades_finais: while (ent["end"] - ent["start"]) > 0 and text[ent["end"] - 1] in SINAIS_REMOVER: ent["end"] -= 1 # ── PASSO 2: Correferência ──────────────────────────────────────────── tokens_doc = text.split() ner_spans_coref = [] char_to_token = {} pos = 0 for tok_idx, tok in enumerate(tokens_doc): for c in range(pos, pos + len(tok)): char_to_token[c] = tok_idx pos += len(tok) + 1 for ent in entidades_finais: s_tok = char_to_token.get(ent["start"]) e_tok = char_to_token.get(min(ent["end"] - 1, len(text) - 1)) if s_tok is not None and e_tok is not None: lbl = ent["label"].replace("PERSONAL-", "") if ent["label"] else "Other" ner_spans_coref.append((s_tok, e_tok, lbl)) clusters = predict_coref_clusters(model_coref, tokenizer_coref, tokens_doc, ner_spans_coref, coref_cfg) # ── PASSO 3: Atribuição de IDs ─────────────────────────────────────── tok_span_to_ent_char = {} for ent in entidades_finais: s_tok = char_to_token.get(ent["start"]) e_tok = char_to_token.get(min(ent["end"] - 1, len(text) - 1)) if s_tok is not None and e_tok is not None: tok_span_to_ent_char[(s_tok, e_tok)] = ent["start"] char_start_to_cluster_id = {} cluster_id_counter = defaultdict(int) for cluster in clusters: cluster_label = None cluster_char_starts = [] for (cs, ce) in cluster: char_s = tok_span_to_ent_char.get((cs, ce)) if char_s is not None: cluster_char_starts.append(char_s) if cluster_label is None: for ent in entidades_finais: if ent["start"] == char_s: cluster_label = ent["label"].replace("PERSONAL-", "") if ent["label"] else "Other" break if not cluster_label or not cluster_char_starts: continue cluster_id_counter[cluster_label] += 1 cid = cluster_id_counter[cluster_label] for char_s in cluster_char_starts: char_start_to_cluster_id[char_s] = cid entity_id_counters = defaultdict(int) text_to_id = defaultdict(dict) entidades_finais.sort(key=lambda x: x["start"]) for ent in entidades_finais: tag = ent["label"].replace("PERSONAL-", "") if ent["label"] else "Other" texto_ent = text[ent["start"]:ent["end"]].strip().lower() if ent["start"] in char_start_to_cluster_id: ent["entity_id"] = char_start_to_cluster_id[ent["start"]] elif texto_ent in text_to_id[tag]: ent["entity_id"] = text_to_id[tag][texto_ent] else: entity_id_counters[tag] += 1 new_id = entity_id_counters[tag] ent["entity_id"] = new_id text_to_id[tag][texto_ent] = new_id # ── PASSO 4: Geração do texto anonimizado ──────────────────────────── entidades_para_substituir = sorted(entidades_finais, key=lambda x: x["start"], reverse=True) texto_anon = text relatorio_json = [] for ent in entidades_para_substituir: tag_interno = ent["label"].replace("PERSONAL-", "") if ent["label"] else "Other" eid = ent.get("entity_id", "") id_part = f"-{eid}" if eid else "" texto_original = text[ent["start"]:ent["end"]].strip() relatorio_json.append({ "category": ent.get("label", "Other"), "text": texto_original, "start": ent["start"], "end": ent["end"], "id": eid, "tag_interno": tag_interno, }) tag_display = label_name(tag_interno) placeholder = f' <{tag_display}{id_part}> ' texto_anon = texto_anon[:ent["start"]] + placeholder + texto_anon[ent["end"]:] relatorio_json.reverse() return re.sub(r' +', ' ', texto_anon).strip(), relatorio_json # ───────────────────────────────────────────────────────────────────────────── # TEXTOS DE EXEMPLO # ───────────────────────────────────────────────────────────────────────────── @st.cache_data # Loads example texts from the JSON file and caches them. def load_example_texts(): json_path = os.path.join(os.path.dirname(__file__), 'example_text.json') try: with open(json_path, 'r', encoding='utf-8') as f: return json.load(f) except Exception: return { "Custom Text": "", "1º Portuguese Meeting Minute": "", "2º Portuguese Meeting Minute": "", "3º Portuguese Meeting Minute": "" } # ───────────────────────────────────────────────────────────────────────────── # INTERFACE PRINCIPAL # ───────────────────────────────────────────────────────────────────────────── # Builds and runs the main Streamlit application interface. def main(): if "lang_pt" not in st.session_state: st.session_state["lang_pt"] = False li = lang_idx() st.markdown( f'
' f'🛡️ PID: {UI["page_title"][li]}
', unsafe_allow_html=True ) st.markdown( f'{UI["page_subtitle"][li]}
', unsafe_allow_html=True ) tokenizer_ner, model_ner, model_coref, tokenizer_coref, coref_cfg, error = load_models() if error: st.error(f"Erro ao carregar modelos: {error}") st.stop() st.sidebar.header(t("sidebar_header")) st.sidebar.markdown("---") st.sidebar.markdown(f"**🌐 {t('lang_toggle_label')}**") col_en, col_pt = st.sidebar.columns(2) with col_en: if st.button( "English (EN)", use_container_width=True, type="secondary" if st.session_state["lang_pt"] else "primary", key="btn_en" ): st.session_state["lang_pt"] = False st.rerun() with col_pt: if st.button( "Portuguese (PT)", use_container_width=True, type="primary" if st.session_state["lang_pt"] else "secondary", key="btn_pt" ): st.session_state["lang_pt"] = True st.rerun() st.sidebar.markdown("---") example_texts = load_example_texts() name_map = UI["example_names"][lang_idx()] options_keys = list(example_texts.keys()) options_display = [name_map.get(k, k) for k in options_keys] selected_display = st.sidebar.selectbox( t("choose_example"), options=options_display, index=0 ) display_to_key = {v: k for k, v in name_map.items()} selected_example = display_to_key.get(selected_display, selected_display) st.sidebar.markdown("{t("entities_title")}
{grid_html}