Spaces:
Runtime error
Runtime error
File size: 1,227 Bytes
36f5dbc 69857da 54b76e3 8e1727c 36f5dbc 54b76e3 36f5dbc 54b76e3 8e1727c 36f5dbc 8e1727c 36f5dbc |
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 |
from fastapi import FastAPI, HTTPException
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
import torch
app = FastAPI()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 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.post("/predict/")
def predict(text: str):
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)
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)}"
)
|