sebdg/unsloth
5
1import gradio as gr2from huggingface_hub import HfApi3from unsloth import FastLanguageModel4from trl import SFTTrainer5from transformers import TrainingArguments, TrainerCallback6from unsloth import is_bfloat16_supported7import torch8from datasets import load_dataset9import logging10from io import StringIO11import time12import asyncio13import psutil14import platform15import os16 17hf_user = None18try:19 hfApi = HfApi()20 hf_user = hfApi.whoami()["name"]21except Exception as e:22 hf_user = "not logged in"23 24def get_human_readable_size(size, decimal_places=2):25 for unit in ['B', 'KB', 'MB', 'GB', 'TB']:26 if size < 1024.0:27 break28 size /= 1024.029 return f"{size:.{decimal_places}f} {unit}"30 31 32# get cpu stats33disk_stats = psutil.disk_usage('.')34print(get_human_readable_size(disk_stats.total))35cpu_info = platform.processor()36print(cpu_info)37os_info = platform.platform()38print(os_info)39 40memory = psutil.virtual_memory()41 42# Dropdown options43model_options = [44 "unsloth/mistral-7b-v0.3-bnb-4bit", # New Mistral v3 2x faster!45 "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",46 "unsloth/llama-3-8b-bnb-4bit", # Llama-3 15 trillion tokens model 2x faster!47 "unsloth/llama-3-8b-Instruct-bnb-4bit",48 "unsloth/llama-3-70b-bnb-4bit",49 "unsloth/Phi-3-mini-4k-instruct", # Phi-3 2x faster!50 "unsloth/Phi-3-medium-4k-instruct",51 "unsloth/mistral-7b-bnb-4bit",52 "unsloth/gemma-2-9b-bnb-4bit",53 "unsloth/gemma-2-9b-bnb-4bit-instruct",54 "unsloth/gemma-2-27b-bnb-4bit", # Gemma 2x faster!55 "unsloth/gemma-2-27b-bnb-4bit-instruct", # Gemma 2x faster!56 "unsloth/Qwen2-1.5B-bnb-4bit", 57 "unsloth/Qwen2-1.5B-bnb-4bit-instruct", 58 "unsloth/Qwen2-7B-bnb-4bit", 59 "unsloth/Qwen2-7B-bnb-4bit-instruct", 60 "unsloth/Qwen2-72B-bnb-4bit", 61 "unsloth/Qwen2-72B-bnb-4bit-instruct", 62 "unsloth/yi-6b-bnb-4bit", 63 "unsloth/yi-34b-bnb-4bit", 64]65gpu_stats = torch.cuda.get_device_properties(0)66start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)67max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)68 69running_on_hf = False70if os.getenv("SYSTEM", None) == "spaces":71 running_on_hf = True72 73system_info = f"""\74- **System:** {os_info}75- **CPU:** {cpu_info} **Memory:** {get_human_readable_size(memory.free)} free of {get_human_readable_size(memory.total)}76- **GPU:** {gpu_stats.name} ({max_memory} GB)77- **Disk:** {get_human_readable_size(disk_stats.free)} free of {get_human_readable_size(disk_stats.total)}78- **Hugging Face:** {running_on_hf}79"""80 81model=None82tokenizer = None83dataset = None84max_seq_length = 204885 86class PrinterCallback(TrainerCallback):87 step = 088 def __init__(self, progress):89 self.progress = progress90 def on_log(self, args, state, control, logs=None, **kwargs):91 _ = logs.pop("total_flos", None)92 if state.is_local_process_zero:93 #print(logs)94 pass95 def on_step_end(self, args, state, control, **kwargs):96 if state.is_local_process_zero:97 self.step = state.global_step98 self.progress(self.step/60, desc=f"Training {self.step}/60")99 #print("**Step ", state.global_step)100 101 102 103def formatting_prompts_func(examples, prompt):104 EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN105 instructions = examples["instruction"]106 inputs = examples["input"]107 outputs = examples["output"]108 texts = []109 for instruction, input, output in zip(instructions, inputs, outputs):110 # Must add EOS_TOKEN, otherwise your generation will go on forever!111 text = prompt.format(instruction, input, output) + EOS_TOKEN112 texts.append(text)113 return { "text" : texts, }114 115def load_model(initial_model_name, load_in_4bit, max_sequence_length):116 global model, tokenizer, max_seq_length117 dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+118 max_seq_length = max_sequence_length119 model, tokenizer = FastLanguageModel.from_pretrained(120 model_name = initial_model_name,121 max_seq_length = max_sequence_length,122 dtype = dtype,123 load_in_4bit = load_in_4bit,124 # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf125 )126 return f"Model {initial_model_name} loaded, using {max_sequence_length} as max sequence length.", gr.update(visible=True, interactive=True), gr.update(interactive=True),gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False)127 128def load_data(dataset_name, data_template_style, data_template):129 global dataset130 dataset = load_dataset(dataset_name, split = "train")131 dataset = dataset.map(lambda examples: formatting_prompts_func(examples, data_template), batched=True)132 return f"Data loaded {len(dataset)} records loaded.", gr.update(visible=True, interactive=True), gr.update(visible=True, interactive=True)133 134def inference(prompt, input_text):135 FastLanguageModel.for_inference(model) # Enable native 2x faster inference136 inputs = tokenizer(137 [138 prompt.format(139 "Continue the fibonnaci sequence.", # instruction140 "1, 1, 2, 3, 5, 8", # input141 "", # output - leave this blank for generation!142 )143 ], return_tensors = "pt").to("cuda")144 145 outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True)146 result = tokenizer.batch_decode(outputs)147 return result[0], gr.update(visible=True, interactive=True)148 149def save_model(model_name, hub_model_name, hub_token, gguf_16bit, gguf_8bit, gguf_4bit, gguf_custom, gguf_custom_value, merge_16bit, merge_4bit, just_lora, push_to_hub, progress=gr.Progress()):150 global model, tokenizer151 152 quants = []153 154 if gguf_custom:155 gguf_custom_value = gguf_custom_value156 quants.append(gguf_custom_value)157 else:158 gguf_custom_value = None159 160 if gguf_16bit:161 quants.append("f16")162 if gguf_8bit:163 quants.append("q8_0")164 if gguf_4bit:165 quants.append("q4_k_m")166 167 if merge_16bit:168 merge = "16bit"169 elif merge_4bit:170 merge = "4bit"171 elif just_lora:172 merge = "lora"173 else:174 merge = None175 176 #model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "f16", token = "")177 if push_to_hub:178 current_quant = 0179 for q in quants:180 progress(current_quant/len(quants), desc=f"Pushing model {model_name} with {q} to HuggingFace Hub")181 model.push_to_hub_gguf(hub_model_name, tokenizer, quantization_method=q, token=hub_token)182 current_quant += 1183 return "Model saved", gr.update(visible=True, interactive=True)184 185def username(profile: gr.OAuthProfile | None):186 hf_user = profile["name"] if profile else "not logged in"187 return hf_user188 189# Create the Gradio interface190with gr.Blocks(title="Unsloth fine-tuning") as demo:191 if (running_on_hf):192 gr.LoginButton()193 # logged_user = gr.Markdown(f"**User:** {hf_user}")194 #demo.load(username, inputs=None, outputs=logged_user)195 with gr.Row():196 with gr.Column(scale=0.5):197 gr.Image("unsloth.png", width="300px", interactive=False, show_download_button=False, show_label=False, show_share_button=False)198 with gr.Column(min_width="550px", scale=1):199 gr.Markdown(system_info) 200 with gr.Column(min_width="250px", scale=0.3):201 gr.Markdown(f"**Links:**\n\n* [Unsloth Hub](https://huggingface.co/unsloth)\n\n* [Unsloth Docs](http://docs.unsloth.com/)\n\n* [Unsloth GitHub](https://github.com/unslothai/unsloth)")202 with gr.Tab("Base Model Parameters"):203 204 with gr.Row():205 initial_model_name = gr.Dropdown(choices=model_options, label="Select Base Model", allow_custom_value=True)206 load_in_4bit = gr.Checkbox(label="Load 4bit model", value=True)207 208 gr.Markdown("### Target Model Parameters")209 with gr.Row():210 max_sequence_length = gr.Slider(minimum=128, value=512, step=64, maximum=128*1024, interactive=True, label="Max Sequence Length")211 load_btn = gr.Button("Load")212 output = gr.Textbox(label="Model Load Status", value="Model not loaded", interactive=False)213 gr.Markdown("---")214 215 with gr.Tab("Data Preparation"):216 with gr.Row():217 dataset_name = gr.Textbox(label="Dataset Name", value="yahma/alpaca-cleaned")218 data_template_style = gr.Dropdown(label="Template", choices=["alpaca","custom"], value="alpaca", allow_custom_value=True)219 with gr.Row():220 data_template = gr.TextArea(label="Data Template", value="""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.221 222### Instruction:223{}224 225### Input:226{}227 228### Response:229{}""")230 gr.Markdown("---")231 output_load_data = gr.Textbox(label="Data Load Status", value="Data not loaded", interactive=False)232 load_data_btn = gr.Button("Load Dataset", interactive=True)233 load_data_btn.click(load_data, inputs=[dataset_name, data_template_style, data_template], outputs=[output_load_data, load_data_btn])234 235 with gr.Tab("Fine-Tuning"):236 gr.Markdown("""### Fine-Tuned Model Parameters""")237 with gr.Row():238 model_name = gr.Textbox(label="Model Name", value=initial_model_name.value, interactive=True)239 240 gr.Markdown("""### Lora Parameters""")241 242 with gr.Row():243 lora_r = gr.Number(label="R", value=16, interactive=True)244 lora_alpha = gr.Number(label="Lora Alpha", value=16, interactive=True)245 lora_dropout = gr.Number(label="Lora Dropout", value=0.1, interactive=True)246 247 gr.Markdown("---")248 gr.Markdown("""### Training Parameters""")249 with gr.Row():250 with gr.Column():251 with gr.Row():252 per_device_train_batch_size = gr.Number(label="Per Device Train Batch Size", value=2, interactive=True)253 warmup_steps = gr.Number(label="Warmup Steps", value=5, interactive=True)254 max_steps = gr.Number(label="Max Steps", value=60, interactive=True)255 gradient_accumulation_steps = gr.Number(label="Gradient Accumulation Steps", value=4, interactive=True)256 with gr.Row():257 logging_steps = gr.Number(label="Logging Steps", value=1, interactive=True)258 log_to_tensorboard = gr.Checkbox(label="Log to Tensorboard", value=True, interactive=True)259 260 with gr.Row():261 # optim = gr.Dropdown(choices=["adamw_8bit", "adamw", "sgd"], label="Optimizer", value="adamw_8bit")262 learning_rate = gr.Number(label="Learning Rate", value=2e-4, interactive=True)263 264 # with gr.Row():265 weight_decay = gr.Number(label="Weight Decay", value=0.01, interactive=True)266 # lr_scheduler_type = gr.Dropdown(choices=["linear", "cosine", "constant"], label="LR Scheduler Type", value="linear")267 gr.Markdown("---")268 269 with gr.Row():270 seed = gr.Number(label="Seed", value=3407, interactive=True)271 output_dir = gr.Textbox(label="Output Directory", value="outputs", interactive=True)272 gr.Markdown("---")273 274 train_output = gr.Textbox(label="Training Status", value="Model not trained", interactive=False)275 train_btn = gr.Button("Train", visible=True)276 277 def train_model(model_name: str, lora_r: int, lora_alpha: int, lora_dropout: float, per_device_train_batch_size: int, warmup_steps: int, max_steps: int,278 gradient_accumulation_steps: int, logging_steps: int, log_to_tensorboard: bool, learning_rate, weight_decay, seed: int, output_dir, progress= gr.Progress()):279 global model, tokenizer280 print(f"$$$ Training model {model_name} with {lora_r} R, {lora_alpha} alpha, {lora_dropout} dropout, {per_device_train_batch_size} per device train batch size, {warmup_steps} warmup steps, {max_steps} max steps, {gradient_accumulation_steps} gradient accumulation steps, {logging_steps} logging steps, {log_to_tensorboard} log to tensorboard, {learning_rate} learning rate, {weight_decay} weight decay, {seed} seed, {output_dir} output dir")281 iseed = seed282 model = FastLanguageModel.get_peft_model(283 model,284 r = lora_r,285 target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",286 "gate_proj", "up_proj", "down_proj",],287 lora_alpha = lora_alpha,288 lora_dropout = lora_dropout,289 bias = "none",290 use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context291 random_state=iseed,292 use_rslora = False, # We support rank stabilized LoRA293 loftq_config = None, # And LoftQ294 )295 progress(0.0, desc="Loading Trainer")296 time.sleep(1)297 trainer = SFTTrainer(298 model = model,299 tokenizer = tokenizer,300 train_dataset = dataset,301 dataset_text_field = "text",302 max_seq_length = max_seq_length,303 dataset_num_proc = 2,304 packing = False, # Can make training 5x faster for short sequences.305 callbacks = [PrinterCallback(progress)],306 args = TrainingArguments(307 per_device_train_batch_size = per_device_train_batch_size,308 gradient_accumulation_steps = gradient_accumulation_steps,309 warmup_steps = warmup_steps,310 max_steps = 60, # Set num_train_epochs = 1 for full training runs311 learning_rate = learning_rate,312 fp16 = not is_bfloat16_supported(),313 bf16 = is_bfloat16_supported(),314 logging_steps = logging_steps,315 optim = "adamw_8bit",316 weight_decay = weight_decay,317 lr_scheduler_type = "linear",318 seed = iseed,319 report_to="tensorboard" if log_to_tensorboard else None,320 output_dir = output_dir321 ),322 )323 trainer.train()324 progress(1, desc="Training completed")325 time.sleep(1)326 return "Model trained 100%",gr.update(visible=True, interactive=False), gr.update(visible=True, interactive=True), gr.update(interactive=True)327 328 329 train_btn.click(train_model, inputs=[model_name, lora_r, lora_alpha, lora_dropout, per_device_train_batch_size, warmup_steps, max_steps, gradient_accumulation_steps, logging_steps, log_to_tensorboard, learning_rate, weight_decay, seed, output_dir], outputs=[train_output, train_btn])330 331 with gr.Tab("Save & Push Options"):332 333 with gr.Row():334 gr.Markdown("### Merging Options")335 with gr.Column():336 merge_16bit = gr.Checkbox(label="Merge to 16bit", value=False, interactive=True)337 merge_4bit = gr.Checkbox(label="Merge to 4bit", value=False, interactive=True)338 just_lora = gr.Checkbox(label="Just LoRA Adapter", value=False, interactive=True)339 gr.Markdown("---")340 341 with gr.Row():342 gr.Markdown("### GGUF Options")343 with gr.Column():344 gguf_16bit = gr.Checkbox(label="Quantize to f16", value=False, interactive=True)345 gguf_8bit = gr.Checkbox(label="Quantize to 8bit (Q8_0)", value=False, interactive=True)346 gguf_4bit = gr.Checkbox(label="Quantize to 4bit (q4_k_m)", value=False, interactive=True)347 with gr.Column():348 gguf_custom = gr.Checkbox(label="Custom", value=False, interactive=True)349 gguf_custom_value = gr.Textbox(label="", value="Q5_K", interactive=True)350 gr.Markdown("---")351 352 with gr.Row():353 gr.Markdown("### Hugging Face Hub Options")354 push_to_hub = gr.Checkbox(label="Push to Hub", value=False, interactive=True)355 with gr.Column():356 hub_model_name = gr.Textbox(label="Hub Model Name", value=f"username/model_name", interactive=True)357 hub_token = gr.Textbox(label="Hub Token", interactive=True, type="password")358 gr.Markdown("---")359 360 # with gr.Row():361 # gr.Markdown("### Ollama options")362 # with gr.Column():363 # ollama_create_local = gr.Checkbox(label="Create in Ollama (local)", value=False, interactive=True)364 # ollama_push_to_hub = gr.Checkbox(label="Push to Ollama", value=False, interactive=True)365 # with gr.Column():366 # ollama_model_name = gr.Textbox(label="Ollama Model Name", value="user/model_name")367 # ollama_pub_key = gr.Button("Ollama Pub Key") 368 save_output = gr.Markdown("---")369 save_button = gr.Button("Save Model", visible=True, interactive=True)370 save_button.click(save_model, inputs=[model_name, hub_model_name, hub_token, gguf_16bit, gguf_8bit, gguf_4bit, gguf_custom, gguf_custom_value, merge_16bit, merge_4bit, just_lora, push_to_hub], outputs=[save_output, save_button])371 372 with gr.Tab("Inference"):373 with gr.Row():374 input_text = gr.Textbox(label="Input Text", lines=4, value="""\375Continue the fibonnaci sequence.376# instruction3771, 1, 2, 3, 5, 8378# input379""", interactive=True)380 output_text = gr.Textbox(label="Output Text", lines=4, value="", interactive=False)381 382 inference_button = gr.Button("Inference", visible=True, interactive=True)383 inference_button.click(inference, inputs=[data_template, input_text], outputs=[output_text, inference_button])384 load_btn.click(load_model, inputs=[initial_model_name, load_in_4bit, max_sequence_length], outputs=[output, load_btn, train_btn, initial_model_name, load_in_4bit, max_sequence_length])385 386demo.launch()387