app.py

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

MODEL_ID = "cardiffnlp/twitter-roberta-base-sentiment-latest"

app = FastAPI()

@app.get("/")
def root():
    return {"status": "ok"}

@app.get("/healthz")
def healthz():
    return {"status": "healthy"}

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
clf = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    return_all_scores=True,
    truncation=True,
)

class Payload(BaseModel):
    sentences: list[str]

def to_label(scores):
    m = {s["label"].lower(): s["score"] for s in scores}
    neg = m.get("negative", m.get("label_0", 0.0))
    neu = m.get("neutral",  m.get("label_1", 0.0))
    pos = m.get("positive", m.get("label_2", 0.0))
    conf = max(neg, neu, pos)
    if conf < 0.55:
        final = "neutral"
    else:
        final = "positive" if pos >= max(neg, neu) else ("negative" if neg >= max(pos, neu) else "neutral")
    return {"label": final, "confidence": conf, "scores": {"positive": pos, "neutral": neu, "negative": neg}}

@app.post("/predict")
def predict(payload: Payload):
    if not payload.sentences:
        return {"results": []}
    outputs = clf(payload.sentences, top_k=None)
    results = [to_label(scores) for scores in outputs]
    return {"results": results}