from dotenv import load_dotenv
from generators import *
import gradio as gr

from utils import async_zip_stream

load_dotenv()


async def handle(system_input: str, user_input: str):
    print(system_input, user_input)
    buffers = ["", "", "", "", ""]
    async for outputs in async_zip_stream(
            generate_gpt2(system_input, user_input),
            generate_mistral_7bvo1(system_input, user_input),
            generate_llama2(system_input, user_input),
            generate_llama3(system_input, user_input),
            generate_t5(system_input, user_input),
    ):
        # gpt_output, mistral_output, llama_output, llama2_output, llama3_output, llama4_output = outputs
        for i, b in enumerate(buffers):
            buffers[i] += str(outputs[i])

        yield list(buffers) + ["", ""]
    yield list(buffers) + [generate_bloom(system_input, user_input)]


with gr.Blocks() as demo:
    system_input = gr.Textbox(label='System Input', value='You are AI assistant', lines=2)
    with gr.Row():
        gpt = gr.Textbox(label='gpt-2', lines=4, interactive=False)
        mistral = gr.Textbox(label='mistral-v01', lines=4, interactive=False)
        t5 = gr.Textbox(label='t5', lines=4, interactive=False)
    with gr.Row():
        llama2 = gr.Textbox(label='llama-2', lines=4, interactive=False)
        llama3 = gr.Textbox(label='llama-3', lines=4, interactive=False)
        bloom = gr.Textbox(label='bloom [GPU]', lines=4, interactive=False)

    user_input = gr.Textbox(label='User Input', lines=2, value='Calculate expression: 7-3=')
    gen_button = gr.Button('Generate')

    gen_button.click(
        fn=handle,
        inputs=[system_input, user_input],
        outputs=[gpt, mistral, llama2, llama3, t5, bloom],
    )

demo.launch()