File size: 7,191 Bytes
026f63e | 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """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()}
|