Spaces:
Running
Running
| from unsloth import FastLanguageModel | |
| from trl import SFTTrainer | |
| from transformers import TrainingArguments | |
| from datasets import load_dataset | |
| import torch | |
| def train(): | |
| max_seq_length = 2048 | |
| print("Loading model and tokenizer...") | |
| # Load Model - 4bit quantization to save memory | |
| model, tokenizer = FastLanguageModel.from_pretrained( | |
| model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit", | |
| max_seq_length = max_seq_length, | |
| dtype = None, | |
| load_in_4bit = True, | |
| ) | |
| print("Adding LoRA adapters...") | |
| # Add LoRA | |
| model = FastLanguageModel.get_peft_model( | |
| model, | |
| r = 16, | |
| target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", | |
| "gate_proj", "up_proj", "down_proj",], | |
| lora_alpha = 16, | |
| lora_dropout = 0, | |
| bias = "none", | |
| use_gradient_checkpointing = "unsloth", | |
| random_state = 3407, | |
| use_rslora = False, | |
| loftq_config = None, | |
| ) | |
| print("Loading dataset...") | |
| # Load human-played dataset | |
| dataset = load_dataset("json", data_files="dataset.jsonl", split="train") | |
| # Format the dataset into the prompt structure | |
| def formatting_prompts_func(examples): | |
| prompts = examples["prompt"] | |
| completions = examples["completion"] | |
| texts = [] | |
| for prompt, completion in zip(prompts, completions): | |
| # We add EOS token so the model knows when to stop | |
| text = f"{prompt}{completion}{tokenizer.eos_token}" | |
| texts.append(text) | |
| return { "text" : texts, } | |
| formatted_dataset = dataset.map(formatting_prompts_func, batched = True) | |
| print("Initializing SFT Trainer...") | |
| # Train | |
| trainer = SFTTrainer( | |
| model = model, | |
| tokenizer = tokenizer, | |
| train_dataset = formatted_dataset, | |
| dataset_text_field = "text", | |
| max_seq_length = max_seq_length, | |
| dataset_num_proc = 2, | |
| packing = False, | |
| args = TrainingArguments( | |
| per_device_train_batch_size = 2, | |
| gradient_accumulation_steps = 4, | |
| warmup_steps = 5, | |
| max_steps = 60, # 60 steps for a quick fine-tune demonstration | |
| learning_rate = 2e-4, | |
| fp16 = not torch.cuda.is_bf16_supported(), | |
| bf16 = torch.cuda.is_bf16_supported(), | |
| logging_steps = 1, | |
| optim = "adamw_8bit", | |
| weight_decay = 0.01, | |
| lr_scheduler_type = "linear", | |
| seed = 3407, | |
| output_dir = "sft_output", | |
| ), | |
| ) | |
| print("Starting training loop...") | |
| trainer_stats = trainer.train() | |
| print("Saving the LoRA model...") | |
| model.save_pretrained("tradestorm-sft-model") | |
| tokenizer.save_pretrained("tradestorm-sft-model") | |
| print("✅ Training complete. Model saved to tradestorm-sft-model/") | |
| if __name__ == "__main__": | |
| train() | |