Update app.py

app.py

@@ -3,47 +3,60 @@
 # ===============================================
 
 from fastapi import FastAPI, Request
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import pipeline
 import uvicorn
 
-# --- Initialize app ---
-app = FastAPI(title="Bloom Check-in Quality Classifier")
-
-# --- Allow web access from anywhere (CORS) ---
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+app = FastAPI(
+    title="Bloom Check-in Quality Classifier",
+    description="A FastAPI app that classifies check-ins as vague, neutral, or descriptive.",
+    version="1.0.0"
 )
 
-# --- Load model once on startup ---
-MODEL_NAME = "user6295018/checkin-quality-classifier"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-label_map = {0: "Vague", 1: "Neutral", 2: "Descriptive"}
+# Load the Hugging Face text classification model
+# This automatically downloads your public model from the Hub
+classifier = pipeline("text-classification", model="user6295018/checkin-quality-classifier")
 
 @app.get("/")
-async def root():
+def read_root():
+    """
+    Root endpoint – simple health check
+    """
     return {"message": "Bloom Check-in Quality Classifier is running!"}
 
 @app.post("/api/predict")
 async def predict(request: Request):
-    data = await request.json()
-    text = data.get("text", "").strip()
+    """
+    Predict endpoint – classify input text using the fine-tuned model
+    """
+    try:
+        data = await request.json()
+    except Exception:
+        return {"error": "Invalid JSON body."}
+
+    # Accept either {"text": "..."} or {"inputs": "..."}
+    text = data.get("text") or data.get("inputs")
+
     if not text:
         return {"error": "No text provided."}
 
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-        pred = torch.argmax(outputs.logits, dim=1).item()
+    try:
+        # Run model inference
+        result = classifier(text)[0]
+
+        # Optional: Map model output labels to human-friendly names
+        label_map = {
+            "LABEL_0": "vague",
+            "LABEL_1": "neutral",
+            "LABEL_2": "descriptive"
+        }
+        label = label_map.get(result["label"], result["label"])
+        score = round(float(result["score"]), 3)
+
+        return {"label": label, "score": score}
 
-    return {"label": label_map.get(pred, "Unknown")}
+    except Exception as e:
+        return {"error": f"Inference failed: {str(e)}"}
 
-# --- ✅ Explicit entry point so Spaces knows what to run ---
+# Optional: local testing entry point (ignored on Spaces)
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)