import gradio as gr

import run

title = "Lit GPT: Pythia 160M "

with gr.Blocks(title=title) as interface:
    with gr.Row():
        prompt = gr.Textbox(label="Input Text")

        temperature = gr.Slider(
            0,
            1,
            value=0.8,
            label="Temperature",
            info="Set the creativity level: Higher values produce more varied results, lower values generate more predictable text.",
        )
        top_k = gr.Slider(
            200,
            300,
            value=200,
            label="Top K",
            info="Control the randomness: Limits the AI to consider only the top K most likely next words.",
        )
        max_new_tokens = gr.Slider(
            10,
            500,
            value=500,
            label="Max Tokens",
            info="top most preferable tokens to consider in the sampling process",
        )

        inputs = [prompt, max_new_tokens, top_k, temperature]

    with gr.Column():
        outputs = gr.Textbox(label="Generated")
        button = gr.Button("Generate")
        button.click(run.generate_from_prompt, inputs=inputs, outputs=outputs)

    # with gr.Row():
    #     gr.Examples(examples=examples, inputs=inputs, outputs=outputs, fn=generate_dialogue, cache_examples=True,)


interface.launch()