Update app.py

app.py

@@ -1,25 +1,33 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
 import torch
 
-# Verify model configuration
-config = AutoConfig.from_pretrained("SrivarshiniGanesan/finetuned-stress-model")
-print(config)
-# Initialize FastAPI app
 app = FastAPI()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model = AutoModelForSequenceClassification.from_pretrained("SrivarshiniGanesan/finetuned-stress-model",config=config)
+# Load config first
+config = AutoConfig.from_pretrained("SrivarshiniGanesan/finetuned-stress-model")
+model = AutoModelForSequenceClassification.from_pretrained(
+    "SrivarshiniGanesan/finetuned-stress-model",
+    config=config
+).to(device)
 tokenizer = AutoTokenizer.from_pretrained("SrivarshiniGanesan/finetuned-stress-model")
 
-@app.get("/")
-def home():
-    return {"message": "Stress Prediction API is running"}
-
 @app.post("/predict/")
 def predict(text: str):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
+    try:
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
         probs = torch.softmax(outputs.logits, dim=-1)
-    stress_prob = probs[:, 1].item()  # Probability of "Stress" class
-    return {"stress_probability": stress_prob}
+        class_labels = config.id2label if config.id2label else {0: "No Stress", 1: "Stress"}
+        stress_idx = list(class_labels.values()).index("Stress")
+        
+        return {"stress_probability": probs[0, stress_idx].item()}
+    
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Prediction failed: {str(e)}"
+        )