Session 6.5 — Lab notebook

# ~1 min. Colab has torch and transformers; this installs nothing new on most runtimes.
import importlib.util, sys
IN_COLAB = 'google.colab' in sys.modules
if importlib.util.find_spec('transformers') is None:
    %pip install -q transformers

import numpy as np, torch, matplotlib.pyplot as plt
import torch.nn.functional as F
from transformers import (GPT2LMHeadModel, GPT2TokenizerFast,
                          AutoTokenizer, AutoModelForSequenceClassification)

torch.set_num_threads(2)
plt.rcParams.update({'figure.figsize': (7, 4.4), 'axes.grid': True, 'grid.alpha': 0.3})

SEED = 0            # change me in Explore ③
N_PAIRS = 400       # preference pairs the proxy trains on
PROJ_DIM = 48       # proxy capacity: features are projected down to this many dimensions
POOL = 48           # completions sampled per evaluation prompt
N_MAX = 32          # largest best-of-n; POOL must exceed it, see the note in ③
print('Colab:', IN_COLAB)
# Record the environment used for this run; results can vary across versions.
import platform, importlib.metadata, os
print('Python:', platform.python_version(), '| CPUs:', os.cpu_count())
for package in ('torch', 'transformers', 'transformer_lens', 'numpy', 'pandas'):
    try:
        print(package, importlib.metadata.version(package))
    except importlib.metadata.PackageNotFoundError:
        pass
try:
    import psutil
    print('RAM GB:', round(psutil.virtual_memory().total / 1e9, 1))
except ImportError:
    pass
print('GPU available:', __import__('torch').cuda.is_available())

gen_tok = GPT2TokenizerFast.from_pretrained('gpt2'); gen_tok.pad_token = gen_tok.eos_token
gen = GPT2LMHeadModel.from_pretrained('gpt2').eval()          # the policy we sample from

GOLD = 'distilbert-base-uncased-finetuned-sst-2-english'
gold_tok = AutoTokenizer.from_pretrained(GOLD)
gold_m = AutoModelForSequenceClassification.from_pretrained(GOLD).eval()   # the gold scorer

@torch.no_grad()
def sample(prompt, k, new_tokens=16):
    """k sampled completions of one prompt."""
    ids = gen_tok([prompt] * k, return_tensors='pt', padding=True)
    out = gen.generate(**ids, do_sample=True, top_k=50, temperature=1.0,
                       max_new_tokens=new_tokens, pad_token_id=gen_tok.eos_token_id)
    return [gen_tok.decode(o, skip_special_tokens=True) for o in out]

@torch.no_grad()
def gold_score(texts):
    """Unbounded 'true quality': positive logit minus negative logit."""
    b = gold_tok(texts, return_tensors='pt', padding=True, truncation=True, max_length=64)
    lg = gold_m(**b).logits
    return (lg[:, 1] - lg[:, 0]).numpy()

@torch.no_grad()
def features(texts):
    """Frozen mean-pooled GPT-2 hidden states: the proxy never sees the gold model."""
    b = gen_tok(texts, return_tensors='pt', padding=True, truncation=True, max_length=64)
    h = gen.transformer(**b).last_hidden_state
    m = b['attention_mask'].unsqueeze(-1).float()
    return ((h * m).sum(1) / m.sum(1)).numpy()

print('gold scorer:', GOLD)

TRAIN_PROMPTS = ['The film was', 'I thought the restaurant', 'My experience with the product',
                 'The hotel we stayed in', 'Their customer service',
                 'The book I finished last night', 'The new phone', 'Our trip to the coast']

torch.manual_seed(SEED); rng = np.random.default_rng(SEED)
pool = [t for p in TRAIN_PROMPTS for t in sample(p, 16)]
pool_gold = gold_score(pool)
pool_feats = features(pool)

print(f'{len(pool)} completions, gold from {pool_gold.min():+.1f} to {pool_gold.max():+.1f}')
order = np.argsort(pool_gold)
print(f'\nworst:  {pool_gold[order[0]]:+.1f}  {pool[order[0]][:88]!r}')
print(f'best:   {pool_gold[order[-1]]:+.1f}  {pool[order[-1]][:88]!r}')

proj = rng.normal(size=(pool_feats.shape[1], PROJ_DIM)) / np.sqrt(pool_feats.shape[1])
X = torch.tensor(pool_feats @ proj, dtype=torch.float32)

# preference pairs, labelled by gold; skip near-ties, which carry no signal
pairs = []
while len(pairs) < N_PAIRS:
    i, j = rng.choice(len(pool), 2, replace=False)
    if abs(pool_gold[i] - pool_gold[j]) < 0.3:
        continue
    pairs.append((i, j) if pool_gold[i] > pool_gold[j] else (j, i))
