"""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) # + one empty row for "absent" 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" # byte keeps the original fixed-width layout 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) # extra all-pad row so choice index -1 maps to "no choice" 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: # randomly hide non-target choices so the model sees 2..N options (variable-N training) 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]) # -1 -> last (all-pad) row 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()}