Upload app.py

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
 from huggingface_hub import HfApi
-from unsloth import FastLanguageModel
+from unsloth import FastLanguageModel, is_bfloat16_supported
+from unsloth.chat_templates import get_chat_template, train_on_responses_only
+
 from trl import SFTTrainer
-from transformers import TrainingArguments, TrainerCallback
-from unsloth import is_bfloat16_supported
+from transformers import TrainingArguments, TrainerCallback, DataCollatorForSeq2Seq
 import torch
 from datasets import load_dataset
-import logging
-from io import StringIO
 import time
-import asyncio
 import psutil
 import platform
 import os
@@ -41,6 +39,11 @@ memory = psutil.virtual_memory()
 
 # Dropdown options
 model_options = [
+    "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
+    "unsloth/Llama-3.2-1B-bnb-4bit",
+    "unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
+    "unsloth/Llama-3.2-3B-bnb-4bit",
+    "unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
     "unsloth/mistral-7b-v0.3-bnb-4bit",      # New Mistral v3 2x faster!
     "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
     "unsloth/llama-3-8b-bnb-4bit",           # Llama-3 15 trillion tokens model 2x faster!
@@ -98,21 +101,40 @@ class PrinterCallback(TrainerCallback):
             self.progress(self.step/60, desc=f"Training {self.step}/60")
             #print("**Step ", state.global_step)
 
+    
 
 
 def formatting_prompts_func(examples, prompt):
-    EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
+    global tokenizer
     instructions = examples["instruction"]
     inputs       = examples["input"]
     outputs      = examples["output"]
+    
     texts = []
     for instruction, input, output in zip(instructions, inputs, outputs):
-        # Must add EOS_TOKEN, otherwise your generation will go on forever!
-        text = prompt.format(instruction, input, output) + EOS_TOKEN
+        conversation = [
+            {
+                "role": "system",
+                "content": instruction + tokenizer.eos_token
+            },
+            {
+                "role": "user",
+                "content": input + tokenizer.eos_token
+            },
+            {
+                "role": "assistant",
+                "content": output + tokenizer.eos_token
+            }
+        ]
+        text = tokenizer.apply_chat_template(
+            conversation, tokenize=False, add_generation_prompt=False
+        )
+        
         texts.append(text)
-    return { "text" : texts, }
+    
+    return { "text" : texts }
 
-def load_model(initial_model_name, load_in_4bit, max_sequence_length):
+def load_model(initial_model_name, load_in_4bit, max_sequence_length, hub_token):
     global model, tokenizer, max_seq_length
     dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
     max_seq_length = max_sequence_length
@@ -121,7 +143,11 @@ def load_model(initial_model_name, load_in_4bit, max_sequence_length):
         max_seq_length = max_sequence_length,
         dtype = dtype,
         load_in_4bit = load_in_4bit,
-        # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
+        token = f"{hub_token}", # use one if using gated models like meta-llama/Llama-2-7b-hf
+    )
+    tokenizer = get_chat_template(
+        tokenizer,
+        chat_template="llama-3.1",
     )
     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)
 
@@ -129,6 +155,7 @@ def load_data(dataset_name, data_template_style, data_template):
     global dataset
     dataset = load_dataset(dataset_name, split = "train")
     dataset = dataset.map(lambda examples: formatting_prompts_func(examples, data_template), batched=True)
+    
     return f"Data loaded {len(dataset)} records loaded.", gr.update(visible=True, interactive=True), gr.update(visible=True, interactive=True)
 
 def inference(prompt, input_text):
@@ -150,6 +177,7 @@ def save_model(model_name, hub_model_name, hub_token, gguf_16bit, gguf_8bit, ggu
     global model, tokenizer
 
     quants = []
+    current_quant = 0
 
     if gguf_custom:
         gguf_custom_value = gguf_custom_value
@@ -165,15 +193,35 @@ def save_model(model_name, hub_model_name, hub_token, gguf_16bit, gguf_8bit, ggu
         quants.append("q4_k_m")
     
     if merge_16bit:
-        merge = "16bit"
+        progress(current_quant/len(quants), desc=f"Pushing model merged 16bit {model_name} to HuggingFace Hub")
+        model.save_pretrained_merged(
+            "model",
+            tokenizer,
+            save_method="merged_16bit",
+        )
+        if push_to_hub:
+            model.push_to_hub_merged(hub_model_name, tokenizer, save_method="merged_16bit", token=hub_token)
+
     elif merge_4bit:
-        merge = "4bit"
+        progress(current_quant/len(quants), desc=f"Pushing model merged 4bit {model_name} to HuggingFace Hub")
+        model.save_pretrained_merged(
+            "model",
+            tokenizer,
+            save_method="merged_4bit",
+        )
+        if push_to_hub:
+            model.push_to_hub_merged(hub_model_name, tokenizer, save_method="merged_4bit",  token=hub_token)
+
     elif just_lora:
-        merge = "lora"
-    else:
-        merge = None
+        progress(current_quant/len(quants), desc=f"Pushing model merged lora {model_name} to HuggingFace Hub")
+        model.save_pretrained_merged(
+            "model",
+            tokenizer,
+            save_method="lora",
+        )
+        if push_to_hub:
+            model.push_to_hub_merged(hub_model_name, tokenizer, save_method="lora",  token=hub_token)
 
-    #model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "f16", token = "")
     if push_to_hub:
         current_quant = 0
         for q in quants:
@@ -298,8 +346,9 @@ with gr.Blocks(title="Unsloth fine-tuning") as demo:
                 model = model,
                 tokenizer = tokenizer,
                 train_dataset = dataset,
-                dataset_text_field = "text",
-                max_seq_length = max_seq_length,
+                dataset_text_field="text",
+                max_seq_length=max_seq_length,
+                data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
                 dataset_num_proc = 2,
                 packing = False, # Can make training 5x faster for short sequences.
                 callbacks = [PrinterCallback(progress)],
@@ -320,6 +369,11 @@ with gr.Blocks(title="Unsloth fine-tuning") as demo:
                     output_dir = output_dir
                 ),
             )
+            trainer = train_on_responses_only(
+                trainer,
+                instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
+                response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
+            )
             trainer.train()
             progress(1, desc="Training completed")
             time.sleep(1)
@@ -381,6 +435,6 @@ Continue the fibonnaci sequence.
 
         inference_button = gr.Button("Inference", visible=True, interactive=True)
         inference_button.click(inference, inputs=[data_template, input_text], outputs=[output_text, inference_button])
-    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])
+    load_btn.click(load_model, inputs=[initial_model_name, load_in_4bit, max_sequence_length, hub_token], outputs=[output, load_btn, train_btn, initial_model_name, load_in_4bit, max_sequence_length])
 
 demo.launch()