v2

app.py

@@ -9,11 +9,8 @@ import os
 from typing import List, Dict, Tuple
 import re
 
-# Constants
-MODELS = {
-    "GolemPII XLM-RoBERTa v1": "CordwainerSmith/GolemPII-xlm-roberta-v1",
-}
 
+MODEL_PATH = "CordwainerSmith/GolemPII-v1"
 
 ENTITY_COLORS = {
     "PHONE_NUM": "#FF9999",
@@ -35,8 +32,8 @@ EXAMPLE_SENTENCES = [
 ]
 
 MODEL_DETAILS = {
-    "name": "GolemPII-xlm-roberta-v1 - Hebrew PII Detection Model",
-    "description": "This model is specifically designed to identify and categorize Personally Identifiable Information (PII) within Hebrew text. It leverages the powerful XLM-RoBERTa base, fine-tuned with a curated Hebrew PII dataset, making it adept at token classification tasks tailored for Hebrew.",
+    "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": [
@@ -53,17 +50,6 @@ MODEL_DETAILS = {
         "DATE",
         "POSTAL_CODE",
     ],
-    "performance_metrics": {
-        "Loss": 0.000729,
-        "Precision": 0.9982,
-        "Recall": 0.9982,
-        "F1-Score": 0.9982,
-        "Accuracy": 0.999795,
-    },
-    "training_details": {
-        "Training language": "Hebrew",
-        # Add other relevant training details if available
-    },
 }
 
 
@@ -133,7 +119,6 @@ class PIIMaskingModel:
         tokens: List[str],
         offset_mapping: List[Tuple[int, int]],
     ) -> Tuple[int, str, int]:
-        """Find the end index and entity type for a span starting at index i"""
         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
@@ -145,19 +130,15 @@ class PIIMaskingModel:
 
             next_label = labels[j]
 
-            # Stop if we hit a new B- tag (except for non-spaced tokens)
             if next_label.startswith("B-") and tokens[j].startswith("▁"):
                 break
 
-            # Stop if we hit a different entity type in I- tags
             if next_label.startswith("I-") and next_label[2:] != current_entity:
                 break
 
-            # Continue if it's a continuation of the same entity
             if next_label.startswith("I-") and next_label[2:] == current_entity:
                 last_valid_end = offset_mapping[j][1]
                 j += 1
-            # Continue if it's a non-spaced B- token
             elif next_label.startswith("B-") and not tokens[j].startswith("▁"):
                 last_valid_end = offset_mapping[j][1]
                 j += 1
@@ -180,7 +161,7 @@ class PIIMaskingModel:
 
         i = 0
         while i < len(tokens):
-            if offset_mapping[i] is None:  # Skip special tokens
+            if offset_mapping[i] is None:
                 i += 1
                 continue
 
@@ -188,23 +169,18 @@ class PIIMaskingModel:
 
             if current_label.startswith(("B-", "I-")):
                 start_char = offset_mapping[i][0]
-
-                # Find the complete entity span
                 next_pos, entity_type, last_valid_end = self._find_entity_span(
                     i, labels, tokens, offset_mapping
                 )
 
-                # Add any text before the entity
                 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)
 
-                # Extract and mask the entity
                 entity_value = original_text[start_char:last_valid_end]
                 mask = self._get_mask_for_entity(entity_type)
 
-                # Add to privacy masks
                 privacy_masks.append(
                     {
                         "label": entity_type,
@@ -215,13 +191,10 @@ class PIIMaskingModel:
                     }
                 )
 
-                # Add masked text
                 masked_text_parts.append(mask)
-
-                # Add colored text
                 color = ENTITY_COLORS.get(entity_type, "#CCCCCC")
                 colored_text_parts.append(
-                    f'<span style="background-color: {color}; padding: 2px; border-radius: 3px;">{mask}</span>'
+                    f'<span style="background-color: {color}; color: black; padding: 2px; border-radius: 3px;">{mask}</span>'
                 )
 
                 current_pos = last_valid_end
@@ -231,7 +204,6 @@ class PIIMaskingModel:
                     start_char = offset_mapping[i][0]
                     end_char = offset_mapping[i][1]
 
-                    # Add any text for this token
                     if current_pos < end_char:
                         text_chunk = original_text[current_pos:end_char]
                         masked_text_parts.append(text_chunk)
@@ -239,7 +211,6 @@ class PIIMaskingModel:
                         current_pos = end_char
                 i += 1
 
-        # Add any remaining text
         if current_pos < len(original_text):
             remaining_text = original_text[current_pos:]
             masked_text_parts.append(remaining_text)
@@ -248,7 +219,6 @@ class PIIMaskingModel:
         return ("".join(masked_text_parts), "".join(colored_text_parts), privacy_masks)
 
     def _get_mask_for_entity(self, entity_type: str) -> str:
