Quantitative evaluation

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()

@torch.no_grad()
def perplexity(text, stride=512, max_len=1024):
    ids = tok(text, return_tensors="pt")["input_ids"]
    nlls, n_tokens = [], 0

    for begin in range(0, ids.size(1), stride):
        end = min(begin + max_len, ids.size(1))
        target_len = end - begin - (max_len - stride if begin else 0)
        chunk = ids[:, begin:end]

        labels = chunk.clone()
        labels[:, :-target_len] = -100      # score only the new tokens

        loss = model(chunk, labels=labels).loss
        nlls.append(loss * target_len)
        n_tokens += target_len

    return torch.exp(torch.stack(nlls).sum() / n_tokens).item()

print(perplexity("The capital of France is Paris. " * 20))   # low: predictable
print(perplexity("qwzx plib fromage tractor helicopter"))    # high: nonsense

# Bits per character — comparable ACROSS tokenizers.
def bits_per_char(text):
    ids = tok(text, return_tensors="pt")["input_ids"]
    loss = model(ids, labels=ids).loss.item()          # nats per token
    return loss * ids.size(1) / (len(text) * 0.6931)   # -> bits per char

import torch, torch.nn.functional as F

logits = torch.tensor([1000.0, 999.0, -1000.0])

print(torch.exp(logits))                       # [inf, inf, 0] — overflow
print(F.softmax(logits, dim=-1))               # fine: max is subtracted

print(torch.log(F.softmax(logits, dim=-1)))    # [..., ..., -inf] — underflow
print(F.log_softmax(logits, dim=-1))           # stable, finite

# Extract the log-prob of specific target tokens — the primitive underneath
# perplexity, HellaSwag and almost every likelihood-based benchmark.
def target_logprobs(logits, targets):
    logp = F.log_softmax(logits, dim=-1)
    return logp.gather(-1, targets.unsqueeze(-1)).squeeze(-1)

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased").eval()

@torch.no_grad()
def top_predictions(sentence, k=5):
    ids = tok(sentence, return_tensors="pt")
    pos = (ids["input_ids"][0] == tok.mask_token_id).nonzero().item()
    logits = model(**ids).logits[0, pos]
    top = logits.topk(k)
    probs = torch.softmax(logits, -1)[top.indices]
    return [(tok.decode([i]), round(p.item(), 3))
            for i, p in zip(top.indices, probs)]

print(top_predictions(f"The capital of France is {tok.mask_token}."))

# The same call, used as a measurement of learned association:
print(top_predictions(f"The nurse said {tok.mask_token} was tired."))
print(top_predictions(f"The engineer said {tok.mask_token} was tired."))
# Comparing the pronoun distributions across matched sentences quantifies
# occupational gender association without any hand-labelling.

import torch, torch.nn.functional as F

@torch.no_grad()
def score_continuation(model, tok, context, continuation):
    """Mean log-probability of the continuation tokens, given the context."""
    ctx = tok(context, return_tensors="pt")["input_ids"]
    full = tok(context + continuation, return_tensors="pt")["input_ids"]

    logits = model(full).logits[0, :-1]        # predicting position i+1
    targets = full[0, 1:]
    logp = F.log_softmax(logits, dim=-1).gather(-1, targets[:, None]).squeeze(-1)

    cont = logp[ctx.size(1) - 1:]              # continuation tokens only
    return cont.mean().item()                  # normalise by length

def hellaswag_accuracy(model, tok, examples):
    correct = 0
    for ex in examples:
        scores = [score_continuation(model, tok, ex["ctx"], e) for e in ex["endings"]]
        correct += int(max(range(len(scores)), key=scores.__getitem__) == ex["label"])
    return correct / len(examples)

# Using .sum() instead of .mean() makes the shortest ending win almost always.
# That single line is the most common bug in home-grown benchmark harnesses.

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,   # compute in bf16, store in 4-bit
    bnb_4bit_quant_type="nf4",               # normal-float 4, better than int4
    bnb_4bit_use_double_quant=True,          # quantise the quantisation constants
)

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/pythia-2.8b", quantization_config=config, device_map="auto")

# Roughly a 4x memory reduction with a small quality cost. Enough to run a
# 7B model on a 12GB card. For evaluation this is usually an acceptable
# trade — but report that you quantised, because it does move the numbers.

import torch, torch.nn.functional as F

@torch.no_grad()
def kl_between_models(model_p, model_q, ids):
    logp = F.log_softmax(model_p(ids).logits, dim=-1)
    logq = F.log_softmax(model_q(ids).logits, dim=-1)
    # KL(P||Q) = sum_x P(x) [log P(x) - log Q(x)]
    return (logp.exp() * (logp - logq)).sum(-1).mean().item()

# PyTorch's version expects log-probabilities for the input and
# probabilities for the target by default — a frequent source of wrong signs.
print(F.kl_div(logq, logp, log_target=True, reduction="batchmean"))

# Asymmetry, demonstrated:
p = torch.tensor([0.5, 0.5, 0.0])
q = torch.tensor([0.5, 0.5, 1e-9])
# KL(p||q) is finite; KL(q||p) blows up, because p rules out what q allows.