Update app.py

app.py

@@ -3,25 +3,28 @@ from pydantic import BaseModel
 import joblib
 import re
 
+# Initialize FastAPI app
 app = FastAPI(
     title="Email Classification API",
     version="1.0.0",
     description="Classifies support emails into categories and masks personal information.",
-    docs_url="/docs",       # Swagger UI enabled here
-    redoc_url="/redoc"      # Optional ReDoc UI
+    docs_url="/docs",
+    redoc_url="/redoc"
 )
 
-# Load model
+# Load pre-trained model
 model = joblib.load("model.joblib")
 
-# Email input structure
+# Input schema
 class EmailInput(BaseModel):
-    subject: str = ""
-    email: str
+    input_email_body: str
 
-# PII masker
+# PII Masking Function
 def mask_and_store_all_pii(text):
+    text = str(text)
     pii_map = {}
+    entity_list = []
+
     patterns = {
         "email": r"\b[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+\b",
         "phone_number": r"\b\d{10}\b",
@@ -34,27 +37,48 @@ def mask_and_store_all_pii(text):
     }
 
     for label, pattern in patterns.items():
-        matches = re.findall(pattern, text)
-        for i, match in enumerate(matches):
-            placeholder = f"[{label}_{i}]"
-            pii_map[placeholder] = match
-            text = text.replace(match, placeholder)
+        for match in re.finditer(pattern, text):
+            original = match.group()
+            start, end = match.start(), match.end()
+            placeholder = f"[{label}_{len(pii_map)}]"
+            pii_map[placeholder] = original
+            entity_list.append({
+                "position": [start, end],
+                "classification": label,
+                "entity": original
+            })
+            text = text.replace(original, placeholder, 1)
+
+    return text, pii_map, entity_list
 
-    return text, pii_map
+# Restore PII
+def restore_pii(masked_text, pii_map):
+    restored = masked_text
+    for placeholder, original in pii_map.items():
+        restored = restored.replace(placeholder, original)
+    return restored
 
-# Main endpoint
+# Classification Endpoint
 @app.post("/classify")
 def classify_email(data: EmailInput):
-    raw_text = f"{data.subject} {data.email}"
-    masked_text, pii_map = mask_and_store_all_pii(raw_text)
-    prediction = model.predict([masked_text])[0]
+    raw_text = data.input_email_body
+
+    # Masking
+    masked_text, pii_map, entity_list = mask_and_store_all_pii(raw_text)
 
+    # Prediction
+    predicted_category = model.predict([masked_text])[0]
+
+    # Response format
     return {
-        "masked_text": masked_text,
-        "predicted_category": prediction,
-        "pii_map": pii_map
+        "input_email_body": raw_text,
+        "list_of_masked_entities": entity_list,
+        "masked_email": masked_text,
+        "category_of_the_email": predicted_category
     }
 
+# Health check endpoint
 @app.get("/")
 def root():
     return {"message": "Email Classification API is running."}
+