from git import Repo
import gradio as gr

from medusa_training import run, DEFAULT_TRAINING_ARGS

# Clone the medusa repo locally
print("Cloning the medusa repo locally...")
Repo.clone_from("https://github.com/FasterDecoding/Medusa.git", "medusa")
print("Cloning the vicuna data locally...")
Repo.clone_from("https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered", "data")
print("Done")


DESCRIPTION = """
The steps to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:

1. Input a public model id from the Hub
2. Click "Submit"
3. That's it! You'll get feedback if it works or not, and if it worked, you'll get the name of the new repo 🔥
"""

title="Create LLM medusa heads in a new repo 🐍"

with gr.Blocks(title=title) as demo:
    description = gr.Markdown(f"""# {title}""")
    description = gr.Markdown(DESCRIPTION)

    with gr.Row() as r:
        with gr.Column() as c:
            model_id = gr.Text(max_lines=1, label="model_id")
            with gr.Accordion("Training arguments (advanced)", open=False):
                training_args = gr.Textbox(DEFAULT_TRAINING_ARGS, interactive=True, lines=14, label="training_args")
            with gr.Row() as c:
                clean = gr.ClearButton()
                submit = gr.Button("Submit", variant="primary")

        with gr.Column() as d:
            status_box = gr.Markdown()

    submit.click(run, inputs=[model_id, training_args], outputs=status_box, concurrency_limit=1)

demo.queue(max_size=10).launch(show_api=True)