| """v6 training contexts: the same item rendered as a separate sequence after no rationale (`direct`), a short stock-model rationale (`brief`) or a longer stock-model |
| thought (`deliberate`). The layouts are exactly those the rationales were generated and validated with (tools/render_for_gen.py, tools/score_rationales.py, |
| dense/think_gate2.py), so training and inference see the same tokens: |
| |
| direct codec prompt, slot read right away (answer_slot.render_ar) |
| brief <|im_start|>system\\n{rat_s32}<|im_end|>\\n + codec prompt + "Rationale:\\n" + rationale (first `brief_cap` tokens) + "\\n\\nAnswer: " -> slot |
| deliberate thinking-template prompt + draft_d prefix + thought + "\\n</think>\\n\\n" -> slot |
| per-option ('parallel') items: the generation prompt is cut after the assistant turn start, the rationale goes there, then "\\n\\n" (brief) or the close (deliberate) |
| and the unchanged slot block ("1:" ... "k:") follows. |
| |
| The rationale text was written for the item as displayed (identity option order, no augmentation), so the rationale contexts never shuffle or augment. |
| """ |
| import random |
| from typing import Any, Dict, List, Optional |
|
|
| import answer_slot as AS |
| from dense.rationale_prompts import SYSTEM |
| from dense.schema import codec_item |
|
|
| CONTEXTS = ("brief", "direct", "deliberate") |
| DRAFT_D = "Okay. Brief reasoning, a few words per step, no double-checking. First," |
|
|
|
|
| def parse_kv(s: str) -> Dict[str, float]: |
| return {k: float(v) for k, v in (p.split("=") for p in s.split(","))} |
|
|
|
|
| class Contexts: |
| def __init__(self, tok: Any, c_off: AS.SlotCodec, prompt: str = "rat_s32", weights: str = "brief=.65,direct=.15,deliberate=.2", |
| loss_weights: str = "brief=1,direct=1,deliberate=.5", parallel_direct: float = 0.5, brief_cap: int = 64): |
| self.tok, self.c_off, self.brief_cap = tok, c_off, brief_cap |
| self.c_on = AS.SlotCodec(tok, state_max=None, option_max=None, template_kwargs={"enable_thinking": True}, mask_id=tok.pad_token_id or 0, |
| strip_label_keys=True, autoregressive=True) |
| enc = lambda t: tok(t, add_special_tokens=False)["input_ids"] |
| self.sys_ids = enc(f"<|im_start|>system\n{SYSTEM[prompt]}<|im_end|>\n") |
| self.pref, self.tail, self.sep = enc("Rationale:\n"), enc("\n\nAnswer: "), enc("\n\n") |
| self.draft_d, self.close = enc(DRAFT_D), enc("\n</think>\n\n") |
| self.weights, self.loss_w, self.parallel_direct = parse_kv(weights), parse_kv(loss_weights), parallel_direct |
|
|
| |
| def mode(self, ex: Dict[str, Any]) -> str: |
| return AS.slot_mode(codec_item(ex)["qtype"], len(ex["target"])) |
|
|
| def draw(self, ex: Dict[str, Any], seed: Any, epoch: int) -> str: |
| rat = ex.get("rationale") or {} |
| w = dict(self.weights) |
| if self.mode(ex) == "parallel" and self.parallel_direct >= 1.0: |
| return "direct" |
| if self.mode(ex) == "parallel": |
| moved = w["brief"] * self.parallel_direct |
| w["brief"] -= moved; w["direct"] += moved |
| if "deliberate" not in rat: |
| w["brief"] += w.pop("deliberate", 0.0); w["deliberate"] = 0.0 |
| if "brief" not in rat: |
| w["direct"] += w["brief"]; w["brief"] = 0.0 |
| names = [c for c in CONTEXTS if w.get(c, 0) > 0] |
| return random.Random(f"{seed}:ctx:{ex['id']}:{epoch}").choices(names, weights=[w[c] for c in names])[0] |
|
|
| def extra_tokens(self, ex: Dict[str, Any], ctx: str) -> int: |
| """Length added by the context (for the solo-micro-batch / skip estimates).""" |
| if ctx == "brief": |
| return len(self.sys_ids) + len(self.pref) + len(self.tail) + min(ex["rationale"]["brief"]["n"], self.brief_cap) |
| if ctx == "deliberate": |
| return len(self.draft_d) + len(self.close) + ex["rationale"]["deliberate"]["n"] |
| return 0 |
|
|
| |
| def gen_prompt(self, ex: Dict[str, Any]) -> List[int]: |
| """Token ids of the stock-model rationale request (brief): system line + codec prompt (cut after the assistant start for per-option items) + "Rationale:\\n".""" |
| r = self.c_off.render_ar(codec_item(ex)) |
| ids = r["ids"][:self._cut(self.c_off, r)] if r["mode"] == "parallel" else r["ids"] |
| return self.sys_ids + ids + self.pref |
|
|
| def clean_brief(self, g: List[int]) -> List[int]: |
| """Generated tokens (stop token removed) -> stored form: a trailing bare-answer line (<= 12 chars after a newline) is dropped so the slot is not told the answer |
| (as tools/gen_rationales.py).""" |
| txt = self.tok.decode(g).rstrip() |
| last = txt.split("\n")[-1].strip() |
| if g and len(last) <= 12 and "\n" in txt: |
| g = self.tok(txt[: -len(last)].rstrip() + "\n", add_special_tokens=False)["input_ids"] |
| return g |
|
|
| |
| @staticmethod |
| def _cut(codec: AS.SlotCodec, r: Dict[str, Any]) -> int: |
| ids, post = r["ids"], codec.post |
| return next(i + len(post) for i in range(r["state_end"], len(ids) - len(post) + 1) if ids[i:i + len(post)] == post) |
|
|
| def render(self, ex: Dict[str, Any], ctx: str) -> Dict[str, Any]: |
| """`brief` / `deliberate` only; `direct` is rendered by the caller (it keeps the augmentations and the option shuffle).""" |
| assert ctx in ("brief", "deliberate"), ctx |
| it = codec_item(ex) |
| rat = ex["rationale"][ctx] |
| if ctx == "brief": |
| codec, head, body = self.c_off, self.sys_ids, rat["tokens"][:self.brief_cap] |
| pre_tail, tail_single, tail_par = self.pref, self.tail, self.sep |
| else: |
| codec, head, body = self.c_on, [], rat["tokens"] |
| pre_tail, tail_single, tail_par = self.draft_d, self.close, self.close |
| r = codec.render_ar(it) |
| if r["mode"] == "parallel": |
| cut = self._cut(codec, r) |
| new = head + r["ids"][:cut] + pre_tail + body + tail_par + r["ids"][cut:] |
| off = len(head) + len(pre_tail) + len(body) + len(tail_par) |
| return {**r, "ids": new, "slot_pos": [p + off for p in r["slot_pos"]]} |
| ids = head + r["ids"] + pre_tail + body + tail_single |
| return {**r, "ids": ids, "slot_pos": [len(ids) - 1]} |
|
|
