#%%
import numpy as np
import gradio as gr

def greet(name):
    return "Hello " + name + "!!"

def calc_extra_token_length(old_seq_len: int = 4096, old_base: int = 10000, new_base: int = 1000000, dim: int = 128):
    """
    from arXiv:2310.05209v1 [cs.CL] 8 Oct 2023
    """
    b = int(dim/2 * np.log(old_seq_len / 2 / np.pi) /
            np.log(old_base)) * 2 / dim
    extra = 2 * np.pi * new_base ** b
    return extra / 1024

def hello_world(name):
    return "你好 " + name + "!!"

#%%
with gr.Blocks() as demo:
    gr.Markdown("# some tools")

    with gr.Tab("greet"):
        inputs = gr.Text(label="input text")
        btn = gr.Button()
        outputs = gr.Text(label="output text")
    btn.click(fn=greet, inputs=inputs, outputs=outputs)

    with gr.Tab("hello world"):
        inputs = gr.Text(label="input text")
        btn = gr.Button()
        outputs = gr.Text(label="output text")
    btn.click(fn=hello_world, inputs=inputs, outputs=outputs)

    with gr.Tab("calc_extra_token_length"):
        with gr.Row():
            with gr.Column():
                inputs = [
                    gr.Number(4096, label="old_seq_len"), 
                    gr.Number(10000, label="old_base"), 
                    gr.Number(1000000, label="new_base"), 
                    gr.Number(128, label="dim"),
                ]
                btn = gr.Button("calc")
            outputs = gr.Textbox(value=0, label="extra max token (k)")
    btn.click(fn=calc_extra_token_length, inputs=inputs, outputs=outputs)

demo.launch()


# %%