| import os |
| import gradio as gr |
| from models.pii_masker import PIIMasker |
| from models.classifier import EmailClassifier |
| |
| masker = PIIMasker() |
| classifier = EmailClassifier() |
|
|
| |
| MODEL_PATH = "models/email_classifier.joblib" |
| if os.path.exists(MODEL_PATH): |
| classifier.load_model(MODEL_PATH) |
| else: |
| print("Warning: Pre-trained model not found. " |
| "Classification will not work properly.") |
|
|
|
|
| def classify_email(email_text): |
| """ |
| Classify the given email text. |
| |
| Args: |
| email_text (str): The email text to classify |
| |
| Returns: |
| str: The classification result |
| """ |
| if not email_text.strip(): |
| return "Please enter an email text to classify." |
|
|
| result = classifier.predict([email_text]) |
| return f"Classification: {result[0]}" |
|
|
|
|
| def mask_pii(email_text): |
| """ |
| Mask personally identifiable information in the email text. |
| |
| Args: |
| email_text (str): The email text to mask PII from |
| |
| Returns: |
| tuple: (masked_email, entities_found) |
| """ |
| if not email_text.strip(): |
| return "Please enter an email text to mask PII.", "No entities found." |
|
|
| masked_email, entities = masker.mask_pii(email_text) |
|
|
| entities_text = "" |
| for entity in entities: |
| entities_text += f"- {entity['classification']}: {entity['entity']}\n" |
|
|
| if not entities: |
| entities_text = "No PII entities detected." |
|
|
| return masked_email, entities_text |
|
|
|
|
| def process_email(email_text, mask_pii_option=True): |
| """ |
| Process the email by optionally masking PII and classifying it. |
| |
| Args: |
| email_text (str): The email text to process |
| mask_pii_option (bool): Whether to mask PII before classification |
| |
| Returns: |
| tuple: (processed_email, entities_found, classification) |
| """ |
| if not email_text.strip(): |
| return ("Please enter an email text.", "No processing performed.", |
| "No classification performed.") |
|
|
| entities_text = "PII masking was not selected." |
| processed_email = email_text |
|
|
| if mask_pii_option: |
| processed_email, entities_text = mask_pii(email_text) |
|
|
| classification = classify_email(processed_email) |
|
|
| return processed_email, entities_text, classification |
|
|
|
|
| |
| with gr.Blocks(title="Email Classification System") as demo: |
| gr.Markdown("# Email Classification System") |
| gr.Markdown("This application classifies emails and can mask personally " |
| "identifiable information (PII).") |
|
|
| with gr.Tab("Process Email"): |
| with gr.Row(): |
| with gr.Column(): |
| input_email = gr.Textbox( |
| label="Input Email Text", |
| placeholder="Enter email text here...", |
| lines=10 |
| ) |
| mask_checkbox = gr.Checkbox( |
| label="Mask PII before classification", |
| value=True |
| ) |
| process_button = gr.Button("Process Email") |
|
|
| with gr.Column(): |
| output_email = gr.Textbox( |
| label="Processed Email", |
| lines=10, |
| interactive=False |
| ) |
| entities_found = gr.Textbox( |
| label="PII Entities Detected", |
| lines=5, |
| interactive=False |
| ) |
| classification_result = gr.Textbox( |
| label="Classification Result", |
| interactive=False |
| ) |
|
|
| with gr.Tab("About"): |
| gr.Markdown(""" |
| ## About This Application |
| |
| This application provides two main functionalities: |
| |
| 1. **PII Masking**: Detects and masks personally identifiable information in |
| emails |
| 2. **Email Classification**: Classifies emails into predefined categories |
| |
| The system can be used to preprocess emails for data privacy compliance |
| and to organize emails by type. |
| """) |
|
|
| process_button.click( |
| fn=process_email, |
| inputs=[input_email, mask_checkbox], |
| outputs=[output_email, entities_found, classification_result] |
| ) |
|
|
|
|
| |
| if __name__ == "__main__": |
| demo.launch() |
|
|
