fixing model output formats

app.py

@@ -1,57 +1,100 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
+from datetime import datetime
 import os
 
-
 os.environ["HF_HOME"] = "/tmp" 
 
 SPAM_MODEL = "valurank/distilroberta-spam-comments-detection"
 TOXIC_MODEL = "s-nlp/roberta_toxicity_classifier"
-SENTIMENT_MODEL =  "nlptown/bert-base-multilingual-uncased-sentiment"
+SENTIMENT_MODEL = "nlptown/bert-base-multilingual-uncased-sentiment"
 NSFW_MODEL = "michellejieli/NSFW_text_classifier"
 
+# Load models
 spam = pipeline("text-classification", model=SPAM_MODEL)
-
 toxic = pipeline("text-classification", model=TOXIC_MODEL)
+sentiment = pipeline("text-classification", model=SENTIMENT_MODEL)
+nsfw = pipeline("text-classification", model=NSFW_MODEL)
 
-sentiment = pipeline("text-classification", model = SENTIMENT_MODEL)
-
-nsfw = pipeline("text-classification", model = NSFW_MODEL)
-
-
-app = FastAPI()
-
-@app.get("/")
-def root():
-    return {"status": "ok"}
+app = FastAPI(title="Plebzs AI Models API")
 
 class Query(BaseModel):
     text: str
 
-@app.post("/spam")
-def predict_spam(query: Query):
-    result = spam(query.text)[0]
-    return {"label": result["label"], "score": result["score"]}
+@app.get("/")
+def root():
+    return {"status": "ok", "message": "Plebzs AI Models API"}
+
+# Required by Plebzs boss
+@app.get("/moderation/ping")
+def moderation_ping():
+    return {
+        "status": "healthy",
+        "models": ["spam", "toxic", "sentiment", "nsfw"],
+        "timestamp": datetime.now().isoformat(),
+        "version": "1.0.0"
+    }
 
-@app.post("/toxic")
-def predict_toxic(query: Query):
+# Main endpoints - formatted for Plebzs
+@app.post("/toxicity")  # Changed name to match Plebzs expectation
+def predict_toxicity(query: Query):
     result = toxic(query.text)[0]
-    return {"label": result["label"], "score": result["score"]}
+    
+    # Convert to 0-1 toxicity scale
+    toxicity_score = result["score"] if result["label"] == "TOXIC" else 1 - result["score"]
+    
+    return {
+        "toxicity_score": round(toxicity_score, 3),
+        "confidence": round(result["score"], 3),
+        "raw_output": result
+    }
 
 @app.post("/sentiment")
 def predict_sentiment(query: Query):
     result = sentiment(query.text)[0]
-    return {"label": result["label"], "score": result["score"]}
+    
+    # Convert star rating to -1 to 1 scale
+    label = result["label"]
+    if "1" in label or "2" in label:  # 1-2 stars = negative
+        sentiment_score = -0.7
+    elif "3" in label:  # 3 stars = neutral
+        sentiment_score = 0.0
+    else:  # 4-5 stars = positive
+        sentiment_score = 0.7
+
+        return {
+        "sentiment_score": round(sentiment_score, 3),
+        "confidence": round(result["score"], 3),
+        "raw_output": result
+    }
+
+# Bonus endpoints (not used by Plebzs yet, but good to have)
+@app.post("/spam")
+def predict_spam(query: Query):
+    result = spam(query.text)[0]
+    spam_score = result["score"] if result["label"] == "SPAM" else 1 - result["score"]
+    
+    return {
+        "spam_score": round(spam_score, 3),
+        "confidence": round(result["score"], 3),
+        "raw_output": result
+    }
 
 @app.post("/nsfw")
 def predict_nsfw(query: Query):
     result = nsfw(query.text)[0]
-    return {"label": result["label"], "score": result["score"]}
+    nsfw_score = result["score"] if result["label"] == "NSFW" else 1 - result["score"]
+    
+    return {
+        "nsfw_score": round(nsfw_score, 3),
+        "confidence": round(result["score"], 3),
+        "raw_output": result
+    }
 
+# Keep your detailed health check
 @app.get("/health")
 def health_check():
-
     status = {
         "server": "running",
         "models": {}
@@ -77,4 +120,4 @@ def health_check():
                 "status": f"error: {str(e)}"
             }
 
-    return status
+    return status

test_health.py

@@ -2,6 +2,12 @@ import requests
 
 url = "https://cjell-Demo.hf.space"
 
-if __name__ == "__main__":
-    response = requests.get(f"{url}/health", timeout=10)
-    print(response.json())
+resp = requests.get(f"{url}/health", timeout=10)
+data = resp.json()
+
+print(f"\nServer status: {data['server'].upper()}")
+print("Model statuses:")
+
+for key, info in data["models"].items():
+    status = info["status"].upper()
+    print(f"  - {key}: {info['model_name']} → {status}")

test_spam.py

@@ -10,4 +10,8 @@ print("Status:", response.status_code)
 try:
     print("JSON:", response.json())
 except Exception:
-    print("Raw text:", response.text)
+    print("Raw text:", response.text)
+
+print("")
+
+print(response.text)