File size: 883 Bytes
9bfe8cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | # api.py
from fastapi import FastAPI
from pydantic import BaseModel
from utils import mask_pii, unmask_pii
from models import load_model
# Create FastAPI app
app = FastAPI()
# Load the trained model at startup
model = load_model("model/classifier.pkl")
# Define request schema
class EmailRequest(BaseModel):
email_body: str
# Define the API endpoint
@app.post("/predict")
def classify_email(req: EmailRequest):
email_text = req.email_body
# Step 1: Mask PII
masked_text, entities = mask_pii(email_text)
# Step 2: Predict Category
category = model.predict([masked_text])[0]
# Step 3: Build the response
response = {
"input_email_body": email_text,
"list_of_masked_entities": entities,
"masked_email": masked_text,
"category_of_the_email": category
}
return response
|