pairs = np.array(pairs)
cut = int(0.8 * len(pairs)); train, test = pairs[:cut], pairs[cut:]

reward = torch.nn.Linear(PROJ_DIM, 1, bias=False)
opt = torch.optim.Adam(reward.parameters(), lr=5e-3)
for step in range(120):
    loss = -F.logsigmoid(reward(X[train[:, 0]]).squeeze(-1)
                         - reward(X[train[:, 1]]).squeeze(-1)).mean()
    opt.zero_grad(); loss.backward(); opt.step()

with torch.no_grad():
    acc = (reward(X[test[:, 0]]).squeeze(-1) > reward(X[test[:, 1]]).squeeze(-1)).float().mean().item()
print(f'trained on {len(train)} pairs, final loss {loss.item():.3f}')
print(f'held-out preference accuracy: {acc:.2f}')
print()
print('Expect roughly 0.6 to 0.75, and treat it as optimistic. The pairs are held out but the')
print('completions are not: each appeared in training and only the pairing is new. Holding out completions instead drops this by about 0.11.')

EVAL_PROMPTS = ['The concert last night', 'This laptop', 'The service at the bank',
                'My flight home', 'The coffee shop downstairs', 'That documentary',
                'The software update', 'Their delivery', 'The meal we ordered',
                'This pair of shoes', 'The museum tour', 'My new headphones',
                'The train service', "This week's episode", 'Our landlord', 'The conference talk']

