Update app.py

app.py

@@ -2,47 +2,62 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
-MODEL_ID = "cardiffnlp/twitter-roberta-base-sentiment-latest"
+PRIMARY_MODEL = "cardiffnlp/twitter-roberta-base-sentiment-latest"
+FALLBACK_MODEL = "distilbert-base-uncased-finetuned-sst-2-english"
 
 app = FastAPI()
+clf = None
+loaded_model_id = None
 
 @app.get("/")
 def root():
-    return {"status": "ok"}
+    return {"status": "ok", "model": loaded_model_id}
 
 @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,
-)
+def health():
+    return {"status": "healthy", "model": loaded_model_id}
 
 class Payload(BaseModel):
     sentences: list[str]
 
-def to_label(scores):
+def load_pipeline():
+    global clf, loaded_model_id
+    if clf is not None:
+        return clf
+    # Try primary (3-class), then fallback (binary)
+    for model_id in (PRIMARY_MODEL, FALLBACK_MODEL):
+        try:
+            tok = AutoTokenizer.from_pretrained(model_id)
+            mdl = AutoModelForSequenceClassification.from_pretrained(model_id)
+            clf = pipeline("text-classification", model=mdl, tokenizer=tok, return_all_scores=True, truncation=True)
+            loaded_model_id = model_id
+            return clf
+        except Exception as e:
+            print(f"Failed to load {model_id}: {e}")
+    raise RuntimeError("No sentiment model could be loaded")
+
+def scores_to_label(scores, binary=False):
     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"
+    if binary or ("neutral" not in m):
+        neg, pos = m.get("negative", 0.0), m.get("positive", 0.0)
+        conf = max(neg, pos)
+        label = "positive" if pos >= neg else "negative"
+        return {"label": label, "confidence": conf, "scores": {"positive": pos, "neutral": 0.0, "negative": neg}}
     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}}
+        neg, neu, pos = m.get("negative", 0.0), m.get("neutral", 0.0), m.get("positive", 0.0)
+        conf = max(neg, neu, pos)
+        # Optional neutral clamp for low confidence
+        label = "positive" if pos >= max(neg, neu) else ("negative" if neg >= max(pos, neu) else "neutral")
+        if conf < 0.55:
+            label = "neutral"
+        return {"label": label, "confidence": conf, "scores": {"positive": pos, "neutral": neu, "negative": neg}}
 
 @app.post("/predict")
 def predict(payload: Payload):
+    classifier = load_pipeline()
     if not payload.sentences:
         return {"results": []}
-    outputs = clf(payload.sentences, top_k=None)
-    results = [to_label(scores) for scores in outputs]
-    return {"results": results}
+    outputs = classifier(payload.sentences, top_k=None)
+    binary = (loaded_model_id == FALLBACK_MODEL)
+    results = [scores_to_label(scores, binary=binary) for scores in outputs]
+    return {"model": loaded_model_id, "results": results}