Instruction tuning and RLHF

# The template is just text with role markers. Any consistent scheme works,
# but it must be identical at training and inference or behaviour collapses.
TEMPLATE = (
    "<|system|>\n{system}<|end|>\n"
    "<|user|>\n{user}<|end|>\n"
    "<|assistant|>\n{assistant}<|end|>"
)

def build_example(tok, system, user, assistant, max_len=512):
    prompt = TEMPLATE.format(system=system, user=user, assistant="")
    full   = TEMPLATE.format(system=system, user=user, assistant=assistant)

    prompt_ids = tok(prompt)["input_ids"]
    full_ids   = tok(full, truncation=True, max_length=max_len)["input_ids"]

    # -100 is CrossEntropyLoss's ignore_index: those positions contribute
    # nothing. Masking the prompt means the model is trained to produce the
    # response, not to reproduce the question it was asked.
    labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids):]
    return {"input_ids": full_ids, "labels": labels[: len(full_ids)]}

# Modern tokenizers ship their own template — use it rather than inventing one:
# tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import DataLoader

tok = AutoTokenizer.from_pretrained("gpt2-large")
tok.add_special_tokens({"additional_special_tokens":
                        ["<|system|>", "<|user|>", "<|assistant|>", "<|end|>"]})
tok.pad_token = tok.eos_token

model = AutoModelForCausalLM.from_pretrained("gpt2-large")
model.resize_token_embeddings(len(tok))    # required after adding tokens

def collate(batch):
    n = max(len(b["input_ids"]) for b in batch)
    pad = lambda seq, fill: seq + [fill] * (n - len(seq))
    return {
        "input_ids": torch.tensor([pad(b["input_ids"], tok.pad_token_id) for b in batch]),
        "labels":    torch.tensor([pad(b["labels"], -100) for b in batch]),
        "attention_mask": torch.tensor(
            [[1] * len(b["input_ids"]) + [0] * (n - len(b["input_ids"])) for b in batch]),
    }

loader = DataLoader(examples, batch_size=4, shuffle=True, collate_fn=collate)
opt = torch.optim.AdamW(model.parameters(), lr=2e-5)

model.train()
for batch in loader:
    loss = model(**batch).loss
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

import torch, torch.nn.functional as F

# Stage 2. Bradley-Terry pairwise loss: the chosen response should outscore
# the rejected one. Only the difference matters, so the scale is arbitrary.
def reward_loss(reward_model, prompt, chosen, rejected):
    r_c = reward_model(prompt, chosen)      # scalar per pair
    r_r = reward_model(prompt, rejected)
    return -F.logsigmoid(r_c - r_r).mean()

# Stage 3. Maximise reward, penalised for straying from the reference policy.
# Without the KL term the policy collapses onto whatever the reward model
# over-scores — reward hacking, and it happens quickly.
def ppo_objective(reward, logprobs_policy, logprobs_ref, beta=0.1):
    kl = (logprobs_policy - logprobs_ref).mean()
    return reward.mean() - beta * kl

# DPO removes the separate reward model entirely, optimising the preference
# data directly. Same data, one training stage, far simpler to run — which is
# why most open work now uses it.
def dpo_loss(pi_chosen, pi_rejected, ref_chosen, ref_rejected, beta=0.1):
    margin = beta * ((pi_chosen - ref_chosen) - (pi_rejected - ref_rejected))
    return -F.logsigmoid(margin).mean()