Update personal_info_identifier.py

personal_info_identifier.py

@@ -1,16 +1,45 @@
 import gradio as gr
 from transformers import pipeline
+import re
 
-# โหลด pipeline
-pii_detector = pipeline("token-classification", model="iiiorg/piiranha-v1-detect-personal-information", aggregation_strategy="simple")
+# โหลดโมเดล PII Detection
+pii_detector = pipeline(
+    "token-classification",
+    model="iiiorg/piiranha-v1-detect-personal-information",
+    aggregation_strategy="simple"
+)
 
-def detect_pii(text):
-    results = pii_detector(text)
-    return results
+def highlight_pii(text):
+    entities = pii_detector(text)
+    highlighted_text = text
+    offset = 0  # For shifting positions when inserting tags
+
+    for entity in sorted(entities, key=lambda x: x['start']):
+        start = entity['start'] + offset
+        end = entity['end'] + offset
+        label = entity['entity_group']
+
+        # wrap entity with a span tag with color (red for demo, can be better)
+        span = f'<span style="background-color:#ffcccc; padding:2px 4px; border-radius:4px;">{text[entity["start"]:entity["end"]]} <b style="color:red;">[{label}]</b></span>'
+        highlighted_text = highlighted_text[:start] + span + highlighted_text[end:]
+        offset += len(span) - (end - start)
+
+    return highlighted_text
 
 def create_personal_info_tab():
     with gr.Column():
-        gr.Markdown("### Detect Personal Information (PII)")
-        textbox = gr.Textbox(label="Enter text", lines=5)
-        output = gr.JSON(label="Detected Entities")
-        textbox.change(fn=detect_pii, inputs=textbox, outputs=output)
+        gr.Markdown("### Personal Information Identifier")
+        input_text = gr.Textbox(label="Enter text", lines=5, placeholder="Type your text here...")
+        output_html = gr.HTML()
+
+        submit_button = gr.Button("Submit")
+
+        submit_button.click(fn=highlight_pii, inputs=input_text, outputs=output_html)
+
+        # Optional: Add example
+        gr.Examples(
+            examples=[
+                "Hugging Face is a company based in Paris and New York City that acquired Gradio in 2021."
+            ],
+            inputs=input_text
+        )