Spaces:
Paused
Paused
File size: 1,618 Bytes
bd89ed8 b6a509f bd89ed8 9471796 bd89ed8 9471796 bd89ed8 9471796 bd89ed8 6508f4c bd89ed8 9471796 bd89ed8 9471796 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
"""
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)
|