Spaces:
Build error
Build error
| from config import SHARE, MODELS, TRAINING_PARAMS, LORA_TRAINING_PARAMS, GENERATION_PARAMS | |
| import os | |
| import gradio as gr | |
| import random | |
| from trainer import Trainer | |
| LORA_DIR = 'lora' | |
| def random_name(): | |
| fruits = [ | |
| "dragonfruit", "kiwano", "rambutan", "durian", "mangosteen", | |
| "jabuticaba", "pitaya", "persimmon", "acai", "starfruit" | |
| ] | |
| return '-'.join(random.sample(fruits, 3)) | |
| class UI(): | |
| def __init__(self): | |
| self.trainer = Trainer() | |
| def load_loras(self): | |
| loaded_model_name = self.trainer.model_name | |
| if os.path.exists(LORA_DIR) and loaded_model_name is not None: | |
| loras = [f for f in os.listdir(LORA_DIR)] | |
| sanitized_model_name = loaded_model_name.replace('/', '_').replace('.', '_') | |
| loras = [f for f in loras if f.startswith(sanitized_model_name)] | |
| loras.insert(0, 'None') | |
| return gr.Dropdown.update(choices=loras) | |
| else: | |
| return gr.Dropdown.update(choices=['None'], value='None') | |
| def training_params_block(self): | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.max_seq_length = gr.Slider( | |
| interactive=True, | |
| minimum=1, maximum=4096, value=TRAINING_PARAMS['max_seq_length'], | |
| label="Max Sequence Length", | |
| ) | |
| self.micro_batch_size = gr.Slider( | |
| minimum=1, maximum=100, step=1, value=TRAINING_PARAMS['micro_batch_size'], | |
| label="Micro Batch Size", | |
| ) | |
| self.gradient_accumulation_steps = gr.Slider( | |
| minimum=1, maximum=128, step=1, value=TRAINING_PARAMS['gradient_accumulation_steps'], | |
| label="Gradient Accumulation Steps", | |
| ) | |
| self.epochs = gr.Slider( | |
| minimum=1, maximum=100, step=1, value=TRAINING_PARAMS['epochs'], | |
| label="Epochs", | |
| ) | |
| self.learning_rate = gr.Slider( | |
| minimum=0.00001, maximum=0.01, value=TRAINING_PARAMS['learning_rate'], | |
| label="Learning Rate", | |
| ) | |
| with gr.Column(): | |
| self.lora_r = gr.Slider( | |
| minimum=1, maximum=64, step=1, value=LORA_TRAINING_PARAMS['lora_r'], | |
| label="LoRA R", | |
| ) | |
| self.lora_alpha = gr.Slider( | |
| minimum=1, maximum=128, step=1, value=LORA_TRAINING_PARAMS['lora_alpha'], | |
| label="LoRA Alpha", | |
| ) | |
| self.lora_dropout = gr.Slider( | |
| minimum=0, maximum=1, step=0.01, value=LORA_TRAINING_PARAMS['lora_dropout'], | |
| label="LoRA Dropout", | |
| ) | |
| def load_model(self, model_name, progress=gr.Progress(track_tqdm=True)): | |
| if model_name == '': return '' | |
| if model_name is None: return self.trainer.model_name | |
| progress(0, desc=f'Loading {model_name}...') | |
| self.trainer.load_model(model_name) | |
| return self.trainer.model_name | |
| def base_model_block(self): | |
| self.model_name = gr.Dropdown(label='Base Model', choices=MODELS) | |
| def training_data_block(self): | |
| training_text = gr.TextArea( | |
| lines=20, | |
| label="Training Data", | |
| info='Paste training data text here. Sequences must be separated with 2 blank lines' | |
| ) | |
| examples_dir = os.path.join(os.getcwd(), 'example-datasets') | |
| def load_example(filename): | |
| with open(os.path.join(examples_dir, filename) , 'r', encoding='utf-8') as f: | |
| return f.read() | |
| example_filename = gr.Textbox(visible=False) | |
| example_filename.change(fn=load_example, inputs=example_filename, outputs=training_text) | |
| gr.Examples("./example-datasets", inputs=example_filename) | |
| self.training_text = training_text | |
| def training_launch_block(self): | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.new_lora_name = gr.Textbox(label='New PEFT Adapter Name', value=random_name()) | |
| with gr.Column(): | |
| train_button = gr.Button('Train', variant='primary') | |
| def train( | |
| training_text, | |
| new_lora_name, | |
| max_seq_length, | |
| micro_batch_size, | |
| gradient_accumulation_steps, | |
| epochs, | |
| learning_rate, | |
| lora_r, | |
| lora_alpha, | |
| lora_dropout, | |
| progress=gr.Progress(track_tqdm=True) | |
| ): | |
| self.trainer.unload_lora() | |
| self.trainer.train( | |
| training_text, | |
| new_lora_name, | |
| max_seq_length=max_seq_length, | |
| micro_batch_size=micro_batch_size, | |
| gradient_accumulation_steps=gradient_accumulation_steps, | |
| epochs=epochs, | |
| learning_rate=learning_rate, | |
| lora_r=lora_r, | |
| lora_alpha=lora_alpha, | |
| lora_dropout=lora_dropout | |
| ) | |
| return new_lora_name | |
| train_button.click( | |
| fn=train, | |
| inputs=[ | |
| self.training_text, | |
| self.new_lora_name, | |
| self.max_seq_length, | |
| self.micro_batch_size, | |
| self.gradient_accumulation_steps, | |
| self.epochs, | |
| self.learning_rate, | |
| self.lora_r, | |
| self.lora_alpha, | |
| self.lora_dropout, | |
| ], | |
| outputs=[self.new_lora_name] | |
| ).then( | |
| fn=lambda x: self.trainer.load_model(x, force=True), | |
| inputs=[self.model_name], | |
| outputs=[] | |
| ) | |
| def inference_block(self): | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.lora_name = gr.Dropdown( | |
| interactive=True, | |
| choices=['None'], | |
| value='None', | |
| label='LoRA', | |
| ) | |
| def load_lora(lora_name, progress=gr.Progress(track_tqdm=True)): | |
| if lora_name == 'None': | |
| self.trainer.unload_lora() | |
| else: | |
| self.trainer.load_lora(f'{LORA_DIR}/{lora_name}') | |
| return lora_name | |
| self.lora_name.change( | |
| fn=load_lora, | |
| inputs=self.lora_name, | |
| outputs=self.lora_name | |
| ) | |
| self.prompt = gr.Textbox( | |
| interactive=True, | |
| lines=5, | |
| label="Prompt", | |
| value="Human: How is cheese made?\nAssistant:" | |
| ) | |
| self.generate_btn = gr.Button('Generate', variant='primary') | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.max_new_tokens = gr.Slider( | |
| minimum=0, maximum=4096, step=1, value=GENERATION_PARAMS['max_new_tokens'], | |
| label="Max New Tokens", | |
| ) | |
| with gr.Column(): | |
| self.do_sample = gr.Checkbox( | |
| interactive=True, | |
| label="Enable Sampling (leave off for greedy search)", | |
| value=True, | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.num_beams = gr.Slider( | |
| minimum=1, maximum=10, step=1, value=GENERATION_PARAMS['num_beams'], | |
| label="Num Beams", | |
| ) | |
| with gr.Column(): | |
| self.repeat_penalty = gr.Slider( | |
| minimum=0, maximum=4.5, step=0.01, value=GENERATION_PARAMS['repetition_penalty'], | |
| label="Repetition Penalty", | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.temperature = gr.Slider( | |
| minimum=0.01, maximum=1.99, step=0.01, value=GENERATION_PARAMS['temperature'], | |
| label="Temperature", | |
| ) | |
| self.top_p = gr.Slider( | |
| minimum=0, maximum=1, step=0.01, value=GENERATION_PARAMS['top_p'], | |
| label="Top P", | |
| ) | |
| self.top_k = gr.Slider( | |
| minimum=0, maximum=200, step=1, value=GENERATION_PARAMS['top_k'], | |
| label="Top K", | |
| ) | |
| with gr.Column(): | |
| self.output = gr.Textbox( | |
| interactive=True, | |
| lines=20, | |
| label="Output" | |
| ) | |
| def generate( | |
| prompt, | |
| do_sample, | |
| max_new_tokens, | |
| num_beams, | |
| repeat_penalty, | |
| temperature, | |
| top_p, | |
| top_k, | |
| progress=gr.Progress(track_tqdm=True) | |
| ): | |
| return self.trainer.generate( | |
| prompt, | |
| do_sample=do_sample, | |
| max_new_tokens=max_new_tokens, | |
| num_beams=num_beams, | |
| repetition_penalty=repeat_penalty, | |
| temperature=temperature, | |
| top_p=top_p, | |
| top_k=top_k | |
| ) | |
| self.generate_btn.click( | |
| fn=generate, | |
| inputs=[ | |
| self.prompt, | |
| self.do_sample, | |
| self.max_new_tokens, | |
| self.num_beams, | |
| self.repeat_penalty, | |
| self.temperature, | |
| self.top_p, | |
| self.top_k | |
| ], | |
| outputs=[self.output] | |
| ) | |
| def layout(self): | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.HTML("""<h2> | |
| <a style="text-decoration: none;" href="https://github.com/lxe/simple-llama-finetuner">🦙 Simple LLM Finetuner</a> <a href="https://huggingface.co/spaces/lxe/simple-llama-finetuner?duplicate=true"><img | |
| src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" style="display:inline"> | |
| </a></h2><p>Finetune an LLM on your own text. Duplicate this space onto a GPU-enabled space to run.</p>""") | |
| with gr.Column(): | |
| self.base_model_block() | |
| with gr.Tab('Finetuning'): | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.training_data_block() | |
| with gr.Column(): | |
| self.training_params_block() | |
| self.training_launch_block() | |
| with gr.Tab('Inference') as inference_tab: | |
| with gr.Row(): | |
| with gr.Column(): | |
| self.inference_block() | |
| inference_tab.select( | |
| fn=self.load_loras, | |
| inputs=[], | |
| outputs=[self.lora_name] | |
| ) | |
| self.model_name.change( | |
| fn=self.load_model, | |
| inputs=[self.model_name], | |
| outputs=[self.model_name] | |
| ).then( | |
| fn=self.load_loras, | |
| inputs=[], | |
| outputs=[self.lora_name] | |
| ) | |
| return demo | |
| def run(self): | |
| self.ui = self.layout() | |
| self.ui.queue().launch(show_error=True, share=SHARE) | |
| if (__name__ == '__main__'): | |
| ui = UI() | |
| ui.run() | |