File size: 3,259 Bytes
ace716e ab3c5c9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e 14d24e9 ace716e 14d24e9 ace716e ab3c5c9 ace716e 14d24e9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e 14d24e9 ace716e ab3c5c9 ace716e ab3c5c9 ace716e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import gradio as gr
import os
from models.pii_masker import PIIMasker
from models.classifier import EmailClassifier
from utils.utils import (preprocess_email, create_sample_dataset,
parse_emails_dataset)
# Check if model exists, if not train it
model_path = "models/email_classifier.joblib"
if not os.path.exists(model_path):
print("Training model as it doesn't exist yet...")
data_path = "data/emails.csv"
# Create sample dataset if needed
if not os.path.exists(data_path):
print("Creating sample dataset...")
create_sample_dataset(data_path)
# Train model
df = parse_emails_dataset(data_path)
X = df['email'].tolist()
y = df['type'].tolist()
classifier = EmailClassifier()
classifier.train(X, y)
# Save trained model
os.makedirs(os.path.dirname(model_path), exist_ok=True)
classifier.save_model(model_path)
print("Model trained and saved successfully!")
# Initialize components
pii_masker = PIIMasker()
classifier = EmailClassifier(model_path=model_path)
def process_email(email_body):
"""
Process email by masking PII and classifying it.
Args:
email_body: Raw email text
Returns:
tuple: (masked_email, entities_text, category)
"""
# Mask PII
masked_email, entities = pii_masker.mask_pii(email_body)
# Preprocess for classification
processed_email = preprocess_email(masked_email)
# Classify email
category = classifier.classify(processed_email)
# Format entities for display - fixed formatting
entities_text = ""
for entity in entities:
entities_text += f"- {entity['classification']}: {entity['entity']}\n"
return masked_email, entities_text, category
def test_masking():
"""Example function to demonstrate PII masking"""
test_email = (
"Hello, my name is John Doe, and my email is johndoe@example.com.\n"
"My phone number is 555-123-4567 and I was born on 15/04/1985.\n"
"My Aadhar number is 1234 5678 9012 and my credit card number is "
"4111 1111 1111 1111 with CVV 123 expiring on 12/25."
)
return test_email
# Create Gradio interface
demo = gr.Interface(
fn=process_email,
inputs=gr.Textbox(
lines=10,
label="Email Content",
placeholder="Enter email text to classify and mask PII..."
),
outputs=[
gr.Textbox(label="Masked Email"),
gr.Textbox(label="Detected PII Entities"),
gr.Textbox(label="Email Category")
],
title="Email Classification System",
description=(
"This application classifies support emails and masks personally "
"identifiable information (PII)."
),
examples=[
["Hello, my name is John Doe, and my email is johndoe@example.com. "
"I need help with my account."],
["I'm having trouble logging in to my account. My username is user"
"123."]
],
article="""
## How It Works
1. **PII Masking**: The system identifies and masks personal information
2. **Email Classification**: The masked email is classified into categories
3. **Results**: View the masked version, detected PII, and email category
"""
)
# Launch the app
demo.launch()
|