RLHF
и Alignment

import torch
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# ── Три модели в памяти при RLHF+PPO ─────────────────────────
# 1. policy: обучаемая модель (π_θ) + value head
policy = AutoModelForCausalLMWithValueHead.from_pretrained("qwen2.5-7b-sft")
# 2. ref_model: замороженная копия SFT (π_ref) — для KL penalty
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained("qwen2.5-7b-sft")
# 3. reward_model: отдельный LLM с linear head → scalar
# (+ value network внутри policy)

config = PPOConfig(
    learning_rate=1.5e-5,
    batch_size=64,
    mini_batch_size=16,
    gradient_accumulation_steps=4,
    optimize_cuda_cache=True,
    # KL penalty coefficient — ключевой гиперпараметр!
    # β слишком мало → reward hacking
    # β слишком велико → не обучается
    kl_penalty="kl",          # "kl" или "abs" или "mse" или "full"
    init_kl_coef=0.2,
    adap_kl_ctrl=True,        # автоматическая адаптация β
    target_kl=6.0,            # целевое значение KL дивергенции
    cliprange=0.2,            # PPO clip: не позволяем большим шагам
    vf_coef=0.1,              # вес value function loss
)

trainer = PPOTrainer(config=config, model=policy, ref_model=ref_model,
                       tokenizer=tokenizer)

for batch in dataset:
    # 1. Генерируем ответ текущей политикой
    query_tensors = batch["input_ids"]
    response_tensors = trainer.generate(query_tensors, max_new_tokens=256)

    # 2. Получаем reward от Reward Model
    texts = tokenizer.batch_decode(response_tensors)
    rewards = [reward_model(q, r) for q, r in zip(batch["queries"], texts)]
    # rewards: список tensor скаляров

    # 3. PPO step: обновляем политику + value head
    # Внутри: вычисляет KL(π_θ||π_ref), clip ratio, advantage
    stats = trainer.step(query_tensors, response_tensors, rewards)
    print(stats["ppo/mean_scores"], stats["ppo/mean_non_score_reward"])
    # mean_non_score_reward ← это и есть KL penalty (отрицательная)

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

MODEL = "Qwen/Qwen2.5-7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model     = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
ref_model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
# ref_model заморожен — служит «якорем» (π_ref)

# ── Датасет: пары (prompt, chosen, rejected) ─────────────────
dataset = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="train_prefs")
# Структура каждого примера:
# {'prompt': 'Explain quantum entanglement...',
#  'chosen': [{'role':'assistant','content':'...хороший ответ...'}],
#  'rejected': [{'role':'assistant','content':'...плохой ответ...'}]}

config = DPOConfig(
    beta=0.1,                  # KL penalty: β=0.1 стандарт
    # Маленький β → модель смелее отходит от ref
    # Большой β → модель консервативнее
    loss_type="sigmoid",        # "sigmoid" (classic) | "ipo" | "kto_pair"
    learning_rate=5e-7,         # DPO: очень маленький LR (меньше чем SFT)
    num_train_epochs=1,         # 1 эпохи обычно достаточно
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    max_length=2048,
    max_prompt_length=512,
    bf16=True,
    logging_steps=10,
    output_dir="./qwen-dpo",
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,     # замороженная копия — для KL
    args=config,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
trainer.train()

# ── Что логируется при обучении ───────────────────────────────
# rewards/chosen:  средний log-prob разница для chosen
# rewards/rejected: средний log-prob разница для rejected
# rewards/margins: chosen_reward - rejected_reward (хотим расти)
# logits/chosen, logits/rejected: неотмасштабированные

from trl import KTOTrainer, KTOConfig
from datasets import Dataset

# Формат KTO: каждый пример = (prompt, completion, label)
# label=True → хороший ответ, label=False → плохой
kto_data = Dataset.from_list([
    {"prompt": "Объясни gradient descent",
     "completion": "Gradient descent — итеративный алгоритм...",
     "label": True},
    {"prompt": "Объясни gradient descent",
     "completion": "Я не могу помочь с этим.",
     "label": False},
    # Не нужно чтобы хорошие и плохие были из одного промпта!
])

config = KTOConfig(
    beta=0.1,
    desirable_weight=1.0,    # вес потерь для хороших примеров
    undesirable_weight=1.0,  # вес потерь для плохих примеров
    # Kahneman-Tversky: люди сильнее реагируют на потери
    # чем на выигрыши (loss aversion). KTO учитывает это.
    max_length=2048,
    learning_rate=5e-7,
)

trainer = KTOTrainer(model=model, ref_model=ref_model,
                       args=config, train_dataset=kto_data)

from trl import GRPOTrainer, GRPOConfig

# ── Reward function: верифицируемые задачи ────────────────────
# Это ключевое в GRPO-Zero: правильность проверяется кодом,
# не reward model. Нет галлюцинации reward, нет reward hacking.

def reward_math(completions: list[str], ground_truth: list[str], **kwargs) -> list[float]:
    """Бинарный reward: правильный ответ = 1.0, неправильный = 0.0."""
    rewards = []
    for completion, gt in zip(completions, ground_truth):
        # Извлекаем ответ из <answer> тегов (DeepSeek-R1 формат)
        import re
        match = re.search(r'<answer>(.*?)</answer>', completion, re.DOTALL)
        if match:
            pred = match.group(1).strip()
            rewards.append(1.0 if pred == gt.strip() else 0.0)
        else:
            rewards.append(0.0)   # нет формата → 0
    return rewards

def reward_format(completions, **kwargs) -> list[float]:
    """Bonus за правильный format с <think> и <answer> тегами."""
    import re
    rewards = []
    for c in completions:
        has_think  = bool(re.search(r'<think>.*?</think>', c, re.DOTALL))
        has_answer = bool(re.search(r'<answer>.*?</answer>', c, re.DOTALL))
        rewards.append(0.5 if has_think and has_answer else 0.0)
    return rewards

config = GRPOConfig(
    num_generations=8,         # G: число ответов на промпт
    max_completion_length=1024,
    learning_rate=5e-7,
    beta=0.04,                  # KL penalty (меньше чем DPO)
    temperature=0.9,           # для генерации G ответов
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
)

trainer = GRPOTrainer(
    model=model,
    reward_funcs=[reward_math, reward_format],   # несколько reward!
    args=config,
    train_dataset=math_dataset,
)
trainer.train()
# GRPO-Zero: точно такой же код но без SFT-преинициализации
# Работает только при чётко верифицируемых задачах (math, code)