import gradio as gr

with gr.Blocks() as app:
    gr.Markdown("# Some calculs about LLMs")

    def calculate_flops(n_params, n_tokens, n_epochs):
        n_params *= 1e9
        n_tokens *= 1e12
        flops = 6 * n_params * n_tokens * n_epochs
        return "%e"%flops, "%e"%flops

    def calculate_time(num_flops, num_flops_gpu):
        num_flops_gpu *= 1e12
        num_flops = float(num_flops)
        time_h = num_flops / num_flops_gpu / 3600
        time_days = time_h / 24
        return f"{round(time_h, 2)} hours or {round(time_days, 2)} days"

    
    gr.Markdown("## FLOPs calculation")
    with gr.Row():
        num_params = gr.Number(label="Number of Parameters (in Billion)", value=1) 
        num_tokens = gr.Number(label="Number of Tokens (in Terra)", value=1)
        num_epochs = gr.Number(label="Number of Epochs", value=1)
    
    button_flops = gr.Button(value="calculate")
    flops = gr.Textbox(label="flops")

    gr.Markdown("## training time calculation")
    with gr.Row():
        num_flops = gr.Number(label="nbr Model FLOPs")
        num_flops_gpu = gr.Number(label="nbr GPU TFLOPs", value=1)

    button_time = gr.Button(value="calculate")
    time = gr.Textbox(label="time")


    button_flops.click(calculate_flops, inputs=[num_params, num_tokens, num_epochs], outputs=[flops, num_flops])
    button_time.click(calculate_time, inputs=[num_flops, num_flops_gpu], outputs=time)



app.launch()