Update train_lora.py

train_lora.py

@@ -1,47 +1,157 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
-from peft import LoraConfig, get_peft_model
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    TrainingArguments, 
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from datasets import load_dataset
+import torch
+import json
 
 MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+MAX_LENGTH = 512
 
+# Load tokenizer and model
+print("Loading model and tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
 
+# Improved LoRA configuration
 lora_config = LoraConfig(
-    r=8,
-    lora_alpha=16,
-    target_modules=["q_proj", "v_proj"],
-    lora_dropout=0.05,
+    r=16,  # Increased from 8 for better capacity
+    lora_alpha=32,  # Increased from 16
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],  # More modules
+    lora_dropout=0.1,  # Increased for better regularization
     bias="none",
     task_type="CAUSAL_LM"
 )
 
 model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
 
+# Load and split dataset
+print("Loading dataset...")
 dataset = load_dataset("json", data_files="train.jsonl")
 
-def tokenize(example):
-    text = tokenizer.apply_chat_template(
-        example["messages"],
-        tokenize=False
+# Split into train/validation (80/20)
+split_dataset = dataset["train"].train_test_split(test_size=0.2, seed=42)
+train_dataset = split_dataset["train"]
+eval_dataset = split_dataset["test"]
+
+print(f"Training samples: {len(train_dataset)}")
+print(f"Validation samples: {len(eval_dataset)}")
+
+def tokenize_function(examples):
+    """Tokenize the examples with proper formatting"""
+    texts = []
+    for messages in examples["messages"]:
+        # Apply chat template
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False
+        )
+        texts.append(text)
+    
+    # Tokenize with padding and truncation
+    tokenized = tokenizer(
+        texts,
+        truncation=True,
+        max_length=MAX_LENGTH,
+        padding="max_length",
+        return_tensors=None
     )
-    return tokenizer(text, truncation=True)
+    
+    # Labels are the same as input_ids for causal LM
+    tokenized["labels"] = tokenized["input_ids"].copy()
+    
+    return tokenized
 
-tokenized = dataset.map(tokenize)
+# Tokenize datasets
+print("Tokenizing datasets...")
+tokenized_train = train_dataset.map(
+    tokenize_function,
+    batched=True,
+    remove_columns=train_dataset.column_names
+)
 
-args = TrainingArguments(
+tokenized_eval = eval_dataset.map(
+    tokenize_function,
+    batched=True,
+    remove_columns=eval_dataset.column_names
+)
+
+# Improved training arguments
+training_args = TrainingArguments(
     output_dir="./brad-ai-lora",
-    per_device_train_batch_size=1,
-    num_train_epochs=3,
-    learning_rate=2e-4,
-    logging_steps=5,
-    save_strategy="epoch"
+    
+    # Training hyperparameters
+    num_train_epochs=5,  # Increased from 3
+    per_device_train_batch_size=2,  # Increased from 1
+    per_device_eval_batch_size=2,
+    gradient_accumulation_steps=4,  # Effective batch size = 8
+    
+    # Learning rate and scheduling
+    learning_rate=3e-4,  # Slightly increased
+    lr_scheduler_type="cosine",  # Better than default
+    warmup_ratio=0.1,  # Warmup for 10% of training
+    
+    # Optimization
+    optim="adamw_torch",
+    weight_decay=0.01,
+    max_grad_norm=1.0,
+    
+    # Logging and evaluation
+    logging_steps=10,
+    eval_strategy="steps",
+    eval_steps=50,
+    save_strategy="steps",
+    save_steps=50,
+    save_total_limit=3,  # Keep only best 3 checkpoints
+    
+    # Performance
+    fp16=True,  # Mixed precision training
+    dataloader_num_workers=2,
+    
+    # Monitoring
+    load_best_model_at_end=True,
+    metric_for_best_model="eval_loss",
+    greater_is_better=False,
+    
+    # Misc
+    report_to="none",  # Change to "tensorboard" if you want logging
+    seed=42
 )
 
+# Create trainer
 trainer = Trainer(
     model=model,
-    args=args,
-    train_dataset=tokenized["train"]
+    args=training_args,
+    train_dataset=tokenized_train,
+    eval_dataset=tokenized_eval,
+    tokenizer=tokenizer
 )
 
+# Train the model
+print("Starting training...")
 trainer.train()
+
+# Save the final model
+print("Saving model...")
+trainer.save_model("./brad-ai-lora-final")
+tokenizer.save_pretrained("./brad-ai-lora-final")
+
+# Evaluate final model
+print("Final evaluation:")
+eval_results = trainer.evaluate()
+print(eval_results)
+
+print("Training complete!")