Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForTokenClassification | |
| import time | |
| import json | |
| import pandas as pd | |
| from datetime import datetime | |
| import os | |
| from typing import List, Dict, Tuple | |
| import re | |
| MODEL_PATH = "CordwainerSmith/GolemPII-v1" | |
| ENTITY_COLORS = { | |
| "PHONE_NUM": "#FF9999", | |
| "ID_NUM": "#99FF99", | |
| "CC_NUM": "#9999FF", | |
| "BANK_ACCOUNT_NUM": "#FFFF99", | |
| "FIRST_NAME": "#FF99FF", | |
| "LAST_NAME": "#99FFFF", | |
| "CITY": "#FFB366", | |
| "STREET": "#B366FF", | |
| "POSTAL_CODE": "#66FFB3", | |
| "EMAIL": "#66B3FF", | |
| "DATE": "#FFB3B3", | |
| "CC_PROVIDER": "#B3FFB3", | |
| } | |
| EXAMPLE_SENTENCES = [ | |
| "שם מלא: תלמה אריאלי מספר תעודת זהות: 61453324-8 תאריך לידה: 15/09/1983 כתובת: ארלוזורוב 22 פתח תקווה מיקוד 2731711 אימייל: [email protected] טלפון: 054-8884771 בפגישה זו נדונו פתרונות טכנולוגיים חדשניים לשיפור תהליכי עבודה. המשתתף יתבקש להציג מצגת בנושא בפגישה הבאה אשר שילם ב 5326-1003-5299-5478 מסטרקארד עם הוראת קבע ל 11-77-352300", | |
| ] | |
| MODEL_DETAILS = { | |
| "name": "GolemPII-v1: Hebrew PII Detection Model", | |
| "description": 'The <a href="https://huggingface.co/CordwainerSmith/GolemPII-v1" target="_blank">GolemPII model</a> was specifically designed to identify and categorize various types of personally identifiable information (PII) present in Hebrew text. Its core intended usage revolves around enhancing privacy protection and facilitating the process of data anonymization. This makes it a good candidate for applications and systems that handle sensitive data, such as legal documents, medical records, or any text data containing PII, where the automatic redaction or removal of such information is essential for ensuring compliance with data privacy regulations and safeguarding individuals\' personal information. The model can be deployed on-premise with a relatively small hardware footprint, making it suitable for organizations with limited computing resources or those prioritizing local data processing.\n\nThe model was trained on the <a href="https://huggingface.co/datasets/CordwainerSmith/GolemGuard" target="_blank">GolemGuard</a> dataset, a Hebrew language dataset comprising over 115,000 examples of PII entities and containing both real and synthetically generated text examples. This data represents various document types and communication formats commonly found in Israeli professional and administrative contexts. GolemGuard covers a wide range of document types and encompasses a diverse array of PII entities, making it ideal for training and evaluating PII detection models.', | |
| "base_model": "xlm-roberta-base", | |
| "training_data": "Custom Hebrew PII dataset", | |
| "detected_pii_entities": [ | |
| "FIRST_NAME", | |
| "LAST_NAME", | |
| "STREET", | |
| "CITY", | |
| "PHONE_NUM", | |
| "EMAIL", | |
| "ID_NUM", | |
| "BANK_ACCOUNT_NUM", | |
| "CC_NUM", | |
| "CC_PROVIDER", | |
| "DATE", | |
| "POSTAL_CODE", | |
| ], | |
| } | |
| class PIIMaskingModel: | |
| def __init__(self, model_name: str): | |
| self.model_name = model_name | |
| hf_token = st.secrets["hf_token"] | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| model_name, use_auth_token=hf_token | |
| ) | |
| self.model = AutoModelForTokenClassification.from_pretrained( | |
| model_name, use_auth_token=hf_token | |
| ) | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.model.to(self.device) | |
| self.model.eval() | |
| def process_text( | |
| self, text: str | |
| ) -> Tuple[str, float, str, List[str], List[str], List[Dict]]: | |
| start_time = time.time() | |
| tokenized_inputs = self.tokenizer( | |
| text, | |
| truncation=True, | |
| padding=False, | |
| return_tensors="pt", | |
| return_offsets_mapping=True, | |
| add_special_tokens=True, | |
| ) | |
| input_ids = tokenized_inputs.input_ids.to(self.device) | |
| attention_mask = tokenized_inputs.attention_mask.to(self.device) | |
| offset_mapping = tokenized_inputs["offset_mapping"][0].tolist() | |
| # Handle special tokens | |
| offset_mapping[0] = None # <s> token | |
| offset_mapping[-1] = None # </s> token | |
| with torch.no_grad(): | |
| outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) | |
| predictions = outputs.logits.argmax(dim=-1).cpu().numpy() | |
| predicted_labels = [ | |
| self.model.config.id2label[label_id] for label_id in predictions[0] | |
| ] | |
| tokens = self.tokenizer.convert_ids_to_tokens(input_ids[0]) | |
| masked_text, colored_text, privacy_masks = self.mask_pii_in_sentence( | |
| tokens, predicted_labels, text, offset_mapping | |
| ) | |
| processing_time = time.time() - start_time | |
| return ( | |
| masked_text, | |
| processing_time, | |
| colored_text, | |
| tokens, | |
| predicted_labels, | |
| privacy_masks, | |
| ) | |
| def _find_entity_span( | |
| self, | |
| i: int, | |
| labels: List[str], | |
| tokens: List[str], | |
| offset_mapping: List[Tuple[int, int]], | |
| ) -> Tuple[int, str, int]: | |
| current_entity = labels[i][2:] if labels[i].startswith("B-") else labels[i][2:] | |
| j = i + 1 | |
| last_valid_end = offset_mapping[i][1] if offset_mapping[i] else None | |
| while j < len(tokens): | |
| if offset_mapping[j] is None: | |
| j += 1 | |
| continue | |
| next_label = labels[j] | |
| if next_label.startswith("B-") and tokens[j].startswith("▁"): | |
| break | |
| if next_label.startswith("I-") and next_label[2:] != current_entity: | |
| break | |
| if next_label.startswith("I-") and next_label[2:] == current_entity: | |
| last_valid_end = offset_mapping[j][1] | |
| j += 1 | |
| elif next_label.startswith("B-") and not tokens[j].startswith("▁"): | |
| last_valid_end = offset_mapping[j][1] | |
| j += 1 | |
| else: | |
| break | |
| return j, current_entity, last_valid_end | |
| def mask_pii_in_sentence( | |
| self, | |
| tokens: List[str], | |
| labels: List[str], | |
| original_text: str, | |
| offset_mapping: List[Tuple[int, int]], | |
| ) -> Tuple[str, str, List[Dict]]: | |
| privacy_masks = [] | |
| current_pos = 0 | |
| masked_text_parts = [] | |
| colored_text_parts = [] | |
| i = 0 | |
| while i < len(tokens): | |
| if offset_mapping[i] is None: | |
| i += 1 | |
| continue | |
| current_label = labels[i] | |
| if current_label.startswith(("B-", "I-")): | |
| start_char = offset_mapping[i][0] | |
| next_pos, entity_type, last_valid_end = self._find_entity_span( | |
| i, labels, tokens, offset_mapping | |
| ) | |
| if current_pos < start_char: | |
| text_before = original_text[current_pos:start_char] | |
| masked_text_parts.append(text_before) | |
| colored_text_parts.append(text_before) | |
| entity_value = original_text[start_char:last_valid_end] | |
| mask = self._get_mask_for_entity(entity_type) | |
| privacy_masks.append( | |
| { | |
| "label": entity_type, | |
| "start": start_char, | |
| "end": last_valid_end, | |
| "value": entity_value, | |
| "label_index": len(privacy_masks) + 1, | |
| } | |
| ) | |
| masked_text_parts.append(mask) | |
| color = ENTITY_COLORS.get(entity_type, "#CCCCCC") | |
| colored_text_parts.append( | |
| f'<span style="background-color: {color}; color: black; padding: 2px; border-radius: 3px;">{mask}</span>' | |
| ) | |
| current_pos = last_valid_end | |
| i = next_pos | |
| else: | |
| if offset_mapping[i] is not None: | |
| start_char = offset_mapping[i][0] | |
| end_char = offset_mapping[i][1] | |
| if current_pos < end_char: | |
| text_chunk = original_text[current_pos:end_char] | |
| masked_text_parts.append(text_chunk) | |
| colored_text_parts.append(text_chunk) | |
| current_pos = end_char | |
| i += 1 | |
| if current_pos < len(original_text): | |
| remaining_text = original_text[current_pos:] | |
| masked_text_parts.append(remaining_text) | |
| colored_text_parts.append(remaining_text) | |
| return ("".join(masked_text_parts), "".join(colored_text_parts), privacy_masks) | |
| def _get_mask_for_entity(self, entity_type: str) -> str: | |
| return { | |
| "PHONE_NUM": "[טלפון]", | |
| "ID_NUM": "[ת.ז]", | |
| "CC_NUM": "[כרטיס אשראי]", | |
| "BANK_ACCOUNT_NUM": "[חשבון בנק]", | |
| "FIRST_NAME": "[שם פרטי]", | |
| "LAST_NAME": "[שם משפחה]", | |
| "CITY": "[עיר]", | |
| "STREET": "[רחוב]", | |
| "POSTAL_CODE": "[מיקוד]", | |
| "EMAIL": "[אימייל]", | |
| "DATE": "[תאריך]", | |
| "CC_PROVIDER": "[ספק כרטיס אשראי]", | |
| "BANK": "[בנק]", | |
| }.get(entity_type, f"[{entity_type}]") | |
| def main(): | |
| st.set_page_config(layout="wide") | |
| st.title("🗿 GolemPII: Hebrew PII Masking Application 🗿") | |
| st.markdown( | |
| """ | |
| <style> | |
| .rtl { direction: rtl; text-align: right; } | |
| .entity-legend { padding: 5px; margin: 2px; border-radius: 3px; display: inline-block; } | |
| .masked-text { | |
| direction: rtl; | |
| text-align: right; | |
| line-height: 2; | |
| padding: 10px; | |
| background-color: #f6f8fa; | |
| border-radius: 5px; | |
| color: black; | |
| white-space: pre-wrap; | |
| } | |
| .main h3 { | |
| margin-bottom: 10px; | |
| } | |
| textarea { | |
| direction: rtl !important; | |
| text-align: right !important; | |
| } | |
| .stTextArea label { | |
| direction: ltr !important; | |
| text-align: left !important; | |
| } | |
| </style> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| # Sidebar with model details | |
| st.sidebar.markdown( | |
| f""" | |
| <div> | |
| <h2>{MODEL_DETAILS['name']}</h2> | |
| <p>{MODEL_DETAILS['description']}</p> | |
| <h3>Supported PII Entities</h3> | |
| <ul> | |
| {" ".join([f'<li><span style="background-color: {ENTITY_COLORS.get(entity, "#CCCCCC")}; color: black; padding: 3px 5px; border-radius: 3px; margin-right: 5px;">{entity}</span></li>' for entity in MODEL_DETAILS['detected_pii_entities']])} | |
| </ul> | |
| </div> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| text_input = st.text_area( | |
| "Enter text to mask (separate multiple texts with commas):", | |
| value="\n".join(EXAMPLE_SENTENCES), | |
| height=200, | |
| ) | |
| show_json = st.checkbox("Show JSON Output", value=True) | |
| if st.button("Process Text"): | |
| texts = [text.strip() for text in text_input.split(",") if text.strip()] | |
| model = PIIMaskingModel() | |
| for text in texts: | |
| st.markdown( | |
| '<h3 style="text-align: center;">Original Text</h3>', | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown(f'<div class="rtl">{text}</div>', unsafe_allow_html=True) | |
| ( | |
| masked_text, | |
| processing_time, | |
| colored_text, | |
| tokens, | |
| predicted_labels, | |
| privacy_masks, | |
| ) = model.process_text(text) | |
| st.markdown( | |
| '<h3 style="text-align: center;">Masked Text</h3>', | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown( | |
| f'<div class="masked-text">{colored_text}</div>', unsafe_allow_html=True | |
| ) | |
| st.markdown(f"Processing Time: {processing_time:.3f} seconds") | |
| if show_json: | |
| st.json( | |
| { | |
| "original": text, | |
| "masked": masked_text, | |
| "processing_time": processing_time, | |
| "tokens": tokens, | |
| "token_classes": predicted_labels, | |
| "privacy_mask": privacy_masks, | |
| "span_labels": [ | |
| [m["start"], m["end"], m["label"]] for m in privacy_masks | |
| ], | |
| } | |
| ) | |
| st.markdown("---") | |
| if __name__ == "__main__": | |
| main() | |