Create app.py

app.py

@@ -1,14 +1,53 @@
-from fastapi import FastAPI, Request
-from pydantic import BaseModel
-from app.model_utils import load_model, predict_label
+import torch
+import pickle
+import torch.nn as nn
+from transformers import BertTokenizer, BertModel
 
-app = FastAPI()
-tokenizer, model, label_encoders = load_model()
+PRETRAINED_MODEL_NAME = 'bert-base-uncased'
+LABEL_COLUMNS = [
+    'Red_Flag_Reason',
+    'Maker_Action',
+    'Escalation_Level',
+    'Risk_Category',
+    'Risk_Drivers',
+    'Investigation_Outcome'
+]
 
-class InputText(BaseModel):
-    text: str
+class BertMultiOutput(nn.Module):
+    def __init__(self, num_labels_per_output):
+        super().__init__()
+        self.bert = BertModel.from_pretrained(PRETRAINED_MODEL_NAME)
+        self.dropout = nn.Dropout(0.3)
+        self.classifiers = nn.ModuleList([
+            nn.Linear(self.bert.config.hidden_size, n_labels)
+            for n_labels in num_labels_per_output
+        ])
 
-@app.post("/predict")
-def predict(input: InputText):
-    predictions = predict_label(input.text, tokenizer, model, label_encoders)
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = self.dropout(outputs.pooler_output)
+        logits = [classifier(pooled_output) for classifier in self.classifiers]
+        return logits
+
+def load_model():
+    with open("bert_model.pkl", "rb") as f:
+        bundle = pickle.load(f)
+    tokenizer = bundle['tokenizer']
+    label_encoders = bundle['label_encoders']
+    num_labels = [len(le.classes_) for le in label_encoders.values()]
+    model = BertMultiOutput(num_labels)
+    model.load_state_dict(bundle['model_state_dict'])
+    model.eval()
+    return tokenizer, model, label_encoders
+
+def predict_label(text, tokenizer, model, label_encoders):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    predictions = {}
+    for i, logits in enumerate(outputs):
+        pred_class = torch.argmax(logits, dim=1).item()
+        label_name = LABEL_COLUMNS[i]
+        pred_label = label_encoders[label_name].inverse_transform([pred_class])[0]
+        predictions[label_name] = pred_label
     return predictions