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