v3

app.py

@@ -1,15 +1,17 @@
+"""
+This module demonstrates a Streamlit application for masking Personally Identifiable
+Information (PII) in Hebrew text using the GolemPII-v1 model.
+"""
+
+import time
+from typing import List, Dict, Tuple
+
 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
 
 
+# Constants for model name and entity colors
 MODEL_NAME = "CordwainerSmith/GolemPII-v1"
 
 ENTITY_COLORS = {
@@ -27,13 +29,36 @@ ENTITY_COLORS = {
     "CC_PROVIDER": "#B3FFB3",
 }
 
+# Example sentences for demonstration
 EXAMPLE_SENTENCES = [
     "שם מלא: תלמה אריאלי מספר תעודת זהות: 61453324-8 תאריך לידה: 15/09/1983 כתובת:  ארלוזורוב 22  פתח תקווה מיקוד 2731711 אימייל: [email protected] טלפון: 054-8884771  בפגישה זו נדונו פתרונות טכנולוגיים חדשניים לשיפור תהליכי עבודה. המשתתף יתבקש להציג מצגת בנושא בפגישה הבאה אשר שילם ב 5326-1003-5299-5478 מסטרקארד עם הוראת קבע ל 11-77-352300",
 ]
 
+# Model details for display in the sidebar
 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.',
+    "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.
+
+        The 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": [
@@ -54,12 +79,20 @@ MODEL_DETAILS = {
 
 
 class PIIMaskingModel:
+    """
+    A class for masking PII in Hebrew text using the GolemPII-v1 model.
+    """
+
     def __init__(self, model_name: str):
+        """
+        Initializes the PIIMaskingModel with the specified model name.
+
+        Args:
+            model_name: The name of the pre-trained model to use.
+        """
         self.model_name = model_name
         hf_token = st.secrets["hf_token"]
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            model_name, token=hf_token
-        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
         self.model = AutoModelForTokenClassification.from_pretrained(
             model_name, token=hf_token
         )
@@ -70,6 +103,23 @@ class PIIMaskingModel:
     def process_text(
         self, text: str
     ) -> Tuple[str, float, str, List[str], List[str], List[Dict]]:
+        """
+        Processes the input text and returns the masked text, processing time,
+        colored text, tokens, predicted labels, and privacy masks.
+
+        Args:
+            text: The input text to process.
+
+        Returns:
+            A tuple containing:
+                - masked_text: The text with PII masked.
+                - processing_time: The time taken to process the text.
+                - colored_text: The text with PII highlighted with colors.
+                - tokens: The tokens of the input text.
+                - predicted_labels: The predicted labels for each token.
+                - privacy_masks: A list of dictionaries containing information about
+                  the masked PII entities.
+        """
         start_time = time.time()
 
         tokenized_inputs = self.tokenizer(
@@ -119,6 +169,21 @@ class PIIMaskingModel:
         tokens: List[str],
         offset_mapping: List[Tuple[int, int]],
     ) -> Tuple[int, str, int]:
+        """
+        Finds the span of an entity starting at the given index.
+
+        Args:
+            i: The starting index of the entity.
+            labels: The list of labels for each token.
+            tokens: The list of tokens.
+            offset_mapping: The offset mapping for each token.
+
+        Returns:
+            A tuple containing:
+                - The index of the next token after the entity.
+                - The type of the entity.
+                - The end character offset of the entity.
+        """
         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
@@ -130,7 +195,7 @@ class PIIMaskingModel:
 
             next_label = labels[j]
 
-            if next_label.startswith("B-") and tokens[j].startswith("▁"):
+            if next_label.startswith("B-") and tokens[j].startswith(" "):
                 break
 
             if next_label.startswith("I-") and next_label[2:] != current_entity:
@@ -139,7 +204,7 @@ class PIIMaskingModel:
             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("▁"):
+            elif next_label.startswith("B-") and not tokens[j].startswith(" "):
                 last_valid_end = offset_mapping[j][1]
                 j += 1
             else:
@@ -154,6 +219,22 @@ class PIIMaskingModel:
         original_text: str,
         offset_mapping: List[Tuple[int, int]],
     ) -> Tuple[str, str, List[Dict]]:
+        """
+        Masks the PII entities in a sentence.
+
+        Args:
+            tokens: The list of tokens in the sentence.
+            labels: The list of labels for each token.
+            original_text: The original text of the sentence.
+            offset_mapping: The offset mapping for each token.
+
+        Returns:
+            A tuple containing:
+                - The masked text.
+                - The colored text.
+                - A list of dictionaries containing information about the masked
+                  PII entities.
+        """
         privacy_masks = []
         current_pos = 0
         masked_text_parts = []
@@ -219,6 +300,15 @@ class PIIMaskingModel:
         return ("".join(masked_text_parts), "".join(colored_text_parts), privacy_masks)
 
     def _get_mask_for_entity(self, entity_type: str) -> str:
+        """
+        Returns the mask for a given entity type.
+
+        Args:
+            entity_type: The type of the entity.
+
+        Returns:
+            The mask for the entity type.
+        """
         return {
             "PHONE_NUM": "[טלפון]",
             "ID_NUM": "[ת.ז]",
@@ -237,6 +327,9 @@ class PIIMaskingModel:
 
 
 def main():
+    """
+    The main function for the Streamlit application.
+    """
     st.set_page_config(layout="wide")
     st.title("🗿 GolemPII: Hebrew PII Masking Application 🗿")
 
@@ -245,17 +338,17 @@ def main():
     <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;             
+        .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 {            
+        .main h3 {
             margin-bottom: 10px;
         }
         textarea {