File size: 6,583 Bytes
96a4100 | 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 117 118 119 120 121 122 123 | """Score complete candidate records in one model call."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Callable, Sequence
import torch
@dataclass(frozen=True, slots=True)
class Row:
index: int
candidate_id: str
targets: tuple[int, ...]
positions: tuple[int, ...]
prefix_tokens: int
input_tokens: int
@dataclass(frozen=True, slots=True)
class CompiledJoint:
input_ids: torch.Tensor
attention_mask: torch.Tensor
position_ids: torch.Tensor
logits_to_keep: torch.Tensor
rows: tuple[Row, ...]
candidate_order: tuple[str, ...]
padded_tokens: int
logit_vectors: int
def _contains(values: tuple[int, ...], needle: tuple[int, ...]) -> bool:
return len(needle) <= len(values) and any(
values[i:i + len(needle)] == needle for i in range(len(values) - len(needle) + 1)
)
def compile_joint(prefix: Sequence[int], candidate_ids: Sequence[str],
suffixes: Sequence[Sequence[int]], terminal: Sequence[int], *,
pad_token_id: int, max_rows: int, max_padded_tokens: int,
max_logit_vectors: int, max_input_tokens: int) -> CompiledJoint:
prefix = tuple(prefix); terminal = tuple(terminal)
ids = tuple(candidate_ids); suffixes = tuple(tuple(v) for v in suffixes)
if not prefix or not terminal or any(type(v) is not int or v < 0 for v in prefix + terminal):
raise ValueError("prefix and terminal must contain nonnegative token IDs")
if not ids or len(ids) != len(set(ids)) or any(not isinstance(v, str) or not v for v in ids):
raise ValueError("candidate IDs must be unique nonempty strings")
if len(ids) != len(suffixes) or len(suffixes) != len(set(suffixes)):
raise ValueError("candidate suffixes must be unique and align with candidate IDs")
if any(prefix[-n:] == terminal[:n] for n in range(1, min(len(prefix), len(terminal)) + 1)):
raise ValueError("prefix overlaps the declared terminal boundary")
for suffix in suffixes:
if (len(suffix) < len(terminal) or suffix[-len(terminal):] != terminal
or _contains(suffix[:-len(terminal)], terminal)):
raise ValueError("every candidate must contain the common terminal exactly at its end")
for i, left in enumerate(suffixes):
for j, right in enumerate(suffixes):
if i != j and len(left) < len(right) and right[:len(left)] == left:
raise ValueError("candidate tokenizations prefix-overlap")
if any(any(type(v) is not int or v < 0 for v in suffix) for suffix in suffixes):
raise ValueError("candidate suffixes must contain nonnegative token IDs")
if len(suffixes) > max_rows:
raise ValueError(f"compiled rows {len(suffixes)} exceed max_rows {max_rows}")
sequences = tuple(prefix + suffix for suffix in suffixes)
if max(map(len,sequences)) > max_input_tokens:
raise ValueError(f"joint request exceeds max_input_tokens {max_input_tokens}; no truncation performed")
width = max(map(len, sequences)); padded = len(sequences) * width
if padded > max_padded_tokens:
raise ValueError(f"compiled padded tokens {padded} exceed max_padded_tokens {max_padded_tokens}")
positions = sorted({p for seq in sequences for p in range(len(prefix) - 1, len(seq) - 1)})
vectors = len(sequences) * len(positions)
if vectors > max_logit_vectors:
raise ValueError(f"compiled logit vectors {vectors} exceed max_logit_vectors {max_logit_vectors}")
input_ids = torch.full((len(sequences), width), pad_token_id, dtype=torch.long)
mask = torch.zeros_like(input_ids)
rows = []
for i, (candidate_id, suffix, sequence) in enumerate(zip(ids, suffixes, sequences, strict=True)):
input_ids[i, :len(sequence)] = torch.tensor(sequence)
mask[i, :len(sequence)] = 1
rows.append(Row(i, candidate_id, suffix,
tuple(range(len(prefix) - 1, len(sequence) - 1)), len(prefix), len(sequence)))
position_ids = mask.cumsum(-1) - 1; position_ids.masked_fill_(mask == 0, 0)
return CompiledJoint(input_ids, mask, position_ids, torch.tensor(positions), tuple(rows), ids,
padded, vectors)
def score_joint(model_forward: Callable[..., Any], compiled: CompiledJoint, *, device: Any) -> dict[str, Any]:
expected_positions = compiled.attention_mask.cumsum(-1) - 1
expected_positions.masked_fill_(compiled.attention_mask == 0, 0)
if not torch.equal(compiled.position_ids, expected_positions):
raise ValueError("compiled position IDs mutated before forward")
keep = [int(value) for value in compiled.logits_to_keep.tolist()]
with torch.inference_mode():
output = model_forward(input_ids=compiled.input_ids.to(device),
attention_mask=compiled.attention_mask.to(device),
position_ids=compiled.position_ids.to(device), logits_to_keep=keep,
use_cache=False, return_dict=True)
logits = output.logits
if logits.ndim != 3 or tuple(logits.shape[:2]) != (len(compiled.rows), len(keep)):
raise ValueError("native sparse logits do not align to compiled joint positions")
logits = logits.float()
if not bool(torch.isfinite(logits).all()):
raise ValueError("model logits contain non-finite values")
columns = {int(p): i for i, p in enumerate(compiled.logits_to_keep.tolist())}
scores = []
for row in compiled.rows:
terms = []
for position, target in zip(row.positions, row.targets, strict=True):
distribution = logits[row.index, columns[position]]
if target >= distribution.numel(): raise ValueError("target token is outside vocabulary")
terms.append(distribution[target] - torch.logsumexp(distribution, -1))
scores.append(torch.stack(terms).sum())
scores_t = torch.stack(scores); probs = torch.exp(scores_t - torch.logsumexp(scores_t, 0))
selected = int(torch.argmax(scores_t))
return {"selected_index": selected, "probabilities": probs.cpu().tolist(),
"log_scores": scores_t.cpu().tolist(),
"row_map": [{"row_index": r.index, "candidate_id": r.candidate_id,
"prediction_positions": list(r.positions)} for r in compiled.rows],
"accounting": {"backbone_calls": 1, "physical_rows": len(compiled.rows),
"padded_tokens": compiled.padded_tokens,
"logit_vectors": compiled.logit_vectors, "truncated": False},
"calibrated": False}
|