Update app.py

app.py

@@ -1,69 +1,78 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline
+from openai import OpenAI
 import uvicorn
 
 app = FastAPI()
 
-# ===== LOAD MODEL (LIGHTWEIGHT MULTILINGUAL TOXICITY) =====
-classifier = pipeline(
-    "text-classification",
-    model="unitary/multilingual-toxic-xlm-roberta",
-    top_k=None
-)
+# =========================
+# OPENAI CLIENT
+# =========================
+client = OpenAI()  # uses OPENAI_API_KEY from environment / HF Secrets
 
-# ===== RESPONSE MODEL =====
+# =========================
+# REQUEST BODY
+# =========================
 class TextIn(BaseModel):
     text: str
 
 
+# =========================
+# CORE MODERATION FUNCTION
+# =========================
 def analyze_text(text: str):
-    result = classifier(text)[0]
-
-    # model output: list of dicts
-    # we normalize it
-    output = []
-    for r in result:
-        output.append({
-            "label": r["label"],
-            "score": float(r["score"])
-        })
-
-    # find toxic score (if label exists)
-    toxic_score = 0.0
-    for r in output:
-        if "toxic" in r["label"].lower():
-            toxic_score = r["score"]
-
-    return {
-        "input": text,
-        "toxic_score": toxic_score,
-        "details": output,
-        "is_flagged": toxic_score > 0.5
-    }
-
-
-# ===== GET ENDPOINT (WHAT YOU WANT) =====
+    try:
+        response = client.moderations.create(
+            model="omni-moderation-latest",
+            input=text
+        )
+
+        result = response.results[0]
+
+        return {
+            "input": text,
+            "flagged": result.flagged,
+            "categories": result.categories.model_dump(),
+            "category_scores": result.category_scores.model_dump()
+        }
+
+    except Exception as e:
+        return {
+            "input": text,
+            "error": str(e),
+            "flagged": False
+        }
+
+
+# =========================
+# GET ENDPOINT
+# =========================
 @app.get("/analyze")
 def analyze_get(text: str):
     return analyze_text(text)
 
 
-# ===== OPTIONAL POST ENDPOINT =====
+# =========================
+# POST ENDPOINT
+# =========================
 @app.post("/analyze")
 def analyze_post(body: TextIn):
     return analyze_text(body.text)
 
 
-# ===== ROOT =====
+# =========================
+# ROOT
+# =========================
 @app.get("/")
 def root():
     return {
         "status": "ok",
-        "usage": "/analyze?text=your_message_here"
+        "usage": "/analyze?text=your_text_here"
     }
 
 
-# ===== RUN (for local testing only) =====
+# =========================
+# RUN (for local testing only)
+# =========================
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)