Update api.py

api.py

@@ -1,80 +1,108 @@
-from fastapi import FastAPI, HTTPException, Body
-from pydantic import BaseModel
-import uvicorn
+import gradio as gr
+import os
 from models.pii_masker import PIIMasker
 from models.classifier import EmailClassifier
-from utils.utils import preprocess_email
+from utils.utils import (preprocess_email, create_sample_dataset,
+                         parse_emails_dataset)
 
-# Initialize FastAPI app
-app = FastAPI(title="Email Classification API",
-              description="API for classifying support emails and masking PII",
-              version="1.0.0")
+# 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"
 
-# Initialize PII masker and classifier
-pii_masker = PIIMasker()
-classifier = EmailClassifier(model_path="models/email_classifier.joblib")
-
-
-class EmailRequest(BaseModel):
-    email_body: str
+    # 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()
 
-class EmailResponse(BaseModel):
-    input_email_body: str
-    list_of_masked_entities: list
-    masked_email: str
-    category_of_the_email: str
+    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!")
 
-@app.post("/classify-email", response_model=EmailResponse)
-async def classify_email(request: EmailRequest = Body(...)):
-    """
-    Classify an email and mask PII information.
-
-    Args:
-        request: Email request object containing the email body
-
-    Returns:
-        dict: Response with masked email and classification
-    """
-    try:
-        email_body = request.email_body
-
-        # 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)
-
-        # Prepare response
-        response = {
-            "input_email_body": email_body,
-            "list_of_masked_entities": entities,
-            "masked_email": masked_email,
-            "category_of_the_email": category
-        }
-
-        return response
-
-    except Exception as e:
-        error_msg = f"Error processing request: {str(e)}"
-        raise HTTPException(status_code=500, detail=error_msg)
+# Initialize components
+pii_masker = PIIMasker()
+classifier = EmailClassifier(model_path=model_path)
 
 
-@app.get("/health")
-async def health_check():
+def process_email(email_body):
     """
-    Health check endpoint.
-
+    Process email by masking PII and classifying it.
+        Args:
+        email_body: Raw email text
     Returns:
-        dict: Status message
+        tuple: (masked_email, entities_text, category)
     """
-    return {"status": "healthy"}
-
+    # 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
+
+
+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
+    """
+)
 
-# For local development
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+# Launch the app
+demo.launch()