Supreetha15/safe-email-classifier

1# test_model.py
2from models import SBERT_SVM_Classifier
3from utils import mask_pii, demask
4import json
5
6# Load trained model
7model = SBERT_SVM_Classifier()
8model.load("sbert_linear_model.joblib")
9
10
11def classify_email(raw_email: str) -> dict:
12    # 1. Mask PII
13    masked_email, entities = mask_pii(raw_email)
14
15    # 2. Predict category using masked email
16    predicted_category = model.predict([masked_email])[0]
17
18    # 3. Demask to restore original
19    demasked_email = demask(masked_email, entities)
20
21    # 4. Return formatted output
22    return {
23        "input_email_body": raw_email,
24        "list_of_masked_entities": entities,
25        "masked_email": masked_email,
26        "category_of_the_email": predicted_category
27    }
28
29
30# 🔬 Test sample
31if __name__ == "__main__":
32    email_text = "Subject: John Doe's Aadhar number is 1234-5678-9012 and email is johnd@example.com. Need access to billing portal."
33    result = classify_email(email_text)
34    print(json.dumps(result, indent=2))
35