Update app.py

app.py

@@ -19,8 +19,9 @@ model.eval()
 
 # Your inference function (adapted from your own)
 def inference(text, model, tokenizer, max_length, id2label):
+    masked_text = mask_pii(text)
     inputs = tokenizer(
-        text,
+        masked_text,
         add_special_tokens=True,
         max_length=MAX_LENGTH,
         padding='max_length',
@@ -35,15 +36,40 @@ def inference(text, model, tokenizer, max_length, id2label):
     probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
     predicted_class_id = torch.argmax(probs).item()
     predicted_class = id2label[predicted_class_id]
+    # Step 4: Extract PII fields (back from masked text)
+    detected_pii = []
+    if "<full_name>" in masked_text:
+        detected_pii.append("full_name")
+    if "<email>" in masked_text:
+        detected_pii.append("email")
+    if "<phone_number>" in masked_text:
+        detected_pii.append("phone_number")
+    if "<dob>" in masked_text:
+        detected_pii.append("dob")
+    if "<aadhar_num>" in masked_text:
+        detected_pii.append("aadhar")
+    if "<credit_debit_no>" in masked_text:
+        detected_pii.append("credit/debit card")
+    if "<cvv_no>" in masked_text:
+        detected_pii.append("cvv")
+    if "<expiry_no>" in masked_text:
+        detected_pii.append("expiry date")
+
+    return {
+        "original_email": text,
+        "masked_email": masked_text,
+        "predicted_email_class": predicted_class,
+        "prediction_confidence": float(probs[predicted_class_id]),
+        "detected_pii": detected_pii
+    }
 
-    return f"Predicted: {predicted_class}", {id2label[i]: float(probs[i]) for i in range(len(probs))}
 
 # Gradio interface
 demo = gr.Interface(
     fn=lambda text: inference(text, model, tokenizer, MAX_LENGTH, id2label),
     inputs=gr.Textbox(lines=8, placeholder="Paste your email here..."),
-    outputs=["text", "label"],
-    title="Email Classifier (DeBERTa v2)",
+    outputs="json",
+    title="Email Classifier",
     description="Classifies emails into Incident, Request, Problem, or Change."
 )