import fn
import gradio as gr
import models

def fn_chat(instruction, input, model, dtype, is_messages, template, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
    args = {
        'instruction': instruction,
        'input': input,
        'model': model,
        'dtype': dtype,
        'is_messages': is_messages,
        'template': template,
        'max_new_tokens': int(max_new_tokens),
        'temperature': float(temperature),
        'top_p': float(top_p),
        'top_k': int(top_k),
        'repetition_penalty': float(repetition_penalty),
    }

    content = fn.infer(args)
    return content

with gr.Blocks() as demo:
    opt = models.get_head_options()

    with gr.Row():
        with gr.Column(scale=1):
            model = gr.Textbox(
                value=opt['model'],
                label='model',
                show_label=True,
                interactive=True,
                show_copy_button=True,
            )

            dtype = gr.Dropdown(
                value=opt['dtype'],
                choices=['int4','int8','fp16', 'bf16'],
                label='dtype',
                show_label=True,
                interactive=True,
                allow_custom_value=True,
            )
            template = gr.Textbox(
                value=opt['template'],
                lines=3,
                label='template',
                show_label=True,
                interactive=True,
                show_copy_button=True,
                )
            is_messages = gr.Checkbox(
                value=opt['is_messages'],
                label='is_messages',
                show_label=True,
                interactive=True,
                )

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

    with gr.Accordion('Preset', open=False):
        gr.Examples(
            models.get_examples(),
            [model, dtype, is_messages, template, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
        )

    with gr.Row():
        with gr.Column(scale=1):
            instruction = gr.Textbox(
                lines=20,
                label='instruction',
                show_label=True,
                interactive=True,
                show_copy_button=True,
                )
            user_input = gr.Textbox(
                lines=1,
                label='input',
                show_label=True,
                interactive=True,
                show_copy_button=True,
                )
            chat_button = gr.Button(value='chat')

        with gr.Column(scale=1):
            said = gr.Textbox(
                label='said',
                lines=15,
                show_label=True,
                show_copy_button=True,
                )

    chat_button.click(
        fn=fn_chat,
        inputs=[instruction, user_input, model, dtype, is_messages, template, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
        outputs=[said],
        )

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