File size: 4,670 Bytes
e69b72a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """Causal-LM sequence scoring used by the zero-shot evaluation tasks.
The core primitive returns, for each input sequence, the model's summed (or
length-normalised) log-probability over a scored token region. Grammaticality
(minimal pairs) scores the whole sentence; multiple-choice scores only the
continuation tokens after a shared context.
Batched scoring is exact: sequences are right-padded, and only real,
in-region target positions contribute. Causal attention means real tokens never
attend to right-padding, so per-token scores are identical to unpadded scoring
(verified in tests).
"""
from __future__ import annotations
import torch
def common_prefix_len(a: list[int], b: list[int]) -> int:
n = 0
for x, y in zip(a, b):
if x != y:
break
n += 1
return n
def encode_with_continuation(tokenizer, context_text: str, continuation_text: str) -> tuple[list[int], int]:
"""Tokenise ``context_text + continuation_text`` and return ``(ids, context_len)``.
``context_len`` is the number of leading tokens that belong to the context
(and are therefore *not* scored). Computed as the shared token prefix between
the context alone and the full string, which is robust to subword-merge
effects at the boundary.
"""
context_ids = list(tokenizer.encode(context_text, add_bos=True).input_ids)
full_ids = list(tokenizer.encode(context_text + continuation_text, add_bos=True).input_ids)
context_len = common_prefix_len(context_ids, full_ids)
# Guarantee at least one scored token and a non-empty context.
context_len = max(1, min(context_len, len(full_ids) - 1))
return full_ids, context_len
def _autocast(device_type: str, precision: str):
if precision == "fp32" or device_type == "cpu":
return torch.autocast(device_type=device_type, enabled=False)
dtype = torch.bfloat16 if precision == "bf16" else torch.float16
return torch.autocast(device_type=device_type, dtype=dtype)
@torch.no_grad()
def score_sequences(
model,
sequences: list[list[int]],
context_lens: list[int],
*,
device: torch.device,
pad_id: int = 0,
batch_size: int = 16,
precision: str = "bf16",
length_normalize: bool = False,
predicate_memory_intervention: str = "none",
predicate_memory_residual_scale: float | None = None,
graph_object_residual_scale: float | None = None,
) -> list[float]:
"""Return a score per sequence: summed log-prob over positions ``[context_len, L)``.
With ``length_normalize=True`` the sum is divided by the number of scored
tokens (recommended for comparing continuations of different lengths).
"""
if len(sequences) != len(context_lens):
raise ValueError("sequences and context_lens must have equal length")
model.eval()
device_type = device.type
scores: list[float] = []
for start in range(0, len(sequences), batch_size):
chunk = sequences[start : start + batch_size]
chunk_ctx = context_lens[start : start + batch_size]
max_len = max(len(seq) for seq in chunk)
batch = torch.full((len(chunk), max_len), pad_id, dtype=torch.long)
real_len = torch.empty(len(chunk), dtype=torch.long)
ctx = torch.tensor(chunk_ctx, dtype=torch.long)
for i, seq in enumerate(chunk):
batch[i, : len(seq)] = torch.tensor(seq, dtype=torch.long)
real_len[i] = len(seq)
batch = batch.to(device)
attention_mask = (batch != pad_id).long()
with _autocast(device_type, precision):
logits = model(
batch,
attention_mask=attention_mask,
predicate_memory_intervention=predicate_memory_intervention,
predicate_memory_residual_scale=predicate_memory_residual_scale,
graph_object_residual_scale=graph_object_residual_scale,
).logits
logprobs = torch.log_softmax(logits[:, :-1, :].float(), dim=-1)
targets = batch[:, 1:]
token_lp = logprobs.gather(-1, targets.unsqueeze(-1)).squeeze(-1) # [B, S-1]
# Target at shifted index j predicts absolute position (j+1). Score it
# when context_len <= j+1 < real_len.
positions = torch.arange(1, max_len, device=device).unsqueeze(0) # [1, S-1]
ctx = ctx.to(device).unsqueeze(1)
rl = real_len.to(device).unsqueeze(1)
mask = (positions >= ctx) & (positions < rl)
summed = (token_lp * mask).sum(dim=1)
counts = mask.sum(dim=1).clamp(min=1)
result = summed / counts if length_normalize else summed
scores.extend(result.tolist())
return scores
|