if __name__ == "__main__":
    import subprocess
    subprocess.run('bash setup.sh', shell=True)
    subprocess.run('bash venv.sh', shell=True)

import fn
import gradio as gr

with gr.Blocks() as demo:
    with gr.Tab('config'):
        info = gr.Markdown()
        with gr.Row():
            with gr.Column(scale=1):
                model = gr.Textbox(
                    value=fn.cfg['model_name'],
                    label='model',
                    interactive=True,
                    show_copy_button=True,
                )
                dtype = gr.Dropdown(
                    value=fn.cfg['dtype'],
                    choices=['4bit'],
                    label='dtype',
                    interactive=True,
                    allow_custom_value=True,
                )

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

        with gr.Row():
            with gr.Column(scale=1):
                inst_template = gr.Textbox(
                    value='',
                    lines=10,
                    label='inst_template',
                    interactive=True,
                    show_copy_button=True,
                    )
            with gr.Column(scale=1):
                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=[model, dtype, instruction, inst_template, 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()