Update app.py

app.py

@@ -1,94 +1,108 @@
+import gradio as gr
 import os
-import argparse
 from models.pii_masker import PIIMasker
 from models.classifier import EmailClassifier
-from utils.utils import create_sample_dataset, parse_emails_dataset
-
-
-def train_model(
-    data_path=None,
-    model_output_path="models/email_classifier.joblib"
-):
-    """
-    Train the email classification model.
-
-    Args:
-        data_path: Path to the email dataset
-        model_output_path: Path to save the trained model
-    """
-    # Create sample dataset if data_path is not provided
-    if not data_path or not os.path.exists(data_path):
-        print("No dataset provided or file not found.")
-        data_path = "data/emails.csv"
+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)
 
-    # Parse and prepare dataset
+    # Train model
     df = parse_emails_dataset(data_path)
     X = df['email'].tolist()
     y = df['type'].tolist()
 
-    # Initialize and train classifier
     classifier = EmailClassifier()
-    report = classifier.train(X, y)
+    classifier.train(X, y)
 
     # Save trained model
-    os.makedirs(os.path.dirname(model_output_path), exist_ok=True)
-    classifier.save_model(model_output_path)
+    os.makedirs(os.path.dirname(model_path), exist_ok=True)
+    classifier.save_model(model_path)
+    print("Model trained and saved successfully!")
+
 
-    return report
+# Initialize components
+pii_masker = PIIMasker()
+classifier = EmailClassifier(model_path=model_path)
 
 
-def test_pii_masking():
+def process_email(email_body):
     """
-    Test the PII masking functionality.
+    Process email by masking PII and classifying it.
+        Args:
+        email_body: Raw email text
+    Returns:
+        tuple: (masked_email, entities_text, category)
     """
-    masker = PIIMasker()
+    # 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
+    entities_text = "\n".join([
+        f"- {entity['classification']}: {entity['entity']}"
+        for entity in entities
+    ])
+
+    return masked_email, entities_text, category
+
 
-    # Test email with PII
+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."
     )
-
-    masked_email, entities = masker.mask_pii(test_email)
-
-    print("Original Email:")
-    print(test_email)
-    print("\nMasked Email:")
-    print(masked_email)
-    print("\nDetected Entities:")
-    for entity in entities:
-        print(f"- {entity['classification']}: {entity['entity']}")
-
-
-def main():
-    """
-    Main function to run the email classification system.
+    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
     """
-    parser = argparse.ArgumentParser(description="Email Classification System")
-    parser.add_argument("--train", action="store_true",
-                        help="Train the classification model")
-    parser.add_argument("--test-masking", action="store_true",
-                        help="Test PII masking functionality")
-    parser.add_argument("--data-path", type=str,
-                        help="Path to email dataset for training")
-
-    args = parser.parse_args()
-
-    if args.train:
-        print("Training classification model...")
-        train_model(args.data_path)
-
-    if args.test_masking:
-        print("Testing PII masking functionality...")
-        test_pii_masking()
-
-    if not args.train and not args.test_masking:
-        print("No action specified. Use --train to train the model or "
-              "--test-masking to test PII masking.")
-
+)
 
-if __name__ == "__main__":
-    main()
+# Launch the app
+demo.launch()