from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset

# ===========================
# Reward Model = classifier: (prompt, response) → score
# Trained on human preferences: score(chosen) > score(rejected)
# ===========================

# 1. Load preference dataset
# Format: {"prompt": "...", "chosen": "good response", "rejected": "bad response"}
dataset = load_dataset("Anthropic/hh-rlhf", split="train[:5000]")

# 2. Load base model as reward model
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=1)  # 1 output = scalar score
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

# 3. Train: reward(chosen) > reward(rejected)
# Loss = -log(sigmoid(reward(chosen) - reward(rejected)))
# → Bradley-Terry model of preferences

# trainer = RewardTrainer(
#     model=model, args=RewardConfig(...),
#     train_dataset=dataset, tokenizer=tokenizer)
# trainer.train()

# CATATAN: Di praktik, SKIP reward model dan gunakan DPO langsung!
# DPO = implicit reward model (lebih sederhana, lebih stabil)

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# ===========================
# PPO for RLHF — conceptual code
# ===========================
# PPO needs 3-4 models in VRAM simultaneously:
# 1. Policy model (LLM being trained)
# 2. Reference model (frozen copy — for KL penalty)
# 3. Reward model (scores responses)
# 4. Value head (estimates future rewards)
# → VERY expensive! 4× VRAM of single model

# config = PPOConfig(
#     model_name="gpt2",
#     learning_rate=1e-5,
#     batch_size=16,
#     mini_batch_size=4,
#     ppo_epochs=4,           # PPO optimization epochs per batch
#     kl_penalty="kl",        # KL divergence from reference model
#     target_kl=0.1,          # target KL (stops if exceeded)
# )

# PPO Training Loop:
# for batch in dataloader:
#     1. Generate responses: policy_model(prompts) → responses
#     2. Score responses: reward_model(prompts, responses) → rewards
#     3. Compute KL penalty: KL(policy || reference) → penalty
#     4. Adjusted reward = reward - β × KL_penalty
#     5. PPO update: maximize adjusted reward

# MASALAH PPO:
# ❌ Butuh 3-4 model di VRAM → sangat mahal
# ❌ Tidak stabil (reward hacking, mode collapse)
# ❌ Banyak hyperparameters (KL target, clip ratio, etc.)
# ❌ Sulit di-reproduce
# → SOLUSI: DPO! (Section 6)

# pip install trl
# TRL = Transformer Reinforcement Learning
# By Hugging Face — the standard library for LLM alignment

from trl import (
    SFTTrainer,      # Supervised Fine-Tuning (Page 8)
    RewardTrainer,   # Train reward model
    PPOTrainer,      # RLHF with PPO
    DPOTrainer,      # Direct Preference Optimization ⭐
    ORPOTrainer,     # Odds Ratio Preference Optimization (newest)
    # KTOTrainer,    # Kahneman-Tversky Optimization
)

# Training pipeline:
# Step 1: SFTTrainer → instruction-following model (Page 8 ✅ done!)
# Step 2: DPOTrainer → aligned model (Page 9 this! ⭐)
# That's it! 2 steps = ChatGPT-like assistant.

#!/usr/bin/env python3
"""
🎯 DPO Fine-Tuning: Align LLM with Human Preferences
Uses QLoRA (Page 8) + DPO → ChatGPT-like alignment on Colab T4!
"""

# !pip install -q transformers datasets accelerate peft bitsandbytes trl

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig
from trl import DPOTrainer, DPOConfig
from datasets import load_dataset

# ═══════════════════════════════════════
# STEP 1: LOAD MODEL (4-bit, same as Page 8!)
# ═══════════════════════════════════════
model_name = "mistralai/Mistral-7B-v0.3"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True)

model = AutoModelForCausalLM.from_pretrained(
    model_name, quantization_config=bnb_config, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# Reference model = same base (DPO needs it for KL penalty)
# DPOTrainer automatically creates reference model internally!

# ═══════════════════════════════════════
# STEP 2: LOAD PREFERENCE DATASET
# ═══════════════════════════════════════
# Format: {"prompt": "...", "chosen": "good response", "rejected": "bad response"}
dataset = load_dataset("Intel/orca_dpo_pairs", split="train[:5000]")

# Inspect
print(dataset[0].keys())  # dict_keys(['system', 'question', 'chosen', 'rejected'])
print(f"Chosen:   {dataset[0]['chosen'][:100]}...")
print(f"Rejected: {dataset[0]['rejected'][:100]}...")

# Format for DPOTrainer
def format_dpo(example):
    return {
        "prompt": f"### Instruction:\n{example['question']}\n\n### Response:\n",
        "chosen": example["chosen"],
        "rejected": example["rejected"],
    }

dataset = dataset.map(format_dpo)

# ═══════════════════════════════════════
# STEP 3: CONFIGURE LORA + DPO
# ═══════════════════════════════════════
peft_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none", task_type="CAUSAL_LM")

dpo_config = DPOConfig(
    output_dir="./dpo-mistral",
    num_train_epochs=1,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=5e-5,          # lower than SFT (more careful!)
    beta=0.1,                    # DPO temperature (0.1-0.5)
    # β controls how much to deviate from reference model
    # β=0.1 → stay close to reference (conservative)
    # β=0.5 → deviate more (aggressive optimization)
    fp16=True,
    gradient_checkpointing=True,
    optim="paged_adamw_8bit",
    max_length=512,
    max_prompt_length=256,
    logging_steps=25,
    save_steps=200,
    report_to="none",
)

# ═══════════════════════════════════════
# STEP 4: TRAIN WITH DPO!
# ═══════════════════════════════════════
trainer = DPOTrainer(
    model=model,
    args=dpo_config,
    train_dataset=dataset,
    processing_class=tokenizer,
    peft_config=peft_config,
    # ref_model=None → auto-created from base model!
)

print("🎯 Training DPO...")
trainer.train()
trainer.save_model("./dpo-mistral-final")

# Test: model should now prefer helpful, safe responses!
print("\\n🏆 DPO alignment complete!")

from trl import ORPOTrainer, ORPOConfig

# ===========================
# ORPO (Odds Ratio Preference Optimization, 2024)
# Combines SFT + DPO into ONE training step!
# → No need for separate SFT then DPO
# → Simpler pipeline, competitive results
# ===========================

# Same dataset format as DPO: prompt, chosen, rejected
# trainer = ORPOTrainer(
#     model=model, args=ORPOConfig(beta=0.1, ...),
#     train_dataset=dataset, peft_config=peft_config)
# trainer.train()

# Evolution of alignment methods:
# 2022: RLHF/PPO (OpenAI)    → complex, 3-4 models
# 2023: DPO (Stanford)        → simple, 2 models ← RECOMMENDED
# 2024: ORPO (KAIST)          → simpler, SFT+DPO in 1 step
# 2024: SimPO (Princeton)     → even simpler, no reference model
# Trend: semakin sederhana, semakin efisien, hasil tetap bagus!