Update app.py

app.py

@@ -2,11 +2,11 @@ import gradio as gr
 import requests
 import os
 import json
+import re
 
 # Configure the endpoint and authentication
 ENDPOINT_URL = os.environ.get("ENDPOINT_URL", "https://dz0eq6vxq3nm0uh7.us-east-1.aws.endpoints.huggingface.cloud")
-# HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # Get API token from environment variable
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()  # Use strip() to remove extra whitespaces and newlines
+HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()
 
 
 def is_token_configured():
@@ -14,14 +14,10 @@ def is_token_configured():
         return "⚠️ Warning: HF_API_TOKEN is not configured. The app won't work until you add this secret in your Space settings."
     return "✅ API token is configured"
 
-
-# Assuming your class with extract_and_parse_json exists
-# Or you can define it separately if needed
-
 class SafetyChecker:
     def __init__(self):
-        self.ENDPOINT_URL = os.environ.get("ENDPOINT_URL", "https://your-endpoint")
-        self.HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()
+        self.ENDPOINT_URL = ENDPOINT_URL
+        self.HF_API_TOKEN = HF_API_TOKEN
 
     def extract_and_parse_json(self, response: str):
         match = re.search(r'```(?:json)?\s*(.*?)\s*```', response, re.DOTALL)
@@ -64,14 +60,13 @@ class SafetyChecker:
             response = requests.post(self.ENDPOINT_URL, json=payload, headers=headers, timeout=30)
 
             if response.status_code == 200:
-                result_raw = response.json()  # still a string inside triple backticks
+                result_raw = response.json()
                 
                 if isinstance(result_raw, str):
                     parsed_result = self.extract_and_parse_json(result_raw)
                 else:
                     parsed_result = result_raw
 
-                # Now parsed_result is a dictionary
                 safety = parsed_result.get("Safety", "Unknown")
                 score = parsed_result.get("Score", "")
                 categories = parsed_result.get("Unsafe Categories", "")
@@ -92,6 +87,10 @@ class SafetyChecker:
         except Exception as e:
             return f"❗ Error: {str(e)}"
 
+# Create an instance of SafetyChecker
+checker = SafetyChecker()
+
+# Build Gradio App
 with gr.Blocks(title="Safety Content Classifier", css="footer {display: none !important}") as demo:
     gr.Markdown("# Safety Content Classifier")
     gr.Markdown("## Connected to external safety model endpoint")
@@ -126,9 +125,9 @@ with gr.Blocks(title="Safety Content Classifier", css="footer {display: none !im
                 lines=8
             )
     
-    # Update event handlers
-    check_button.click(fn=check_safety, inputs=[input_text], outputs=output)
-    input_text.submit(fn=check_safety, inputs=[input_text], outputs=output)
+    # Notice: Now we use checker.check_safety instead of check_safety directly
+    check_button.click(fn=checker.check_safety, inputs=[input_text], outputs=output)
+    input_text.submit(fn=checker.check_safety, inputs=[input_text], outputs=output)
     
     gr.Examples(
         [
@@ -141,4 +140,3 @@ with gr.Blocks(title="Safety Content Classifier", css="footer {display: none !im
     )
 
 demo.launch()
-