File size: 2,037 Bytes
8661759 6f388b3 8661759 02b4129 8661759 6f388b3 8661759 6f388b3 8661759 6f388b3 8661759 6f388b3 8661759 6f388b3 8661759 6f388b3 8661759 6f388b3 |
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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# -------------------------------------------------
# CONFIG
# -------------------------------------------------
MODEL_NAME = "j-hartmann/emotion-english-distilroberta-base"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# -------------------------------------------------
# LOAD MODEL (ONCE AT STARTUP)
# -------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.to(device)
model.eval()
# -------------------------------------------------
# FASTAPI APP
# -------------------------------------------------
app = FastAPI(title="Emotion Detection API")
# ✅ ROOT ROUTE (CRITICAL FOR HF SPACES)
@app.get("/")
def health():
"""
Health / wake-up endpoint.
Hugging Face uses this to wake the Space.
"""
return {"status": "ok"}
# -------------------------------------------------
# REQUEST SCHEMA
# -------------------------------------------------
class EmotionRequest(BaseModel):
text: str
# -------------------------------------------------
# EMOTION ENDPOINT
# -------------------------------------------------
@app.post("/emotion")
def classify_emotion(payload: EmotionRequest):
text = payload.text.strip()
if not text:
raise HTTPException(status_code=400, detail="Text cannot be empty")
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=128,
)
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=-1)
pred_id = torch.argmax(probs, dim=-1).item()
return {
"emotion": model.config.id2label[pred_id],
"confidence": round(probs[0][pred_id].item(), 4),
}
|