Add retry logic and wait_for_model

app.py

@@ -1,14 +1,23 @@
 import gradio as gr
 import requests
 import os
+import time
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 API_URL = "https://api-inference.huggingface.co/models/Hate-speech-CNERG/dehatebert-mono-english"
 
 def query(text):
     headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    response = requests.post(API_URL, headers=headers, json={"inputs": text})
-    return response.json()
+    # Retry up to 3 times if model is loading
+    for i in range(3):
+        response = requests.post(API_URL, headers=headers, json={
+            "inputs": text,
+            "options": {"wait_for_model": True}
+        })
+        if response.status_code == 200:
+            return response.json()
+        time.sleep(3)
+    return {"error": f"Status: {response.status_code}, Response: {response.text}"}
 
 def predict(comment):
     if not comment.strip():
@@ -17,9 +26,8 @@ def predict(comment):
     try:
         result = query(comment)
         
-        # Show raw result for debugging
-        if isinstance(result, dict):
-            return f"API Error: {result}", ""
+        if isinstance(result, dict) and "error" in result:
+            return f"API Error: {result['error']}", ""
         
         if isinstance(result, list):
             result = result[0][0]
@@ -35,7 +43,7 @@ def predict(comment):
             
             return prediction, f"{probability}%"
         
-        return f"Unexpected response: {result}", ""
+        return f"Unexpected: {result}", ""
     
     except Exception as e:
         return f"Exception: {str(e)}", ""