Update app.py

app.py

@@ -1,9 +1,9 @@
-from flask import Flask, jsonify
-from threading import Thread
+import streamlit as st
 import time
+from threading import Thread
 from gradio_client import Client
-
-app = Flask(__name__)
+from tenacity import retry, stop_after_attempt, wait_fixed, retry_if_exception_type
+import httpx
 
 # Initialize clients
 client1 = Client("orionai/training-data-collection_2")
@@ -18,37 +18,46 @@ state = {
     "total_growth_speed": 0
 }
 
+@retry(stop=stop_after_attempt(3), wait=wait_fixed(2), retry=retry_if_exception_type(httpx.ReadTimeout))
 def get_token_count():
-    result1 = client1.predict(api_name="/update_token_display")
-    result2 = client2.predict(api_name="/update_token_display")
-    result3 = client3.predict(api_name="/update_token_display")
+    result1 = client1.predict(api_name="/update_token_display", timeout=10.0)
+    result2 = client2.predict(api_name="/update_token_display", timeout=10.0)
+    result3 = client3.predict(api_name="/update_token_display", timeout=10.0)
     return int(result1), int(result2), int(result3)
 
 def monitor_growth():
     while True:
-        curr_count1, curr_count2, curr_count3 = get_token_count()
-        growth_speed1 = curr_count1 - state["prev_count1"]
-        growth_speed2 = curr_count2 - state["prev_count2"]
-        growth_speed3 = curr_count3 - state["prev_count3"]
-        total_tokens = curr_count1 + curr_count2 + curr_count3
-        total_growth_speed = growth_speed1 + growth_speed2 + growth_speed3
-
-        state["prev_count1"] = curr_count1
-        state["prev_count2"] = curr_count2
-        state["prev_count3"] = curr_count3
-        state["prev_total_tokens"] = total_tokens
-        state["total_growth_speed"] = total_growth_speed
+        try:
+            curr_count1, curr_count2, curr_count3 = get_token_count()
+            growth_speed1 = curr_count1 - state["prev_count1"]
+            growth_speed2 = curr_count2 - state["prev_count2"]
+            growth_speed3 = curr_count3 - state["prev_count3"]
+            total_tokens = curr_count1 + curr_count2 + curr_count3
+            total_growth_speed = growth_speed1 + growth_speed2 + growth_speed3
+
+            state["prev_count1"] = curr_count1
+            state["prev_count2"] = curr_count2
+            state["prev_count3"] = curr_count3
+            state["prev_total_tokens"] = total_tokens
+            state["total_growth_speed"] = total_growth_speed
+
+            time.sleep(1)
+        except httpx.ReadTimeout:
+            print("Timeout occurred while fetching token count.")
+
+# Start the monitor thread
+Thread(target=monitor_growth, daemon=True).start()
 
-        time.sleep(1)
+# Streamlit UI
+st.title("Token Count Dashboard")
 
-Thread(target=monitor_growth, daemon=True).start()
+token_count_placeholder = st.empty()
+growth_speed_placeholder = st.empty()
 
-@app.route('/stats', methods=['GET'])
-def get_stats():
-    return jsonify({
-        "total_tokens": state["prev_total_tokens"],
-        "total_growth_speed": state["total_growth_speed"]
-    })
+def update_ui():
+    while True:
+        token_count_placeholder.metric("Total Token Count", state["prev_total_tokens"])
+        growth_speed_placeholder.metric("Total Growth Speed", state["total_growth_speed"])
+        time.sleep(1)
 
-if __name__ == '__main__':
-    app.run(debug=True)
+Thread(target=update_ui, daemon=True).start()