app.py

"""
Holds the gradio app itself
"""

import gradio as gr

from src.train_workflow import run, DEFAULT_TRAINING_ARGS
from src.calibration_datasets import CalibrationDataset


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. Select a dataset to train the medusa heads on. The dataset should be representative of the downstream use case.
3. Click "Submit"
4. 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")
            dataset_names = [
                cls.dataset for cls in CalibrationDataset.__subclasses__()
            ]
            dataset = gr.Dropdown(dataset_names, label="dataset")
            with gr.Accordion("Training arguments (advanced)", open=False):
                training_args = gr.Textbox(DEFAULT_TRAINING_ARGS, interactive=True, lines=20, 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, dataset], outputs=status_box, concurrency_limit=1)

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