File size: 4,483 Bytes
df57b75 | 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 | """Option likelihoods under one prefix, with the prefix computed once.
The likelihood readout scores each option's text as a continuation of the same chat-template prompt. The
first implementation re-ran the whole prompt for every option, so a hard item with a 1,400-token prompt and
four options put about 5,600 tokens through the model while the report counted 1,400. Here the prompt goes
through once, its key/value cache is reused for every option, and `processed_tokens` is exactly what the
model saw: the prompt plus the option continuations. The numbers are the same as the uncached computation
(`tests/test_native_scoring.py`); only the work and the honest token count change.
"""
import copy
import math
def score_options(model, prefix_ids, pieces, device, max_tokens, batch=8):
"""(probabilities, processed_tokens) over the options, or (None, widest) when an option would not fit.
Probabilities are the softmax of the length-normalised log-likelihood of each option's tokens given the
prefix. `pieces` are the options' token ids; options of equal length are scored in one batch of `batch`.
"""
import torch
from transformers import DynamicCache
if not pieces or any(not piece for piece in pieces):
raise ValueError('every option needs at least one token')
widest = len(prefix_ids) + max(len(piece) for piece in pieces)
if widest > max_tokens:
return None, widest
totals = [None] * len(pieces)
with torch.no_grad():
prefix = torch.tensor([list(prefix_ids)], device=device)
out = model(input_ids=prefix, past_key_values=DynamicCache(), use_cache=True, logits_to_keep=1)
first = torch.log_softmax(out.logits[0, -1].float(), dim=-1) # predicts each option's first token
cache = out.past_key_values
by_length = {}
for i, piece in enumerate(pieces):
by_length.setdefault(len(piece), []).append(i)
for length, indices in by_length.items():
for start in range(0, len(indices), batch):
chunk = indices[start:start + batch]
ids = torch.tensor([pieces[i] for i in chunk], device=device)
shared = copy.deepcopy(cache)
if len(chunk) > 1:
shared.batch_repeat_interleave(len(chunk))
logprobs = None
if length > 1:
logits = model(input_ids=ids, past_key_values=shared, use_cache=True).logits.float()
logprobs = torch.log_softmax(logits, dim=-1) # position j predicts token j + 1
for row, i in enumerate(chunk):
piece = pieces[i]
total = float(first[piece[0]])
for j in range(1, length):
total += float(logprobs[row, j - 1, piece[j]])
totals[i] = total
means = [t / len(piece) for t, piece in zip(totals, pieces)]
top = max(means)
weights = [math.exp(m - top) for m in means]
processed = len(prefix_ids) + sum(len(piece) for piece in pieces)
return [w / sum(weights) for w in weights], processed
def score_options_uncached(model, prefix_ids, pieces, device, max_tokens, batch=8):
"""The original computation, kept as the reference the cached one is tested against."""
import torch
widest = len(prefix_ids) + max(len(piece) for piece in pieces)
if widest > max_tokens:
return None, widest
totals = [None] * len(pieces)
by_length = {}
for i, piece in enumerate(pieces):
by_length.setdefault(len(piece), []).append(i)
prefix_tensor = torch.tensor(list(prefix_ids))
for length, indices in by_length.items():
for start in range(0, len(indices), batch):
chunk = indices[start:start + batch]
ids = torch.stack([torch.cat([prefix_tensor, torch.tensor(pieces[i])]) for i in chunk]).to(device)
with torch.no_grad():
logits = model(input_ids=ids, use_cache=False, logits_to_keep=length + 1).logits.float()
logprobs = torch.log_softmax(logits, dim=-1)
for row, i in enumerate(chunk):
totals[i] = float(sum(logprobs[row, j, token] for j, token in enumerate(pieces[i])))
means = [t / len(piece) for t, piece in zip(totals, pieces)]
top = max(means)
weights = [math.exp(m - top) for m in means]
return [w / sum(weights) for w in weights], sum(len(prefix_ids) + len(piece) for piece in pieces)
|