File size: 6,778 Bytes
0648c43 | 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 | """Turn a jev record into one packed sequence: the state once, then every question after it.
<state> state tokens | <q> instr <opt> option </opt> ... <decide> | <q> ... <decide> | ...
Isolation: a question token may attend to the state and to earlier tokens of its own question, never to another
question. Position ids restart at the end of the state for every question, so each question sees exactly the
positions it would see if it were asked alone.
Marker tokens are unused reserved ids of the K2 tokenizer, inserted by id (never through text tokenization).
"""
import json
import torch
MARKERS = {"state": "reserved_special_token_100", "q": "reserved_special_token_101", "opt": "reserved_special_token_102",
"end_opt": "reserved_special_token_103", "decide": "reserved_special_token_104"}
def render(v, indent=0):
pad = " " * indent
if isinstance(v, dict):
return "\n".join(f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {x}"
for k, x in v.items())
if isinstance(v, list):
return "\n".join(f"{pad}- {render(x, indent + 1).strip() if isinstance(x, (dict, list)) else x}" for x in v)
return f"{pad}{v}"
def options_and_target(q):
"""Option texts and the target distribution for one question."""
t = q["type"]
if t == "choice":
keys = list(q["criteria"])
opts = [f"{k}: {d}" if d else str(k) for k, d in q["criteria"].items()]
if q.get("soft") is not None:
target = [float(q["soft"][k]) for k in keys]
else:
target = [1.0 if k == q["label"] else 0.0 for k in keys]
elif t == "score":
opts = [str(x) for x in q["criteria"]]
target = [float(x) for x in q["soft"]] if q.get("soft") is not None else \
[1.0 if i == q["label"] else 0.0 for i in range(len(opts))]
else:
crit = q.get("criteria") or {} # optional definitions of the two outcomes: {"false": ..., "true": ...}
opts = [f"{k}: {render(crit[k]).strip()}" if crit.get(k) else k for k in ("false", "true")]
p = float(q["soft"]) if q.get("soft") is not None else float(bool(q["label"]))
target = [1.0 - p, p]
s = sum(target)
return opts, [x / s for x in target]
def teacher_target(q):
"""The question's `teacher` field as a distribution in canonical option order (None if absent)."""
t = q.get("teacher")
if t is None:
return None
if q["type"] == "choice":
return [float(t[k]) for k in q["criteria"]]
if q["type"] == "score":
return [float(x) for x in t]
return [1.0 - float(t), float(t)]
class Encoder:
def __init__(self, tok, max_len=4096, max_state=3072, teacher_mix=None):
"""teacher_mix = alpha: target = alpha * gold + (1 - alpha) * teacher when a question has `teacher`."""
self.tok, self.max_len, self.max_state, self.teacher_mix = tok, max_len, max_state, teacher_mix
vocab = tok.get_vocab()
self.ids = {k: vocab[v] for k, v in MARKERS.items()}
self.bos = tok.bos_token_id
def text(self, s):
return self.tok(s, add_special_tokens=False).input_ids
def encode(self, rec, rng=None):
"""Returns a dict of python lists, or None when not even one question fits.
With `rng`, choice and noul options are shown in a random order (targets permuted to match)."""
state = self.text(render(rec["state"]))[: self.max_state]
ids = [self.bos, self.ids["state"]] + state
seg = [0] * len(ids)
pos = list(range(len(ids)))
S = len(ids)
decide, opt_ends, targets, qtypes, qkeys, perms = [], [], [], [], [], []
for j, (k, q) in enumerate(rec["questions"].items()):
opts, target = options_and_target(q)
tt = teacher_target(q) if self.teacher_mix is not None else None
if tt is not None:
a = self.teacher_mix
target = [a * g + (1 - a) * t for g, t in zip(target, tt)]
order = list(range(len(opts)))
if rng is not None and q["type"] != "score":
rng.shuffle(order)
opts, target = [opts[i] for i in order], [target[i] for i in order]
instr = q["instructions"] if isinstance(q["instructions"], str) else render(q["instructions"]) # Kev uses dicts too
qi = [self.ids["q"]] + self.text(instr)
ends = []
for o in opts:
qi += [self.ids["opt"]] + self.text(o)[:128] + [self.ids["end_opt"]]
ends.append(len(qi) - 1)
qi.append(self.ids["decide"])
if len(ids) + len(qi) > self.max_len:
continue # skip questions that do not fit; the others still train
base = len(ids)
ids += qi
seg += [j + 1] * len(qi)
pos += list(range(S, S + len(qi)))
decide.append(base + len(qi) - 1)
opt_ends.append([base + e for e in ends])
targets.append(target)
qtypes.append(q["type"])
qkeys.append(k)
perms.append(order) # shown position -> canonical option index
if not decide:
return None
return {"ids": ids, "seg": seg, "pos": pos, "decide": decide, "opt_ends": opt_ends,
"targets": targets, "qtypes": qtypes, "qkeys": qkeys, "perms": perms}
def collate(items, pad_id):
"""Right-pad a batch and build the [B, 1, L, L] boolean isolation mask (True = may attend)."""
B, L = len(items), max(len(x["ids"]) for x in items)
ids = torch.full((B, L), pad_id, dtype=torch.long)
pos = torch.zeros((B, L), dtype=torch.long)
seg = torch.full((B, L), -1, dtype=torch.long)
for b, x in enumerate(items):
n = len(x["ids"])
ids[b, :n] = torch.tensor(x["ids"])
pos[b, :n] = torch.tensor(x["pos"])
seg[b, :n] = torch.tensor(x["seg"])
causal = torch.ones(L, L, dtype=torch.bool).tril()
sq, sk = seg[:, :, None], seg[:, None, :]
mask = causal[None] & (sk >= 0) & (sq >= 0) & ((sk == 0) | (sk == sq))
mask |= torch.eye(L, dtype=torch.bool)[None] # padding rows attend to themselves only (avoids NaN softmax)
# flat question index: (batch row, decide position, option end positions, target)
qs = [(b, d, e, t, ty) for b, x in enumerate(items)
for d, e, t, ty in zip(x["decide"], x["opt_ends"], x["targets"], x["qtypes"])]
return {"input_ids": ids, "position_ids": pos, "mask": mask[:, None], "questions": qs}
def load_jsonl(path, limit=0):
out = []
with open(path) as fh:
for i, line in enumerate(fh):
if limit and i >= limit:
break
out.append(json.loads(line))
return out
|