app.py

import gradio as gr
from gradio_client import Client

# 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,
    "prev_count2": 0,
    "prev_count3": 0,
    "prev_total_tokens": 0,
    "total_growth_speed": 0
}

def get_token_count(client):
    try:
        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 update_dashboard():
    try:
        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"]
        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

        return total_tokens, total_growth_speed
    except Exception as e:
        print(f"Error in update dashboard: {e}")
        return 0, 0

def refresh_metrics():
    tokens, speed = update_dashboard()
    return tokens, speed

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Token Count Dashboard")
    
    total_tokens = gr.Number(label="Total Token Count")
    growth_speed = gr.Number(label="Total Growth Speed")
    
    update_button = gr.Button("Update Stats")
    update_button.click(fn=refresh_metrics, inputs=[], outputs=[total_tokens, growth_speed])
    
demo.launch()