Update app.py

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
-import requests
+from gradio_client import Client
 import time
 from threading import Thread
 
-# URLs for the APIs
-client1_url = 'https://orionai-training-data-collection-2.hf.space/api/update_token_display'
-client2_url = 'https://orionai-training-data-collection-3.hf.space/api/update_token_display'
-client3_url = 'https://orionai-training-data-collection.hf.space/api/update_token_display'
+# Initialize clients
+client1 = Client("orionai/training-data-collection-2")
+client2 = Client("orionai/training-data-collection-3")
+client3 = Client("orionai/training-data-collection")
 
 state = {
     "prev_count1": 0,
@@ -16,22 +16,20 @@ state = {
     "total_growth_speed": 0
 }
 
-def get_token_count(url):
+def get_token_count(client):
     try:
-        response = requests.post(url, headers={'Content-Type': 'application/json'})
-        response.raise_for_status()  # Raise an error for bad status codes
-        result = response.json()
-        return int(result.get('token_count', 0))  # Adjust according to API response
-    except (requests.RequestException, ValueError) as e:
-        print(f"Error fetching token count from {url}: {e}")
+        result = client.predict(api_name="/update_token_display")
+        return int(result)
+    except Exception as e:
+        print(f"Error fetching token count: {e}")
         return 0
 
 def monitor_growth():
     while True:
         try:
-            curr_count1 = get_token_count(client1_url)
-            curr_count2 = get_token_count(client2_url)
-            curr_count3 = get_token_count(client3_url)
+            curr_count1 = get_token_count(client1)
+            curr_count2 = get_token_count(client2)
+            curr_count3 = get_token_count(client3)
             growth_speed1 = curr_count1 - state["prev_count1"]
             growth_speed2 = curr_count2 - state["prev_count2"]
             growth_speed3 = curr_count3 - state["prev_count3"]
@@ -62,9 +60,10 @@ with gr.Blocks() as demo:
     
     def update_dashboard():
         tokens, speed = get_dashboard_metrics()
-        return tokens, speed
+        total_tokens.update(value=tokens)
+        growth_speed.update(value=speed)
     
     # Automatically update the metrics every 5 seconds
-    gr.Timer(5000, update_dashboard, [total_tokens, growth_speed], every=5000)
+    gr.Timer(5, update_dashboard)
 
 demo.launch()