Upload app.py

app.py

@@ -0,0 +1,49 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import torch
+from transformers import RobertaTokenizer
+from config import DEVICE, MAX_LEN, LABEL_COLUMNS, ROBERTA_MODEL_NAME, MODEL_SAVE_DIR, LABEL_ENCODERS_PATH
+from models.roberta_model import RobertaMultiOutputModel
+from dataset_utils import load_label_encoders
+import numpy as np
+
+app = FastAPI()
+
+# Load label encoders
+label_encoders = load_label_encoders()
+
+# Load tokenizer and model
+tokenizer = RobertaTokenizer.from_pretrained(ROBERTA_MODEL_NAME)
+num_labels = [len(label_encoders[col].classes_) for col in LABEL_COLUMNS]
+model = RobertaMultiOutputModel(num_labels)
+model.load_state_dict(torch.load(MODEL_SAVE_DIR + "ROBERTA_model.pth", map_location=DEVICE))
+model.to(DEVICE)
+model.eval()
+
+class RequestText(BaseModel):
+    text: str
+
+@app.post("/predict")
+def predict_labels(request: RequestText):
+    try:
+        inputs = tokenizer(
+            request.text,
+            padding='max_length',
+            truncation=True,
+            max_length=MAX_LEN,
+            return_tensors="pt"
+        )
+        input_ids = inputs['input_ids'].to(DEVICE)
+        attention_mask = inputs['attention_mask'].to(DEVICE)
+
+        with torch.no_grad():
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            predictions = [torch.argmax(logits, dim=1).cpu().numpy()[0] for logits in outputs]
+
+        decoded_predictions = {
+            label: label_encoders[label].inverse_transform([pred])[0]
+            for label, pred in zip(LABEL_COLUMNS, predictions)
+        }
+        return decoded_predictions
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))