Spaces:
Running
Running
| """ | |
| 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) | |