| """Prompt construction. One context, N typed questions, N answer slots. |
| |
| All N decisions are read from a single forward pass: the logits at each |
| "Answer k: (" slot are restricted to the option-letter tokens. No answer |
| letters are ever inserted, so slot k sees the context and all questions but |
| no earlier answers (the decisions are conditionally independent given input). |
| """ |
| import random |
|
|
| LETTERS = "ABCDEFGHIJ" |
| NARROW = len(LETTERS) |
| MAX_OPTIONS = 255 |
| |
| ABSTAIN_PREFIXES = ("none of the above", "none of these", "not listed", "no suitable", "does not apply", "cannot tell") |
| ABSTAIN_EXACT = ("other", "unsure", "something else", "neither of these", "other / not covered") |
|
|
|
|
| def is_abstain_option(o): |
| o = o.strip().lower() |
| return o.startswith(ABSTAIN_PREFIXES) or o in ABSTAIN_EXACT |
|
|
|
|
| _LABELS = {} |
| _OPT_CACHE = {} |
|
|
|
|
| def _enc_opt(tok, text): |
| """Token ids of ") <option text>" (cached: fixed label sets repeat the same strings millions of times).""" |
| key = (id(tok), text) |
| v = _OPT_CACHE.get(key) |
| if v is None: |
| v = tok.encode(f") {text}", add_special_tokens=False) |
| if len(_OPT_CACHE) < 2_000_000: |
| _OPT_CACHE[key] = v |
| return v |
|
|
|
|
| def label_table(tok): |
| """(label strings, label token ids), MAX_OPTIONS entries; the first NARROW are A..J so narrow questions are unchanged.""" |
| key = id(tok) |
| if key not in _LABELS: |
| import string |
| U = string.ascii_uppercase |
| names = list(U) + [a + b for a in U for b in U] |
| out = [] |
| for n in names: |
| t = tok.encode(n, add_special_tokens=False) |
| if len(t) == 1: |
| out.append((n, t[0])) |
| if len(out) == MAX_OPTIONS: |
| break |
| assert len(out) == MAX_OPTIONS and len({i for _, i in out}) == MAX_OPTIONS |
| _LABELS[key] = ([n for n, _ in out], [i for _, i in out], |
| tok.encode("\n(", add_special_tokens=False)) |
| return _LABELS[key] |
|
|
|
|
| def _select(q, rng, max_options): |
| opts = list(range(len(q.options))) |
| if len(opts) > max_options: |
| |
| forced = {q.gold} | {i for i, o in enumerate(q.options) if is_abstain_option(o)} |
| others = [i for i in opts if i not in forced] |
| opts = rng.sample(others, max_options - len(forced)) + list(forced) |
| rng.shuffle(opts) |
| return opts |
|
|
|
|
| def _options_ids(tok, q, opts): |
| if len(opts) <= NARROW: |
| return tok.encode("".join(f"\n({LETTERS[j]}) {q.options[oi]}" for j, oi in enumerate(opts)), add_special_tokens=False) |
| _, lab_ids, open_ids = label_table(tok); out = [] |
| for j, oi in enumerate(opts): |
| out += open_ids + [lab_ids[j]] + _enc_opt(tok, q.options[oi]) |
| return out |
|
|
|
|
| def build_schema_first(example, tok, rng=None, max_options=NARROW, max_ctx_tokens=1536): |
| """Schema-first layout: all question/option blocks, then the context, then one answer slot per question. |
| |
| Question 1: ...\nOptions:\n(A) ... <- prefix: depends only on the questions, so its cache (attention KV and |
| \n\nQuestion 2: ... delta-net states) is computed once per schema and reused for every state |
| \n\nContext:\n<state>\n\nAnswer 1: (\nAnswer 2: ( |
| |
| The three parts are tokenized separately, so `ids[:prefix_len]` is identical for every state.""" |
| rng = rng or random |
| perms = [_select(q, rng, max_options) for q in example.qs] |
| pre = schema_prefix_ids(tok, example.qs, perms) |
| suf, slots = schema_suffix_ids(tok, example.context, len(example.qs), max_ctx_tokens) |
| return dict(ids=pre + suf, slots=[len(pre) + s for s in slots], golds=[p.index(q.gold) if q.gold in p else -1 for p, q in zip(perms, example.qs)], |
| nopts=[len(p) for p in perms], perms=perms, prefix_len=len(pre)) |
|
|
|
|
| def schema_prefix_ids(tok, qs, perms=None): |
| """Token ids of the question/option blocks (the cacheable part of the schema-first layout).""" |
| multi = len(qs) > 1; pre = [] |
| for k, q in enumerate(qs): |
| opts = perms[k] if perms is not None else list(range(len(q.options))) |
| pre += tok.encode(f"{chr(10) * 2 if k else ''}Question{' ' + str(k + 1) if multi else ''}: {q.text}\nOptions:", add_special_tokens=False) + _options_ids(tok, q, opts) |
| return pre |
|
|
|
|
| def schema_suffix_ids(tok, context, n_q, max_ctx_tokens=1536): |
| """Token ids after the schema prefix: the context and one answer slot per question. Returns (ids, slot positions in ids).""" |
| ids = tok.encode("\n\nContext:\n", add_special_tokens=False) + tok.encode(context, add_special_tokens=False)[:max_ctx_tokens]; slots = [] |
| for k in range(n_q): |
| ids += tok.encode(f"{chr(10) * 2 if k == 0 else chr(10)}Answer{' ' + str(k + 1) if n_q > 1 else ''}: (", add_special_tokens=False); slots.append(len(ids) - 1) |
| return ids, slots |
|
|
|
|
| def build(example, tok, rng=None, max_options=NARROW, max_ctx_tokens=1536, layout="state_first"): |
| """Returns dict(ids=list[int], slots=list[int], golds=list[int], nopts=list[int], perms=list[list[int]]).""" |
| if layout == "schema_first": |
| return build_schema_first(example, tok, rng, max_options, max_ctx_tokens) |
| rng = rng or random |
| ctx_ids = tok.encode("Context:\n" + example.context, add_special_tokens=False)[:max_ctx_tokens] |
| ids = list(ctx_ids) |
| slots, golds, nopts, perms = [], [], [], [] |
| multi = len(example.qs) > 1 |
| for k, q in enumerate(example.qs): |
| opts = list(range(len(q.options))) |
| if len(opts) > max_options: |
| |
| forced = {q.gold} | {i for i, o in enumerate(q.options) if is_abstain_option(o)} |
| others = [i for i in opts if i not in forced] |
| keep = rng.sample(others, max_options - len(forced)) + list(forced) |
| opts = keep |
| rng.shuffle(opts) |
| head = f"\n\nQuestion{' ' + str(k + 1) if multi else ''}: {q.text}\nOptions:" |
| tail = f"\nAnswer{' ' + str(k + 1) if multi else ''}: (" |
| if len(opts) <= NARROW: |
| lines = [head] + [f"\n({LETTERS[j]}) {q.options[oi]}" for j, oi in enumerate(opts)] + [tail] |
| piece = tok.encode("".join(lines), add_special_tokens=False) |
| else: |
| _, lab_ids, open_ids = label_table(tok) |
| piece = tok.encode(head, add_special_tokens=False) |
| for j, oi in enumerate(opts): |
| piece += open_ids + [lab_ids[j]] + _enc_opt(tok, q.options[oi]) |
| piece += tok.encode(tail, add_special_tokens=False) |
| ids.extend(piece) |
| slots.append(len(ids) - 1) |
| golds.append(opts.index(q.gold) if q.gold in opts else -1) |
| nopts.append(len(opts)) |
| perms.append(opts) |
| return dict(ids=ids, slots=slots, golds=golds, nopts=nopts, perms=perms) |
|
|
|
|
| def letter_ids(tok): |
| ids = label_table(tok)[1] |
| for j, L in enumerate(LETTERS): |
| assert tok.encode(L, add_special_tokens=False) == [ids[j]], L |
| return ids |
|
|
|
|
| def render(example, tok, **kw): |
| b = build(example, tok, **kw) |
| return tok.decode(b["ids"]) |
|
|
