# train_fft.py

import os
import torch

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig

# ============================================================
# PATHS
# ============================================================

MODEL_PATH = "./Qwen3-0.6B"
DATASET_PATH = "./dataset/train.jsonl"
OUTPUT_DIR = "./outputs/qwen3_0.6b_fft"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ============================================================
# DATASET
# ============================================================

print("Loading dataset...")
dataset = load_dataset(
    "json",
    data_files=DATASET_PATH,
    split="train"
)
print(f"Dataset size: {len(dataset)}")

# ============================================================
# TOKENIZER
# ============================================================

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True
)

# Fix for models that don't have pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# ============================================================
# MODEL
# ============================================================

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto"          # Optional: helps with memory on single GPU
)

model.config.use_cache = False

# ============================================================
# TRAINING CONFIG
# ============================================================

training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,
    learning_rate=5e-6,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    bf16=True,
    logging_steps=10,
    save_strategy="steps",
    save_steps=200,
    save_total_limit=2,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    max_length=512,
    packing=False,
    gradient_checkpointing=True,
    report_to="none",
    # Optional but recommended:
    dataloader_num_workers=2,
    remove_unused_columns=False,
)

# ============================================================
# TRAINER
# ============================================================

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,   # Newer TRL uses processing_class
    # tokenizer=tokenizer,        # You can use this if processing_class doesn't work
)

# ============================================================
# TRAIN
# ============================================================

print("Starting full fine-tuning...")
trainer.train()

# ============================================================
# SAVE MODEL
# ============================================================

print("Saving model...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print("=" * 60)
print("✅ FULL FINE TUNING COMPLETED")
print(f"Model saved to: {OUTPUT_DIR}")
print("=" * 60)