import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_from_disk
from trl import SFTTrainer, SFTConfig

def train():
    dataset_path = "/home/jason/Downloads/xiaozhi-mcp/local_processed_dataset"
    model_id = "Qwen/Qwen2.5-3B-Instruct"
    output_dir = "./digital_twin_adapters"

    if not os.path.exists(dataset_path):
        print(f"Error: Processed dataset not found at '{dataset_path}'. Run prepare_dataset.py first.")
        return

    print("Loading dataset from local cache...")
    dataset = load_from_disk(dataset_path)
    train_dataset = dataset["train"]
    eval_dataset = dataset["test"]

    print("Initializing 4-bit quantization config (NF4)...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16
    )

    print(f"Loading Base Model: '{model_id}'...")
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )
    model.config.use_cache = False

    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    model = prepare_model_for_kbit_training(model)

    print("Configuring LoRA Adapter target parameters...")
    peft_config = LoraConfig(
        r=16,
        lora_alpha=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )

    training_args = SFTConfig(
        output_dir=output_dir,
        num_train_epochs=8,                     # Increased epochs to fully bake in details
        per_device_train_batch_size=2,          # Dropped to 2 to safely fit 3B model in 12GB VRAM
        gradient_accumulation_steps=4,          # Keeps an effective batch size of 8
        gradient_checkpointing=True,
        optim="paged_adamw_8bit",
        logging_steps=5,
        learning_rate=2e-4,                     # Higher learning rate to lock in memories
        bf16=True,
        fp16=False,
        max_grad_norm=0.3,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",             # Smoothly decay learning rate
        eval_strategy="no",                     # Disabled evaluation to avoid evaluation-phase OOM
        save_strategy="no",                     # Disabled intermediate checkpoint saving to save disk space
        report_to="none",
        max_length=512,
        packing=False
    )

    print("Building local training harness...")
    trainer = SFTTrainer(
        model=model,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=peft_config,
        processing_class=tokenizer,
        args=training_args
    )

    print("Starting QLoRA Fine-tuning...")
    trainer.train()

    print(f"Saving adapters to: {output_dir}")
    trainer.model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    print("Training successfully finished.")

if __name__ == "__main__":
    train()