| """Turn rows of a DecisionData into model-ready numpy batches (+ torch loaders). |
| |
| Questions and choice labels are tokenised once per *vocabulary entry* (a dataset |
| has few distinct questions/choices). With a frozen pretrained backbone a |
| `BackboneTable` goes further and stores the backbone output per vocabulary entry, |
| so training never re-runs the text encoder on questions or choices. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from minijev.tokenizers import ByteTokenizer |
| from minijev.tasks import Task |
|
|
| from .decision_data import DecisionData |
|
|
| NEG = -1e9 |
|
|
|
|
| class BackboneTable: |
| """Frozen-backbone outputs for a list of strings: tok (V, L, dt) fp16, mask (V, L), pooled (V, dt).""" |
|
|
| def __init__(self, texts: list[str], model, tokenizer, max_len: int, batch: int = 256): |
| import torch |
|
|
| dev = next(model.parameters()).device |
| toks, masks, pooled = [], [], [] |
| L = 1 |
| with torch.inference_mode(): |
| for i in range(0, len(texts), batch): |
| ids = torch.as_tensor(tokenizer.encode_batch(texts[i:i + batch], max_len, dynamic=True), device=dev) |
| t, m, p = model.backbone(ids) |
| toks.append(t.half().cpu().numpy()) |
| masks.append(m.cpu().numpy()) |
| pooled.append(p.float().cpu().numpy()) |
| L = max(L, t.shape[1]) |
| dt = pooled[0].shape[1] if pooled else 1 |
| self.tok = np.zeros((len(texts) + 1, L, dt), np.float16) |
| self.mask = np.zeros((len(texts) + 1, L), bool) |
| k = 0 |
| for t, m in zip(toks, masks): |
| self.tok[k:k + len(t), : t.shape[1]] = t |
| self.mask[k:k + len(t), : m.shape[1]] = m |
| k += len(t) |
| self.pooled = np.concatenate(pooled + [np.zeros((1, dt), np.float32)]) |
|
|
|
|
| class BatchEncoder: |
| def __init__(self, data: DecisionData, task: Task, max_question_len: int = 32, max_choice_len: int = 16, |
| choice_dropout: float = 0.0, seed: int = 0, tokenizer=None, |
| question_table: BackboneTable | None = None, choice_table: BackboneTable | None = None): |
| self.data, self.task = data, task |
| self.tok = tokenizer or getattr(task, "tokenizer", None) or ByteTokenizer() |
| self.dynamic = self.tok.kind != "byte" |
| self.q_tok = self.tok.encode_batch(data.questions, max_question_len, dynamic=self.dynamic) |
| labels = [task.model_label(c) for c in data.choices] |
| c_tok = self.tok.encode_batch(labels, max_choice_len, dynamic=self.dynamic) |
| |
| self.c_tok = np.concatenate([c_tok, np.full((1, c_tok.shape[1]), self.tok.pad_id, np.int64)]) |
| self.q_table, self.c_table = question_table, choice_table |
| self.choice_dropout = choice_dropout |
| self.seed = seed |
|
|
| def _trim(self, ids: np.ndarray) -> np.ndarray: |
| if not self.dynamic: |
| return ids |
| n = int((ids != self.tok.pad_id).any(tuple(range(ids.ndim - 1))).nonzero()[0].max(initial=0)) + 1 |
| return ids[..., :n] |
|
|
| def __call__(self, idx: np.ndarray, rng: np.random.Generator | None = None) -> dict[str, np.ndarray]: |
| d = self.data |
| idx = np.asarray(idx) |
| batch = dict(self.task.encode_arrays(d.state, idx)) |
| ci = d.choice_idx[idx].astype(np.int64) |
| present = ci >= 0 |
| target = d.target[idx].astype(np.int64) |
| soft = d.soft[idx].astype(np.float32) if d.soft is not None else None |
| teacher = d.teacher_logits[idx].astype(np.float32) if d.teacher_logits is not None else None |
| if self.choice_dropout > 0 and rng is not None: |
| |
| drop = (rng.random(ci.shape) < self.choice_dropout) & present |
| drop[np.arange(len(idx)), target.clip(0)] &= target < 0 |
| keep_count = (present & ~drop).sum(1) |
| drop[keep_count < 2] = False |
| present = present & ~drop |
| ci = np.where(present, ci, -1) |
| if soft is not None: |
| soft = np.where(present, soft, 0.0) |
| soft /= np.maximum(soft.sum(1, keepdims=True), 1e-8) |
| if teacher is not None: |
| teacher = np.where(present, teacher, NEG) |
| qi = d.question_idx[idx] |
| if self.q_table is not None: |
| qt = self.q_table |
| L = max(1, int(qt.mask[qi].any(0).nonzero()[0].max(initial=0)) + 1) |
| batch.update(question_bb_tok=qt.tok[qi, :L].astype(np.float32), question_bb_mask=qt.mask[qi, :L], |
| question_bb_pooled=qt.pooled[qi]) |
| else: |
| batch["question_ids"] = self._trim(self.q_tok[qi]) |
| if self.c_table is not None: |
| batch.update(choice_bb_pooled=self.c_table.pooled[ci], choice_bb_mask=ci >= 0) |
| else: |
| batch["choice_ids"] = self._trim(self.c_tok[ci]) |
| batch["target"] = target |
| if soft is not None: |
| batch["soft"] = soft |
| if teacher is not None: |
| batch["teacher_logits"] = teacher |
| if d.score is not None: |
| batch["score_target"] = d.score[idx].astype(np.float32) |
| return batch |
|
|
|
|
| def model_inputs(b: dict) -> dict: |
| """Batch (torch tensors) -> keyword arguments of MiniJEVNet.forward.""" |
| kw = {"state": {k: v for k, v in b.items() if k.startswith("state_")}} |
| if "question_bb_pooled" in b: |
| kw["question_bb"] = (b["question_bb_tok"], b["question_bb_mask"], b["question_bb_pooled"]) |
| else: |
| kw["question_ids"] = b["question_ids"] |
| if "choice_bb_pooled" in b: |
| kw["choice_bb"] = (b["choice_bb_pooled"], b["choice_bb_mask"]) |
| else: |
| kw["choice_ids"] = b["choice_ids"] |
| return kw |
|
|
|
|
| class EpochBatches: |
| """Map-style dataset of *batches* (use with DataLoader(batch_size=None)). |
| |
| The order is a deterministic permutation of (seed, epoch), so a resumed run |
| can skip exactly the batches it already consumed. |
| """ |
|
|
| def __init__(self, encoder: BatchEncoder, batch_size: int, shuffle: bool, seed: int = 0, epoch: int = 0, |
| start_batch: int = 0, drop_last: bool = False, augment: bool = True): |
| self.encoder, self.bs, self.augment = encoder, batch_size, augment |
| n = len(encoder.data) |
| self.seed, self.epoch = seed, epoch |
| order = np.random.default_rng((seed, epoch)).permutation(n) if shuffle else np.arange(n) |
| nb = n // batch_size if drop_last else (n + batch_size - 1) // batch_size |
| self.batches = [order[i * batch_size:(i + 1) * batch_size] for i in range(start_batch, nb)] |
|
|
| def __len__(self): |
| return len(self.batches) |
|
|
| def __getitem__(self, i): |
| rng = np.random.default_rng((self.seed, self.epoch, i, 7)) if self.augment else None |
| return self.encoder(self.batches[i], rng) |
|
|
|
|
| def to_torch(batch: dict[str, np.ndarray], device, non_blocking: bool = False) -> dict: |
| import torch |
|
|
| return {k: torch.as_tensor(v).to(device, non_blocking=non_blocking) for k, v in batch.items()} |
|
|