-        """Get the mask text for a given entity type"""
         return {
             "PHONE_NUM": "[טלפון]",
             "ID_NUM": "[ת.ז]",
@@ -266,24 +236,10 @@ class PIIMaskingModel:
         }.get(entity_type, f"[{entity_type}]")
 
 
-def save_results_to_file(results: Dict):
-    """
-    Save processing results to a JSON file
-    """
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    filename = f"pii_masking_results_{timestamp}.json"
-
-    with open(filename, "w", encoding="utf-8") as f:
-        json.dump(results, f, ensure_ascii=False, indent=2)
-
-    return filename
-
-
 def main():
     st.set_page_config(layout="wide")
     st.title("🗿 GolemPII: Hebrew PII Masking Application 🗿")
 
-    # Add CSS styles
     st.markdown(
         """
     <style>
@@ -299,155 +255,91 @@ def main():
             color: black;             
             white-space: pre-wrap;
         }
-        /* Red headers for sections */
-        .main h3 {
-            color: #d73a49;
+        .main h3 {            
             margin-bottom: 10px;
         }
-        /* Styles for the model details sidebar */
-        .model-details-sidebar h2 {
-            margin-top: 0;
-        }
-        .model-details-sidebar table {
-            width: 100%;
-            border-collapse: collapse;
+        textarea {
+            direction: rtl !important;
+            text-align: right !important;
         }
-        .model-details-sidebar td, .model-details-sidebar th {
-            padding: 8px;
-            border: 1px solid #ddd;
-            text-align: left;
+        .stTextArea label {
+            direction: ltr !important;
+            text-align: left !important;
         }
     </style>
     """,
         unsafe_allow_html=True,
     )
 
-    # Sidebar configuration
-    st.sidebar.header("Configuration")
-    selected_model = st.sidebar.selectbox("Select Model", list(MODELS.keys()))
-    show_json = st.sidebar.checkbox("Show JSON Output", value=True)
-    run_all_models = st.sidebar.checkbox("Run All Models")
-
-    # Display Model Details in Sidebar
+    # Sidebar with model details
     st.sidebar.markdown(
         f"""
-        <div class="model-details-sidebar">
-            <h2>Model Details: {MODEL_DETAILS['name']}</h2>
+        <div>
+            <h2>{MODEL_DETAILS['name']}</h2>
             <p>{MODEL_DETAILS['description']}</p>
-            <table>
-                <tr><td>Base Model:</td><td>{MODEL_DETAILS['base_model']}</td></tr>
-                <tr><td>Training Data:</td><td>{MODEL_DETAILS['training_data']}</td></tr>            
-            </table>
-            <h3>Detected PII Entities</h3>
+            <h3>Supported PII Entities</h3>
             <ul>
-                {" ".join([f'<li><span class="entity-badge" style="background-color: {ENTITY_COLORS.get(entity, "#CCCCCC")}; padding: 3px 5px; border-radius: 3px; margin-right: 5px;">{entity}</span></li>' for entity in MODEL_DETAILS['detected_pii_entities']])}
+                {" ".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
     text_input = st.text_area(
         "Enter text to mask (separate multiple texts with commas):",
         value="\n".join(EXAMPLE_SENTENCES),
         height=200,
     )
 
-    # Process button
+    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()]
-
-        if run_all_models:
-            all_results = {}
-            progress_bar = st.progress(0)
-
-            for idx, (model_name, model_path) in enumerate(MODELS.items()):
-                st.subheader(f"Results for {model_name}")
-                model = PIIMaskingModel(model_path)
-                model_results = {}
-
-                for text_idx, text in enumerate(texts):
-                    (
-                        masked_text,
-                        processing_time,
-                        colored_text,
-                        tokens,
-                        predicted_labels,
-                        privacy_masks,
-                    ) = model.process_text(text)
-                    model_results[f"text_{text_idx+1}"] = {
+        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
                         ],
                     }
-
-                all_results[model_name] = model_results
-                progress_bar.progress((idx + 1) / len(MODELS))
-
-            # Save and display results
-            filename = save_results_to_file(all_results)
-            st.success(f"Results saved to {filename}")
-
-            # Show comparison table
-            comparison_data = []
-            for model_name, results in all_results.items():
-                avg_time = sum(
-                    text_data["processing_time"] for text_data in results.values()
-                ) / len(results)
-                comparison_data.append(
-                    {"Model": model_name, "Avg Processing Time": f"{avg_time:.3f}s"}
-                )
-
-            st.subheader("Model Comparison")
-            st.table(pd.DataFrame(comparison_data))
-
-        else:
-            # Process with single selected model
-            model = PIIMaskingModel(MODELS[selected_model])
-
-            for text in texts:
-                st.markdown("### Original Text", 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("### Masked Text", 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("---")
+            st.markdown("---")
 
 
 if __name__ == "__main__":