import fn
import gradio as gr

with gr.Blocks() as demo:
    gr.Markdown('# gemma2')
    with gr.Tab('config'):
        info = gr.Markdown()
        with gr.Row():
            with gr.Column(scale=1):
                size = gr.Dropdown(
                    value=fn.cfg['size'],
                    choices=['9b','27b'],
                    label='size',
                    interactive=True,
                )

            with gr.Column(scale=1):
                max_new_tokens = gr.Textbox(
                    value=fn.default_args['max_new_tokens'],
                    label='max_new_tokens',
                    interactive=True,
                    show_copy_button=True,
                    )
                temperature = gr.Textbox(
                    value=fn.default_args['temperature'],
                    label='temperature',
                    interactive=True,
                    show_copy_button=True,
                    )
                top_p = gr.Textbox(
                    value=fn.default_args['top_p'],
                    label='top_p',
                    interactive=True,
                    show_copy_button=True,
                    )
                top_k = gr.Textbox(
                    value=fn.default_args['top_k'],
                    label='top_k',
                    interactive=True,
                    show_copy_button=True,
                    )
                repetition_penalty = gr.Textbox(
                    value=fn.default_args['repetition_penalty'],
                    label='repetition_penalty',
                    interactive=True,
                    show_copy_button=True,
                    )

        with gr.Row():
            with gr.Column(scale=1):
                first_assistant = gr.Textbox(
                    value='',
                    label='first_assistant',
                    interactive=True,
                    show_copy_button=True,
                    )
                chat_template = gr.Textbox(
                    value='',
                    lines=10,
                    label='chat_template',
                    interactive=True,
                    show_copy_button=True,
                    )

        set_button = gr.Button(value='Save')

    with gr.Tab('instruct'):
        with gr.Row():
            with gr.Column(scale=1):
                instruction = gr.Textbox(
                    lines=20,
                    label='instruction',
                    interactive=True,
                    show_copy_button=True,
                    )
                input = gr.Textbox(
                    lines=1,
                    label='input',
                    interactive=True,
                    show_copy_button=True,
                    )
            with gr.Column(scale=1):
                said = gr.Textbox(
                    label='said',
                    lines=20,
                    show_copy_button=True,
                    )
                numel = gr.Textbox(
                    lines=1,
                    label='numel',
                    show_copy_button=True,
                    )
        inst_button = gr.Button(value='inst')
        numel_button = gr.Button(value='numel')

    with gr.Tab('chat'):
        gr.ChatInterface(fn.chat)

    set_button.click(
        fn=fn.set_config,
        inputs=[size, instruction, first_assistant, chat_template, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
        outputs=[info],
        )

    inst_button.click(
        fn=fn.chat,
        inputs=[input, input, instruction],
        outputs=[said],
        )

    numel_button.click(
        fn=fn.numel,
        inputs=[input, input, instruction],
        outputs=[numel],
        )

if __name__ == '__main__':
    demo.launch()