| """Native right-padded batches. Supervision is returned separately.""" |
| from __future__ import annotations |
| from rook.render import encode |
|
|
| MODEL_FIELDS = {"input_ids", "attention_mask", "candidate_positions", "candidate_mask", "query_positions"} |
|
|
|
|
| def collate_lists(encoded, pad_id: int, max_length: int): |
| if not encoded or type(pad_id) is not int or pad_id < 0: |
| raise ValueError("empty batch or missing padding token") |
| length = ((max(len(e.ids) for e in encoded) + 7) // 8) * 8 |
| if length > max_length: |
| raise ValueError("padded sequence exceeds explicit limit; no truncation") |
| width = max(len(e.keys) for e in encoded) |
| result = {key: [] for key in MODEL_FIELDS} |
| for e in encoded: |
| n, c = len(e.ids), len(e.keys) |
| if not 2 <= c <= 255 or len(e.candidate_positions) != c or len(set(e.candidate_positions)) != c or not all(0 <= p < e.query_position == n-1 for p in e.candidate_positions): |
| raise ValueError("invalid native endpoints/query") |
| result["input_ids"].append(list(e.ids) + [pad_id] * (length-n)) |
| result["attention_mask"].append([1] * n + [0] * (length-n)) |
| result["candidate_positions"].append(list(e.candidate_positions) + [0] * (width-c)) |
| result["candidate_mask"].append([True] * c + [False] * (width-c)) |
| result["query_positions"].append(e.query_position) |
| return result |
|
|
|
|
| def tensor_batch(inputs, tokenizer, device, max_length=1024): |
| import torch |
| encoded = [encode(item, tokenizer, max_length) for item in inputs] |
| pad = tokenizer.pad_token_id |
| if pad is None: |
| pad = tokenizer.eos_token_id |
| values = collate_lists(encoded, pad, max_length) |
| return {k: torch.tensor(v, dtype=torch.bool if k == "candidate_mask" else torch.long, device=device) for k, v in values.items()}, encoded |
|
|
|
|
| def masked_ce(logits, mask, labels): |
| import torch |
| import torch.nn.functional as F |
| if logits.dtype != torch.float32 or mask.dtype != torch.bool or labels.dtype != torch.long or logits.shape != mask.shape or labels.shape != (logits.shape[0],): |
| raise ValueError("invalid candidate CE tensor contract") |
| if not mask.any(1).all() or (labels < 0).any() or (labels >= logits.shape[1]).any() or not mask.gather(1, labels[:, None]).all(): |
| raise ValueError("label selects an invalid candidate") |
| if not torch.isfinite(logits[mask]).all(): |
| raise FloatingPointError("nonfinite valid candidate logits") |
| loss = F.cross_entropy(logits.masked_fill(~mask, -float("inf")), labels) |
| if not torch.isfinite(loss): |
| raise FloatingPointError("nonfinite candidate CE") |
| return loss |
|
|
