Update app.py

app.py

@@ -3,22 +3,27 @@ from pydantic import BaseModel
 import joblib
 import re
 
+# Initialize the FastAPI app with default docs enabled
+app = FastAPI(
+    title="Email Classification API",
+    version="1.0.0",
+    description="Classifies emails and masks PII/PCI information."
+)
+
 # Load the model once when the app starts
 model = joblib.load("model.joblib")
 
-app = FastAPI(title="Email Classification API")
-
+# Define the root endpoint
 @app.get("/")
 def root():
     return {"message": "Email Classification API is running."}
 
-
-# Input data schema
+# Define input data schema
 class EmailInput(BaseModel):
     subject: str = ""
     email: str
 
-# PII masking function (same as your training)
+# Function to mask and store PII
 def mask_and_store_all_pii(text):
     text = str(text)
     pii_map = {}
@@ -43,24 +48,13 @@ def mask_and_store_all_pii(text):
 
     return text, pii_map
 
-# Restore PII function
-def restore_pii(masked_text, pii_map):
-    for placeholder, original in pii_map.items():
-        masked_text = masked_text.replace(placeholder, original)
-    return masked_text
-
+# Endpoint to classify email
 @app.post("/classify")
 def classify_email(data: EmailInput):
-    # Combine subject + email text
     raw_text = f"{data.subject} {data.email}"
-
-    # Mask PII
     masked_text, pii_map = mask_and_store_all_pii(raw_text)
-
-    # Predict class
     prediction = model.predict([masked_text])[0]
 
-    # Return prediction and masked email
     return {
         "predicted_category": prediction,
         "masked_text": masked_text,