Spaces:
Sleeping
Sleeping
File size: 2,812 Bytes
9556312 |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
"""
This module is used to launch Axolotl with user defined configurations.
"""
import gradio as gr
import yaml
def config(
base_model,
dataset,
dataset_type,
learn_rate,
gradient_accumulation_steps,
micro_batch_size,
seq_length,
num_epochs,
output_dir,
val_size,
):
"""
This function generates a configuration dictionary and saves it as a yaml file.
"""
config_dict = {
"base_model": base_model,
"datasets": [{"path": dataset, "type": dataset_type}],
"learning_rate": learn_rate,
"gradient_accumulation_steps": gradient_accumulation_steps,
"micro_batch_size": micro_batch_size,
"sequence_len": seq_length,
"num_epochs": num_epochs,
"output_dir": output_dir,
"val_set_size": val_size,
}
with open("config.yml", "w", encoding="utf-8") as file:
yaml.dump(config_dict, file)
print(config_dict)
return yaml.dump(config_dict)
with gr.Blocks(title="Axolotl Launcher") as demo:
gr.Markdown(
"""
# Axolotl Launcher
Fill out the required fields below to create a training run.
"""
)
base_model_name = gr.Textbox(
"TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", label="Base model"
)
with gr.Row():
dataset_path = gr.Textbox("mhenrichsen/alpaca_2k_test", label="Dataset")
dataset_type_name = gr.Dropdown(
choices=["alpaca", "sharegpt"], label="Dataset type", value="alpaca"
)
with gr.Row():
learning_rate = gr.Number(0.000001, label="Learning rate")
gradient_accumulation_steps_count = gr.Number(
1, label="Gradient accumulation steps"
)
val_set_size_count = gr.Number(0, label="Validation size")
with gr.Row():
micro_batch_size_count = gr.Number(1, label="Micro batch size")
sequence_length = gr.Number(1024, label="Sequence length")
num_epochs_count = gr.Number(1, label="Epochs")
output_dir_path = gr.Textbox("./model-out", label="Output directory")
mode = gr.Radio(
choices=["Full finetune", "QLoRA", "LoRA"],
value="Full finetune",
label="Training mode",
info="FFT = 16 bit, Qlora = 4 bit, Lora = 8 bit",
)
create_config = gr.Button("Create config")
output = gr.TextArea(label="Generated config")
create_config.click(
config,
inputs=[
base_model_name,
dataset_path,
dataset_type_name,
learning_rate,
gradient_accumulation_steps_count,
micro_batch_size_count,
sequence_length,
num_epochs_count,
output_dir_path,
val_set_size_count,
],
outputs=output,
)
demo.launch(share=True)
|