| """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) |
| 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) |
| 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) |
|
|
