import gradio as gr

from mini_llm.config import LLMConfig
from mini_llm.service import LocalMiniLLMService


config = LLMConfig()
service = LocalMiniLLMService(config=config)


def generate_text(prompt, max_new_tokens, temperature, top_k):
    text, status = service.generate(
        prompt=prompt,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        top_k=int(top_k),
    )
    return text, status


def train_model(extra_text, steps):
    return service.train(extra_text=extra_text, steps=int(steps))


def reset_model():
    return service.reset()


with gr.Blocks(
    title="Local Small LLM",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray"),
) as demo:
    gr.Markdown(
        """
        # Local Small LLM
        A tiny language model written in Python and trained from scratch on local text.

        - No external pretrained LLM
        - Tiny transformer decoder
        - Character-level tokenizer
        - CPU-friendly local training
        """
    )

    with gr.Tab("Generate"):
        prompt_input = gr.Textbox(
            label="Prompt",
            value="User: hello\nAssistant:",
            lines=6,
        )
        with gr.Row():
            max_tokens_input = gr.Slider(20, 240, value=120, step=10, label="Max New Tokens")
            temperature_input = gr.Slider(0.2, 1.5, value=0.8, step=0.05, label="Temperature")
            top_k_input = gr.Slider(1, 32, value=16, step=1, label="Top-K")
        generate_button = gr.Button("Generate", variant="primary")
        output_text = gr.Textbox(label="Output", lines=10)
        output_status = gr.Textbox(label="Status", lines=3)

    with gr.Tab("Train"):
        extra_text_input = gr.Textbox(
            label="Extra Training Text",
            placeholder="Add your own local text here to continue training the model.",
            lines=10,
        )
        steps_input = gr.Slider(10, 400, value=120, step=10, label="Training Steps")
        train_button = gr.Button("Train / Continue Training", variant="primary")
        reset_button = gr.Button("Reset Model")
        train_status = gr.Textbox(label="Training Status", lines=6)

    generate_button.click(
        fn=generate_text,
        inputs=[prompt_input, max_tokens_input, temperature_input, top_k_input],
        outputs=[output_text, output_status],
    )

    train_button.click(
        fn=train_model,
        inputs=[extra_text_input, steps_input],
        outputs=[train_status],
    )

    reset_button.click(
        fn=reset_model,
        outputs=[train_status],
    )


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