diff --git a/modules/training.py b/modules/training.py index 656a8b3a..1949fa4e 100644 --- a/modules/training.py +++ b/modules/training.py @@ -7,7 +7,7 @@ from pathlib import Path import gradio as gr import torch import transformers -from datasets import load_dataset +from datasets import Dataset, load_dataset from peft import (LoraConfig, get_peft_model, get_peft_model_state_dict, prepare_model_for_int8_training) @@ -18,8 +18,8 @@ CURRENT_STEPS = 0 MAX_STEPS = 0 CURRENT_GRADIENT_ACCUM = 1 -def get_json_dataset(path: str): - return ['None'] + sorted(set(map(lambda x : '.'.join(str(x.name).split('.')[:-1]), Path(path).glob('*.json'))), key=str.lower) +def get_dataset(path: str, ext: str): + return ['None'] + sorted(set(map(lambda x : '.'.join(str(x.name).split('.')[:-1]), Path(path).glob(f'*.{ext}'))), key=str.lower) def create_train_interface(): with gr.Tab('Train LoRA', elem_id='lora-train-tab'): @@ -40,20 +40,26 @@ def create_train_interface(): lora_dropout = gr.Slider(label='LoRA Dropout', minimum=0.0, maximum=1.0, step=0.025, value=0.05, info='Percentage probability for dropout of LoRA layers.') cutoff_len = gr.Slider(label='Cutoff Length', minimum=0, maximum=2048, value=256, step=32, info='Cutoff length for text input. Essentially, how long of a line of text to feed in at a time. Higher values require drastically more VRAM.') - with gr.Row(): - dataset = gr.Dropdown(choices=get_json_dataset('training/datasets'), value='None', label='Dataset', info='The dataset file to use for training.') - ui.create_refresh_button(dataset, lambda : None, lambda : {'choices': get_json_dataset('training/datasets')}, 'refresh-button') - eval_dataset = gr.Dropdown(choices=get_json_dataset('training/datasets'), value='None', label='Evaluation Dataset', info='The dataset file used to evaluate the model after training.') - ui.create_refresh_button(eval_dataset, lambda : None, lambda : {'choices': get_json_dataset('training/datasets')}, 'refresh-button') - format = gr.Dropdown(choices=get_json_dataset('training/formats'), value='None', label='Data Format', info='The format file used to decide how to format the dataset input.') - ui.create_refresh_button(format, lambda : None, lambda : {'choices': get_json_dataset('training/formats')}, 'refresh-button') + with gr.Tab(label="Formatted Dataset"): + with gr.Row(): + dataset = gr.Dropdown(choices=get_dataset('training/datasets', 'json'), value='None', label='Dataset', info='The dataset file to use for training.') + ui.create_refresh_button(dataset, lambda : None, lambda : {'choices': get_dataset('training/datasets', 'json')}, 'refresh-button') + eval_dataset = gr.Dropdown(choices=get_dataset('training/datasets', 'json'), value='None', label='Evaluation Dataset', info='The dataset file used to evaluate the model after training.') + ui.create_refresh_button(eval_dataset, lambda : None, lambda : {'choices': get_dataset('training/datasets', 'json')}, 'refresh-button') + format = gr.Dropdown(choices=get_dataset('training/formats', 'json'), value='None', label='Data Format', info='The format file used to decide how to format the dataset input.') + ui.create_refresh_button(format, lambda : None, lambda : {'choices': get_dataset('training/formats', 'json')}, 'refresh-button') + with gr.Tab(label="Raw Text File"): + with gr.Row(): + raw_text_file = gr.Dropdown(choices=get_dataset('training/datasets', 'txt'), value='None', label='Text File', info='The raw text file to use for training.') + ui.create_refresh_button(raw_text_file, lambda : None, lambda : {'choices': get_dataset('training/datasets', 'txt')}, 'refresh-button') + overlap_len = gr.Slider(label='Overlap Length', minimum=0,maximum=512, value=32, step=8, info='Overlap length - ie how many tokens from the prior chunk of text to include into the next chunk. (The chunks themselves will be of a size determined by Cutoff Length above)') with gr.Row(): start_button = gr.Button("Start LoRA Training") stop_button = gr.Button("Interrupt") output = gr.Markdown(value="Ready") - start_button.click(do_train, [lora_name, micro_batch_size, batch_size, epochs, learning_rate, lora_rank, lora_alpha, lora_dropout, cutoff_len, dataset, eval_dataset, format], [output]) + start_button.click(do_train, [lora_name, micro_batch_size, batch_size, epochs, learning_rate, lora_rank, lora_alpha, lora_dropout, cutoff_len, dataset, eval_dataset, format, raw_text_file, overlap_len], [output]) stop_button.click(do_interrupt, [], [], cancels=[], queue=False) def do_interrupt(): @@ -84,7 +90,8 @@ def clean_path(base_path: str, path: str): return path return f'{Path(base_path).absolute()}/{path}' -def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int, learning_rate: float, lora_rank: int, lora_alpha: int, lora_dropout: float, cutoff_len: int, dataset: str, eval_dataset: str, format: str): +def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int, learning_rate: str, lora_rank: int, + lora_alpha: int, lora_dropout: float, cutoff_len: int, dataset: str, eval_dataset: str, format: str, raw_text_file: str, overlap_len: int): global WANT_INTERRUPT, CURRENT_STEPS, MAX_STEPS, CURRENT_GRADIENT_ACCUM WANT_INTERRUPT = False CURRENT_STEPS = 0 @@ -93,20 +100,17 @@ def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int # == Input validation / processing == yield "Prepping..." lora_name = f"{shared.args.lora_dir}/{clean_path(None, lora_name)}" - if dataset is None: - return "**Missing dataset choice input, cannot continue.**" - if format is None: - return "**Missing format choice input, cannot continue.**" + actual_lr = float(learning_rate) + + if cutoff_len <= 0 or micro_batch_size <= 0 or batch_size <= 0 or actual_lr <= 0 or lora_rank <= 0 or lora_alpha <= 0: + yield f"Cannot input zeroes." + return + gradient_accumulation_steps = batch_size // micro_batch_size CURRENT_GRADIENT_ACCUM = gradient_accumulation_steps - actual_lr = float(learning_rate) shared.tokenizer.pad_token = 0 shared.tokenizer.padding_side = "left" - # == Prep the dataset, format, etc == - with open(clean_path('training/formats', f'{format}.json'), 'r') as formatFile: - format_data: dict[str, str] = json.load(formatFile) - def tokenize(prompt): result = shared.tokenizer(prompt, truncation=True, max_length=cutoff_len + 1, padding="max_length") return { @@ -114,27 +118,55 @@ def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int "attention_mask": result["attention_mask"][:-1], } - def generate_prompt(data_point: dict[str, str]): - for options, data in format_data.items(): - if set(options.split(',')) == set(x[0] for x in data_point.items() if len(x[1].strip()) > 0): - for key, val in data_point.items(): - data = data.replace(f'%{key}%', val) - return data - raise RuntimeError(f'Data-point "{data_point}" has no keyset match within format "{list(format_data.keys())}"') - - def generate_and_tokenize_prompt(data_point): - prompt = generate_prompt(data_point) - return tokenize(prompt) - - print("Loading datasets...") - data = load_dataset("json", data_files=clean_path('training/datasets', f'{dataset}.json')) - train_data = data['train'].shuffle().map(generate_and_tokenize_prompt) - - if eval_dataset == 'None': + # == Prep the dataset, format, etc == + if raw_text_file is not None: + print("Loading raw text file dataset...") + with open(clean_path('training/datasets', f'{raw_text_file}.txt'), 'r') as file: + raw_text = file.read() + tokens = shared.tokenizer.encode(raw_text) + del raw_text # Note: could be a gig for a large dataset, so delete redundant data as we go to be safe on RAM + tokens = list(split_chunks(tokens, cutoff_len - overlap_len)) + for i in range(1, len(tokens)): + tokens[i] = tokens[i - 1][-overlap_len:] + tokens[i] + text_chunks = [shared.tokenizer.decode(x) for x in tokens] + del tokens + data = Dataset.from_list([tokenize(x) for x in text_chunks]) + train_data = data.shuffle() eval_data = None + del text_chunks + else: - eval_data = load_dataset("json", data_files=clean_path('training/datasets', f'{eval_dataset}.json')) - eval_data = eval_data['train'].shuffle().map(generate_and_tokenize_prompt) + with open(clean_path('training/formats', f'{format}.json'), 'r') as formatFile: + format_data: dict[str, str] = json.load(formatFile) + + if dataset is None: + yield "**Missing dataset choice input, cannot continue.**" + return + if format is None: + yield "**Missing format choice input, cannot continue.**" + return + + def generate_prompt(data_point: dict[str, str]): + for options, data in format_data.items(): + if set(options.split(',')) == set(x[0] for x in data_point.items() if len(x[1].strip()) > 0): + for key, val in data_point.items(): + data = data.replace(f'%{key}%', val) + return data + raise RuntimeError(f'Data-point "{data_point}" has no keyset match within format "{list(format_data.keys())}"') + + def generate_and_tokenize_prompt(data_point): + prompt = generate_prompt(data_point) + return tokenize(prompt) + + print("Loading JSON datasets...") + data = load_dataset("json", data_files=clean_path('training/datasets', f'{dataset}.json')) + train_data = data['train'].shuffle().map(generate_and_tokenize_prompt) + + if eval_dataset == 'None': + eval_data = None + else: + eval_data = load_dataset("json", data_files=clean_path('training/datasets', f'{eval_dataset}.json')) + eval_data = eval_data['train'].shuffle().map(generate_and_tokenize_prompt) # == Start prepping the model itself == if not hasattr(shared.model, 'lm_head') or hasattr(shared.model.lm_head, 'weight'): @@ -229,3 +261,7 @@ def do_train(lora_name: str, micro_batch_size: int, batch_size: int, epochs: int else: print("Training complete!") yield f"Done! LoRA saved to `{lora_name}`" + +def split_chunks(arr, step): + for i in range(0, len(arr), step): + yield arr[i:i + step]