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()