# The slowest cell: 16 prompts x 48 completions = 768 short generations. About six minutes on the standard Colab CPU tested on 11 September 2026; allow longer on a slower runtime.
proxy_pools, gold_pools = [], []
for p in EVAL_PROMPTS:
    outs = []
    for _ in range(POOL // 8):
        outs += sample(p, 8)
    f = torch.tensor(features(outs) @ proj, dtype=torch.float32)
    with torch.no_grad():
        proxy_pools.append(reward(f).squeeze(-1).numpy())
    gold_pools.append(gold_score(outs))
print(f'{len(EVAL_PROMPTS)} prompts x {POOL} completions scored by both models')

NS = [n for n in (1, 2, 4, 8, 16, 32, 64) if n <= N_MAX]
REPEATS = 40

proxy_curve, gold_curve = [], []
for n in NS:
    pv, gv = [], []
    for pr, go in zip(proxy_pools, gold_pools):
        for _ in range(REPEATS):
            subset = rng.choice(len(pr), n, replace=False)
            pick = subset[int(np.argmax(pr[subset]))]      # the proxy chooses
            pv.append(pr[pick]); gv.append(go[pick])       # gold only watches
    proxy_curve.append(np.mean(pv)); gold_curve.append(np.mean(gv))

kl = np.array([np.log(n) - (n - 1) / n for n in NS])
print(f"{'n':>4}  {'KL':>5}  {'proxy':>8}  {'gold':>8}")
for i, n in enumerate(NS):
    print(f'{n:>4}  {kl[i]:>5.2f}  {proxy_curve[i]:>8.2f}  {gold_curve[i]:>8.2f}')
print()
print(f'proxy rose {proxy_curve[-1] - proxy_curve[0]:+.2f} across the range; gold rose {gold_curve[-1] - gold_curve[0]:+.2f}.')
print(f'you optimised the proxy to KL = {kl[-1]:.2f} nats.')

# Gao et al. measure optimisation by d = sqrt(KL), NOT by the KL itself, and fit
# R(d) = d(alpha - beta d) with R(0) := 0. The square root is not decoration: fit a
# downward parabola in the wrong variable and you will invent curvature that is not there.
d = np.sqrt(kl)
delta = np.array(gold_curve) - gold_curve[0]         # R(0) := 0, as in the paper
A = np.stack([d, -d ** 2], axis=1)
(alpha, beta), *_ = np.linalg.lstsq(A, delta, rcond=None)

print(f'fitted alpha = {alpha:.3f}, beta = {beta:+.4f}')
if beta > 0:
    peak_d = alpha / (2 * beta)
    peak_kl = peak_d ** 2
    if peak_kl > 12:                                 # past anything Gao et al. validated
        print(f'predicted peak at d = {peak_d:.2f}, i.e. KL = {peak_kl:.0f} nats. That is far')
        digits = int((peak_kl + 1) / np.log(10)) + 1
        print(f'beyond the 10 nats the published form was checked at, and needs an n of about')
        print(f'{digits} digits. The fit has not located a peak; it has run out of curvature and')
        print('put the vertex somewhere off the map.')
    else:
        n_needed = np.exp(peak_kl + 1)               # KL ~ log n - 1 for large n
        print(f'predicted peak at d = {peak_d:.2f}, i.e. KL = {peak_kl:.1f} nats, n of about {n_needed:,.0f}')
else:
    print('beta came out negative, so the fitted parabola turns the wrong way and has no peak')
    print('at any positive distance. Over the range you covered, gold reward is still rising')
    print('and shows no sign of bending down. That is a result about your data. Report it.')
print()
print(f'You measured out to {kl[-1]:.1f} nats. Gao et al. fitted this form on data to n = 1,000')
print('(about 6 nats) and then validated it at n = 60,000 (about 10 nats). Best-of-n on a CPU')
print('cannot reach that: n = 1,000 here would need about 1,500 generations per prompt.')
print()
print('So the turnover in 6.4 is real and you cannot see it from here. What you can do is fit')
print('the published form to your own data and extrapolate, which Session 2.5 told you to')
print('distrust. Explore ③ asks you to find out how much to distrust it.')

# The picture the numbers describe. Proxy and gold are shifted to start at zero so they
# share an axis; the dashed line is Gao's form with YOUR fitted alpha and beta.
fig, ax = plt.subplots()
ax.plot(d, np.array(proxy_curve) - proxy_curve[0], 'o-', label='proxy (what you optimise)')
ax.plot(d, delta, 's-', label='gold (what you actually want)')
grid = np.linspace(0, d.max() * 1.35, 200)
ax.plot(grid, grid * (alpha - beta * grid), '--', color='grey', label="Gao's form, fitted")
for x, y, nn in zip(d, delta, NS):
    ax.annotate(f'n={nn}', (x, y), textcoords='offset points', xytext=(0, 7), fontsize=8)
ax.set_xlabel(r'$d = \sqrt{\mathrm{KL}}$  (optimisation pressure)')
ax.set_ylabel('score, relative to $n=1$')
ax.set_title('Best-of-$n$: what you optimise against what you want')
ax.legend()
plt.show()

import pandas as pd
MAFAND = ('https://raw.githubusercontent.com/masakhane-io/lafand-mt/'
          'main/data/tsv_files/en-zul/dev.tsv')
try:
    pairs_df = pd.read_csv(MAFAND, sep='\t').dropna().head(120)
except Exception as e:
    print('download failed:', e)
    print('Using a two-row stand-in; the percentages below mean nothing on it.')
    print('Falling back to a few hand-written pairs so the section still runs.')
    pairs_df = pd.DataFrame({'en': ['The service was good.', 'The food was cold.'],
                             'zul': ['Insizakalo ibinhle.', 'Ukudla bekubandayo.']})

# the non-English column is named for the language, so pick it up generically:
# swapping en-zul for en-yor in the URL also needs LANG_NAME below changed
OTHER = [c for c in pairs_df.columns if c != 'en'][0]
LANG_NAME = {'zul': 'isiZulu', 'yor': 'Yoruba', 'swa': 'Kiswahili',
             'hau': 'Hausa', 'xho': 'isiXhosa', 'amh': 'Amharic'}.get(OTHER, OTHER)

@torch.no_grad()
def gold_prob(texts):
    b = gold_tok(texts, return_tensors='pt', padding=True, truncation=True, max_length=64)
    return gold_m(**b).logits.softmax(-1)[:, 1].numpy()

en = gold_prob(list(pairs_df['en'])); zu = gold_prob(list(pairs_df[OTHER]))

# before reading anything into the gap: can the gold model see this language at all?
ids = gold_tok(list(pairs_df[OTHER]), truncation=True, max_length=64)['input_ids']
flat = [t for row in ids for t in row[1:-1]]
unk = flat.count(gold_tok.unk_token_id) / max(1, len(flat))
print(f'{unk:.0%} of {LANG_NAME} tokens are [UNK] to the gold tokeniser')
if unk > 0.5:
    print('  Many tokens are unknown. Inspect tokenised inputs before interpreting scores.')
    print('  Token coverage alone does not establish language comprehension.')
print()
agree = np.mean((en > 0.5) == (zu > 0.5))
conf = lambda p: np.mean((p > 0.9) | (p < 0.1))

print(f'mean score, English : {en.mean():.3f}')
print(f'mean score, {LANG_NAME} : {zu.mean():.3f}')
print(f'same predicted sentiment on parallel sentences: {agree:.0%}')
print(f'confident (>0.9 or <0.1): English {conf(en):.0%}, {LANG_NAME} {conf(zu):.0%}')
print()
flips = np.where((en > 0.5) != (zu > 0.5))[0][:3]
for i in flips:
    print(f"EN {en[i]:.2f}  {pairs_df['en'].iloc[i][:86]}")
    print(f"{OTHER.upper()[:2]} {zu[i]:.2f}  {pairs_df[OTHER].iloc[i][:86]}\n")