File size: 2,322 Bytes
065ec46 | 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 | """Batches of rendered decisions, and the per-example augmentation for training."""
from __future__ import annotations
import random
import torch
from ..render import Renderer, Rendered
from ..schema import Decision
def collate(items: list[Rendered], pad_id: int, labels: list[int] | None = None) -> dict:
"""Right-padded ids, and candidate positions padded with -1.
Right padding is safe for a causal backbone: padding sits after every
position that is read, so nothing read has attended to it.
"""
L = max(len(r.input_ids) for r in items)
K = max(len(r.option_positions) for r in items)
ids = torch.full((len(items), L), pad_id, dtype=torch.long)
mask = torch.zeros((len(items), L), dtype=torch.long)
pos = torch.full((len(items), K), -1, dtype=torch.long)
for i, r in enumerate(items):
ids[i, :len(r.input_ids)] = torch.tensor(r.input_ids)
mask[i, :len(r.input_ids)] = 1
pos[i, :len(r.option_positions)] = torch.tensor(r.option_positions)
out = {"input_ids": ids, "attention_mask": mask, "option_pos": pos}
if labels is not None:
out["labels"] = torch.tensor(labels, dtype=torch.long)
return out
def training_view(d: Decision, rng: random.Random, max_options: int = 24
) -> tuple[list[str], int]:
"""(candidates, target index) as the model sees this record in one epoch.
Choice candidates are shuffled, so no position carries the answer, and a
long label set (CLINC has 151 intents) is subsampled to between 4 and
`max_options`, gold always included, so the model learns from lists of
every length. Scales and yes/no keep their order: that order is meaning.
"""
cands, t = d.candidates(), d.target()
if d.type != "choice":
return cands, t
gold = cands[t]
others = [c for i, c in enumerate(cands) if i != t]
if len(cands) > max_options:
k = rng.randint(min(4, max_options), max_options) - 1
others = rng.sample(others, k)
picked = others + [gold]
rng.shuffle(picked)
return picked, picked.index(gold)
def render_for_training(r: Renderer, d: Decision, rng: random.Random,
max_options: int = 24) -> tuple[Rendered, int]:
cands, t = training_view(d, rng, max_options)
return r.render(d, cands), t
|