Update main.py

main.py

@@ -0,0 +1,38 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import List, Dict, Any
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+from utils import mask_pii_multilingual
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("./model")
+model = AutoModelForSequenceClassification.from_pretrained("./model")
+model.eval()
+
+app = FastAPI()
+
+class EmailInput(BaseModel):
+    input_email_body: str
+
+@app.post("/classify")
+def classify_email(input: EmailInput):
+    email_body = input.input_email_body
+
+    # Step 1: Mask PII
+    masked_text, masked_entities = mask_pii_multilingual(email_body)
+
+    # Step 2: Classification
+    inputs = tokenizer(masked_text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+        pred = torch.argmax(logits, dim=1).item()
+
+    label_map = {0: "Incident", 1: "Request", 2: "Change", 3: "Problem"}
+
+    return {
+        "input_email_body": email_body,
+        "list_of_masked_entities": masked_entities,
+        "masked_email": masked_text,
+        "category_of_the_email": label_map[pred]
+    }