RLHF

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

class RewardModel(nn.Module):
    """
    Reward Model: base LLM + một head tuyến tính cho ra scalar reward.
    """

    def __init__(self, base_model_name: str):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(base_model_name)
        hidden = self.backbone.config.hidden_size
        self.value_head = nn.Linear(hidden, 1, bias=False)

    def forward(self, input_ids, attention_mask):
        out = self.backbone(input_ids, attention_mask=attention_mask)
        last_hidden = out.last_hidden_state
        # Lấy hidden state của token cuối cùng (cuối response)
        lengths = attention_mask.sum(dim=1) - 1
        last = last_hidden[torch.arange(last_hidden.size(0)), lengths]
        reward = self.value_head(last).squeeze(-1)  # (batch,)
        return reward


def rm_loss(reward_chosen, reward_rejected):
    """Bradley-Terry loss: khuyến khích reward(chosen) > reward(rejected)."""
    return -F.logsigmoid(reward_chosen - reward_rejected).mean()


# Training loop
def train_step(rm, tokenizer, batch, optimizer):
    # batch gồm danh sách (prompt, chosen, rejected)
    chosen_texts = [p + c for p, c, _ in batch]
    rejected_texts = [p + r for p, _, r in batch]

    chosen_enc = tokenizer(chosen_texts, padding=True, truncation=True,
                           return_tensors="pt").to("cuda")
    rejected_enc = tokenizer(rejected_texts, padding=True, truncation=True,
                             return_tensors="pt").to("cuda")

    r_chosen = rm(**chosen_enc)
    r_rejected = rm(**rejected_enc)

    loss = rm_loss(r_chosen, r_rejected)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item(), (r_chosen > r_rejected).float().mean().item()

# Dùng thư viện TRL (Hugging Face) — abstraction phổ biến cho RLHF.
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer
import torch

# 1. Load policy (π_θ) — khởi tạo từ π_SFT
policy = AutoModelForCausalLMWithValueHead.from_pretrained("sft-model")
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained("sft-model")
ref_model.eval()  # ref_model = π_SFT, đóng băng để tính KL

tokenizer = AutoTokenizer.from_pretrained("sft-model")
reward_model = load_reward_model("rm-model")

# 2. Cấu hình PPO
config = PPOConfig(
    model_name="sft-model",
    learning_rate=1e-5,
    batch_size=64,
    mini_batch_size=8,
    ppo_epochs=4,
    init_kl_coef=0.05,     # β ban đầu
    adap_kl_ctrl=True,     # adaptive KL control
    target=6.0,            # target KL (nats)
    cliprange=0.2,         # PPO clip range
    cliprange_value=0.2,
    vf_coef=0.1,
    gamma=1.0,
    lam=0.95,
)

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

# 3. Training loop
for epoch in range(num_epochs):
    for batch in prompt_dataloader:
        prompts = batch["prompt"]
        prompt_tensors = [tokenizer.encode(p, return_tensors="pt")[0] for p in prompts]

        # (a) Policy sinh response
        response_tensors = [
            trainer.generate(p, max_new_tokens=200, do_sample=True, top_p=0.9)[0]
            for p in prompt_tensors
        ]
        responses = [tokenizer.decode(r) for r in response_tensors]

        # (b) Reward model chấm
        rewards = [
            torch.tensor(reward_model.score(p, r))
            for p, r in zip(prompts, responses)
        ]

        # (c) PPO update — tự động áp dụng KL penalty với ref_model
        stats = trainer.step(prompt_tensors, response_tensors, rewards)

        # (d) Log
        print(f"reward={stats['ppo/mean_scores']:.3f}, "
              f"KL={stats['objective/kl']:.2f}, "
              f"β={stats['objective/kl_coef']:.4f}")