remove __pycache__ and a decoder backup file that were uploaded by mistake

maram_inference/maram/decode/decode.py.pre_events_20261031T0701

@@ -1,472 +0,0 @@
-"""MARAM decoding: greedy1 | bestofK | native, sharing the prompt KV cache.
-
-The physical sequence grows "as generated": each step appends one token per stream slot (or <pad> for a
-finished/dead slot), block_idx = current block, and a <chk> query after every N content tokens. Attention
-for the new queries is computed over the whole physical cache with the SAME mask_mod as training
-(`maram.model.mask.build_mask`) evaluated on the concatenated (cache + new) metadata, materialised densely
-(queries are few, so this is exact and cheap).
-
-native: every N tokens a <chk> forward; score_k = sigmoid(correct) - 0.1*sigmoid(remaining);
-kill the lowest-scored unfinished stream; fork the highest-scored unfinished stream's KV entries into the
-killed slot (copies of its K/V with the new slot's seg id) with header <s_k><strat:z>, z = a strategy the
-killed stream did not have (preferring one the parent does not have either). --no_respawn = prune only.
-"""
-from __future__ import annotations
-import math, re
-from dataclasses import dataclass, field
-from typing import List, Optional
-import numpy as np, torch, torch.nn.functional as F
-
-from maram.model.mask import build_mask, PAD_SEG
-from maram.model.model import MaramModel, ModelConfig
-from maram.tokenizer.maram_tokenizer import (PAD, PROMPT, CHK, END, STREAM_IDS, STRAT_IDS, STRAT_NAMES, FIRST_MERGE_ID)
-
-ANSWER_RE = re.compile(r"Answer\s*:\s*(.+)", re.I)
-
-
-class LayerKV:
-    def __init__(self, B, H, S, hd, device, dtype):
-        self.k = torch.zeros(B, H, S, hd, device=device, dtype=dtype)
-        self.v = torch.zeros(B, H, S, hd, device=device, dtype=dtype)
-        self.len = 0
-
-    def update(self, k_new, v_new):
-        n = k_new.shape[2]
-        self.k[:, :, self.len:self.len + n] = k_new
-        self.v[:, :, self.len:self.len + n] = v_new
-        return self.k[:, :, :self.len + n], self.v[:, :, :self.len + n]
-
-
-class KVCache:
-    def __init__(self, model: MaramModel, B: int, S: int, device, dtype):
-        cfg = model.cfg
-        self.layers = [LayerKV(B, cfg.n_heads, S, cfg.head_dim, device, dtype) for _ in range(cfg.n_layers)]
-        self.seg = torch.full((B, S), PAD_SEG, dtype=torch.long, device=device)
-        self.pos = torch.zeros((B, S), dtype=torch.long, device=device)
-        self.block = torch.zeros((B, S), dtype=torch.long, device=device)
-        self.B, self.S_alloc, self.len, self.device = B, S, 0, device
-
-    def commit(self, n, seg, pos, block):
-        self.seg[:, self.len:self.len + n] = seg
-        self.pos[:, self.len:self.len + n] = pos
-        self.block[:, self.len:self.len + n] = block
-        self.len += n
-        for l in self.layers:
-            l.len = self.len
-
-    def expand_rows(self, reps: int):
-        """Repeat every batch row `reps` times (share a prefilled prompt across samples)."""
-        for l in self.layers:
-            l.k = l.k.repeat_interleave(reps, dim=0); l.v = l.v.repeat_interleave(reps, dim=0)
-        self.seg = self.seg.repeat_interleave(reps, dim=0); self.pos = self.pos.repeat_interleave(reps, dim=0)
-        self.block = self.block.repeat_interleave(reps, dim=0); self.B *= reps
-
-    def invalidate(self, row, seg_id):
-        """Mask every cache entry of (row, seg): a killed stream disappears for all later queries."""
-        m = self.seg[row, :self.len] == seg_id
-        self.seg[row, :self.len][m] = PAD_SEG
-        return int(m.sum())
-
-    def fork(self, rows, src_seg, dst_seg):
-        """For each (row, src, dst): append copies of all cache entries with seg==src relabelled seg=dst."""
-        if not rows:
-            return 0
-        counts = []
-        idxs = []
-        for r, s, d in zip(rows, src_seg, dst_seg):
-            ix = torch.nonzero(self.seg[r, :self.len] == s).flatten()
-            idxs.append(ix); counts.append(len(ix))
-        m = max(counts)
-        if self.len + m > self.S_alloc:
-            raise RuntimeError("KV cache too small for fork")
-        for r, s, d, ix in zip(rows, src_seg, dst_seg, idxs):
-            n = len(ix)
-            if n == 0:
-                continue
-            for l in self.layers:
-                l.k[r, :, self.len:self.len + n] = l.k[r, :, ix]
-                l.v[r, :, self.len:self.len + n] = l.v[r, :, ix]
-            self.seg[r, self.len:self.len + n] = d
-            self.pos[r, self.len:self.len + n] = self.pos[r, ix]
-            self.block[r, self.len:self.len + n] = self.block[r, ix]
-        self.len += m
-        for l in self.layers:
-            l.len = self.len
-        return m
-
-
-@dataclass
-class Stream:
-    slot: int                      # 0-based -> seg = slot+1, header STREAM_IDS[slot]
-    strat: int
-    tokens: List[int] = field(default_factory=list)   # content tokens incl <end>
-    n_fed: int = 0                 # tokens fed so far (headers + content) -> next pos = P + n_fed
-    finished: bool = False
-    dead: bool = False
-    score: float = float("nan")
-    p_correct: float = float("nan")
-    p_remain: float = float("nan")
-    last_h: Optional[torch.Tensor] = None
-    pending: Optional[int] = None  # token to feed next (sampled last step)
-    parent: Optional[int] = None
-    born_block: int = 0
-    generated: int = 0             # sampled tokens (counted for cost), incl <end>
-    gen: int = 0                   # occupant index of this slot (0 = original, +1 per respawn)
-    temp: Optional[float] = None   # sampling temperature of this stream (None = run-level `temperature`)
-
-
-class Decoder:
-    def __init__(self, model: MaramModel, tokenizer, device="cuda", dtype=torch.bfloat16, judge_mode="full"):
-        self.model = model.eval().to(device)
-        self.tk = tokenizer
-        self.device = device; self.dtype = dtype
-        self.cfg = model.cfg
-        self.judge_mode = judge_mode
-
-    # ------------------------------------------------------------------ core forward over the cache
-    @torch.no_grad()
-    def _forward(self, cache: KVCache, tokens, seg, pos, block):
-        """tokens/seg/pos/block: [B, n] new queries. Appends them to the cache. Returns hidden [B,n,d]."""
-        B, n = tokens.shape
-        S = cache.len
-        seg_all = torch.cat([cache.seg[:, :S], seg], 1); pos_all = torch.cat([cache.pos[:, :S], pos], 1)
-        blk_all = torch.cat([cache.block[:, :S], block], 1)
-        mm = build_mask(seg_all, pos_all, blk_all, self.cfg.K, self.cfg.N)
-        b = torch.arange(B, device=self.device)[:, None, None]
-        q = (S + torch.arange(n, device=self.device))[None, :, None]
-        kv = torch.arange(S + n, device=self.device)[None, None, :]
-        dense = mm(b, 0, q, kv)                                   # [B, n, S+n]
-        with torch.autocast(self.device if self.device == "cuda" else "cpu", dtype=self.dtype, enabled=self.dtype != torch.float32):
-            h = self.model.hidden_states(tokens, pos, None, dense, cache.layers)
-        cache.commit(n, seg, pos, block)
-        return h
-
-    def _logits(self, h):
-        return self.model.lm_head(h).float()
-
-    def _judge(self, h_chk, h_last):
-        with torch.autocast(self.device if self.device == "cuda" else "cpu", dtype=self.dtype, enabled=self.dtype != torch.float32):
-            return self.model.judge(h_chk, h_last, self.judge_mode).float()
-
-    def _sample(self, logits, temperature):
-        if temperature <= 0:
-            return logits.argmax(-1)
-        return torch.multinomial(torch.softmax(logits / temperature, -1), 1).squeeze(-1)
-
-    # ------------------------------------------------------------------ prompt prefill (shared)
-    def prefill(self, prompts: List[List[int]], K_slots: int, max_new: int, extra: int = 0):
-        """prompts: list of token id lists (<=256). Builds a cache sized for the run and prefills the prompt
-        region (left-aligned, pad masked). Returns cache and P per row."""
-        B = len(prompts)
-        Pm = max(len(p) for p in prompts) + 1
-        S = Pm + K_slots * (max_new + 2 + 2 * 8) + 2 * 8 + extra
-        cache = KVCache(self.model, B, S, self.device, self.dtype)
-        toks = torch.full((B, Pm), PAD, dtype=torch.long, device=self.device)
-        seg = torch.full((B, Pm), PAD_SEG, dtype=torch.long, device=self.device)
-        pos = torch.zeros((B, Pm), dtype=torch.long, device=self.device)
-        P = torch.zeros(B, dtype=torch.long, device=self.device)
-        for i, p in enumerate(prompts):
-            toks[i, 0] = PROMPT; toks[i, 1:1 + len(p)] = torch.tensor(p, device=self.device)
-            seg[i, :1 + len(p)] = 0; pos[i, :1 + len(p)] = torch.arange(1 + len(p), device=self.device)
-            P[i] = 1 + len(p)
-        self._forward(cache, toks, seg, pos, torch.zeros_like(pos))
-        return cache, P
-
-    # ------------------------------------------------------------------ generic multi-stream loop
-    @torch.no_grad()
-    def run_streams(self, prompts, strats: List[List[int]], temperature: float, max_new: int = 512,
-                    use_judge: bool = False, respawn: bool = True, N: Optional[int] = None, sample_seed: int = 1337,
-                    slot_temps: Optional[List[float]] = None, strat_pool: Optional[List[int]] = None,
-                    fork_temp: Optional[float] = None, row_temps: Optional[List[List[float]]] = None,
-                    prune: str = "always", tau: float = 0.0, delta: float = 0.0, immune_slots: Optional[List[int]] = None,
-                    dump_tokens: bool = False):
-        """strats[i] = strategy token per slot for problem i (len K_slots). Returns per-problem dict.
-        prune (use_judge only): "always" = at every <chk> with >=2 unfinished streams kill the lowest-scored one (the
-        original rule); "gated" = kill it only if its p_correct < tau and the best unfinished stream's p_correct exceeds
-        it by >= delta, otherwise do nothing at that checkpoint; "none" = never kill (every stream runs to completion and
-        the judge only picks at the final <chk>). immune_slots: slots that are never killed (pruning happens among the
-        other slots only; the fork source may still be an immune slot). dump_tokens: keep every stream's token ids
-        (live and killed) in the records.
-        slot_temps: optional per-slot sampling temperature (overrides `temperature`; a forked child inherits its
-        slot's temperature unless fork_temp is given). row_temps[i][k]: per-problem, per-slot temperature (overrides
-        slot_temps). fork_temp: temperature of every respawned child. strat_pool: strategies a forked child may take
-        (default all three)."""
-        K = len(strats[0]); N = N or self.cfg.N
-        pool = list(strat_pool) if strat_pool else list(STRAT_IDS)
-        assert prune in ("always", "gated", "none"), prune
-        immune = set(immune_slots or [])
-        gate_skips = [0] * len(prompts)   # checkpoints where the gate blocked a kill
-        if slot_temps is not None:
-            assert len(slot_temps) == K, (slot_temps, K)
-        per_stream_temps = slot_temps is not None or row_temps is not None or fork_temp is not None
-        def init_temp(i, k):
-            if row_temps is not None: return row_temps[i][k]
-            if slot_temps is not None: return slot_temps[k]
-            return temperature
-        assert K <= len(STREAM_IDS)
-        B = len(prompts)
-        g = torch.Generator(device=self.device); g.manual_seed(sample_seed)
-        nb = math.ceil(max_new / N)
-        cache, P = self.prefill(prompts, K, max_new, extra=((nb - 1) * (max_new + nb + 2) if use_judge and respawn else 0))
-        Pl = P.tolist()   # per-step token/seg/pos tensors are built on the CPU and moved once (avoids per-element CUDA launches)
-        streams = [[Stream(slot=k, strat=strats[i][k], temp=init_temp(i, k)) for k in range(K)] for i in range(B)]
-        chk_count = [0] * B; forks = [0] * B; kills = [0] * B; dead_generated = [0] * B
-        judge_log = [[] for _ in range(B)]; dead_recs = [[] for _ in range(B)]
-        dev = self.device
-        n_blocks = nb
-        for b in range(n_blocks):
-            # ---- headers for this block
-            for hdr_step in range(2 if b == 0 else 1):
-                toks = torch.full((B, K), PAD, dtype=torch.long)
-                seg = torch.full((B, K), PAD_SEG, dtype=torch.long)
-                pos = torch.zeros((B, K), dtype=torch.long)
-                for i in range(B):
-                    for k, s in enumerate(streams[i]):
-                        if s.finished or s.dead:
-                            continue
-                        if b == 0 or s.born_block == b:
-                            tid = STREAM_IDS[k] if hdr_step == 0 else s.strat
-                        else:
-                            if hdr_step == 1: continue
-                            tid = STREAM_IDS[k]
-                        toks[i, k] = tid; seg[i, k] = k + 1; pos[i, k] = Pl[i] + s.n_fed; s.n_fed += 1
-                if (seg != PAD_SEG).any():
-                    toks, seg, pos = toks.to(dev), seg.to(dev), pos.to(dev)
-                    h = self._forward(cache, toks, seg, pos, torch.full_like(pos, b))
-                    for i in range(B):
-                        for k, s in enumerate(streams[i]):
-                            if seg[i, k] != PAD_SEG:
-                                s.last_h = h[i, k]
-            # for streams born this block (forks) the two headers were fed above when hdr_step covered them
-            # ---- N content tokens
-            for t in range(N):
-                if b * N + t >= max_new:
-                    break
-                # sample next token for every live stream from the logits at its last fed position
-                toks = torch.full((B, K), PAD, dtype=torch.long)
-                seg = torch.full((B, K), PAD_SEG, dtype=torch.long)
-                pos = torch.zeros((B, K), dtype=torch.long)
-                live = [(i, k) for i in range(B) for k, s in enumerate(streams[i]) if not (s.finished or s.dead)]
-                if not live:
-                    break
-                hs = torch.stack([streams[i][k].last_h for i, k in live])
-                logits = self._logits(hs)
-                # never sample control tokens: <pad> <prompt> <s1..s4> <strat:*> <chk> (256..265) and <right>/<wrong>
-                # (267, 268). Byte tokens 0..255 (every single character incl. digits) and <end> (266) stay allowed.
-                logits[:, PAD:END] = -float("inf")
-                logits[:, END + 1:FIRST_MERGE_ID] = -float("inf")
-                if per_stream_temps:
-                    tt = torch.tensor([streams[i][k].temp for i, k in live], device=logits.device, dtype=logits.dtype)[:, None]
-                    samp = torch.multinomial(torch.softmax(logits / tt.clamp_min(1e-6), -1), 1, generator=g).squeeze(-1)
-                    nxt = torch.where(tt[:, 0] <= 0, logits.argmax(-1), samp)
-                elif temperature <= 0:
-                    nxt = logits.argmax(-1)
-                else:
-                    nxt = torch.multinomial(torch.softmax(logits / temperature, -1), 1, generator=g).squeeze(-1)
-                for (i, k), tid in zip(live, nxt.tolist()):
-                    s = streams[i][k]
-                    if len(s.tokens) + 1 >= max_new and tid != END:
-                        tid = END                                 # forced finish at the budget
-                    s.tokens.append(tid); s.generated += 1
-                    toks[i, k] = tid; seg[i, k] = k + 1; pos[i, k] = Pl[i] + s.n_fed; s.n_fed += 1
-                toks, seg, pos = toks.to(dev), seg.to(dev), pos.to(dev)
-                h = self._forward(cache, toks, seg, pos, torch.full_like(pos, b))
-                for (i, k) in live:
-                    s = streams[i][k]
-                    s.last_h = h[i, k]
-                    if s.tokens[-1] == END:
-                        s.finished = True
-            # ---- checkpoint
-            if use_judge:
-                toks = torch.full((B, 1), CHK, dtype=torch.long, device=dev)
-                seg = torch.full((B, 1), K + 1, dtype=torch.long, device=dev)
-                pos = (P + (b + 1) * N)[:, None]
-                h = self._forward(cache, toks, seg, pos, torch.full_like(pos, b))
-                fork_rows, fork_src, fork_dst = [], [], []   # one batched fork per checkpoint (shared cache length)
-                for i in range(B):
-                    chk_count[i] += 1
-                    cand = [s for s in streams[i] if not s.dead and s.last_h is not None]
-                    if not cand:
-                        continue
-                    hl = torch.stack([s.last_h for s in cand])
-                    hc = h[i, 0][None].expand(len(cand), -1)
-                    j = self._judge(hc, hl)
-                    pc = torch.sigmoid(j[:, 0]).tolist(); pr = torch.sigmoid(j[:, 1]).tolist()
-                    for s, a, r in zip(cand, pc, pr):
-                        s.p_correct, s.p_remain = a, r
-                        s.score = a - 0.1 * r
-                        judge_log[i].append({"block": b, "slot": s.slot, "gen": s.gen, "strat": STRAT_NAMES[s.strat], "p_correct": a,
-                                             "p_remain": r, "finished": s.finished, "len": len(s.tokens)})
-                    last_block = (b + 1) * N >= max_new
-                    unfinished = [s for s in cand if not s.finished]
-                    killable = [s for s in unfinished if s.slot not in immune]
-                    if len(unfinished) >= 2 and killable and not last_block and prune != "none":
-                        worst = min(killable, key=lambda s: s.score)
-                        best = max(unfinished, key=lambda s: s.score)
-                        allowed = worst is not best
-                        if allowed and prune == "gated":
-                            allowed = worst.p_correct < tau and (best.p_correct - worst.p_correct) >= delta
-                            if not allowed:
-                                gate_skips[i] += 1
-                        if allowed:
-                            worst.dead = True; kills[i] += 1; dead_generated[i] += worst.generated
-                            cache.invalidate(i, worst.slot + 1)      # killed stream vanishes from the cache
-                            dead_recs[i].append({"slot": worst.slot, "gen": worst.gen, "strat": STRAT_NAMES[worst.strat],
-                                                 "text": self.tk.decode(worst.tokens), "len": len(worst.tokens),
-                                                 "generated_own": worst.generated, "killed_block": b, "parent": worst.parent,
-                                                 "p_correct": worst.p_correct, "p_remain": worst.p_remain,
-                                                 **({"tokens": list(worst.tokens)} if dump_tokens else {})})
-                            if respawn:
-                                fork_rows.append(i); fork_src.append(best.slot + 1); fork_dst.append(worst.slot + 1)
-                                forks[i] += 1
-                                taken = {s.strat for s in streams[i] if not s.dead}
-                                bi_ = pool.index(best.strat) if best.strat in pool else -1
-                                order = pool[bi_ + 1:] + pool[:bi_ + 1]
-                                choices = ([z for z in order if z != worst.strat and z != best.strat]
-                                           or [z for z in order if z != worst.strat] or order)
-                                child = Stream(slot=worst.slot, strat=choices[0], tokens=list(best.tokens), n_fed=best.n_fed,
-                                               last_h=best.last_h, parent=best.slot, born_block=b + 1, gen=worst.gen + 1,
-                                               temp=fork_temp if fork_temp is not None else init_temp(i, worst.slot))
-                                streams[i][worst.slot] = child
-                if fork_rows:
-                    cache.fork(fork_rows, fork_src, fork_dst)
-            if all(s.finished or s.dead for i in range(B) for s in streams[i]):
-                break
-        # ---- collect
-        out = []
-        for i in range(B):
-            recs = []
-            for s in streams[i]:
-                recs.append({"slot": s.slot, "gen": s.gen, "strat": STRAT_NAMES[s.strat], "text": self.tk.decode(s.tokens), "finished": s.finished,
-                             "dead": s.dead, "score": s.score, "p_correct": s.p_correct, "p_remain": s.p_remain, "len": len(s.tokens),
-                             "parent": s.parent, "generated_own": s.generated, "temp": s.temp,
-                             **({"tokens": list(s.tokens)} if dump_tokens else {})})
-            out.append({"streams": recs, "dead_streams": dead_recs[i], "chk_forwards": chk_count[i], "kills": kills[i], "forks": forks[i],
-                        "judge_log": judge_log[i], "cache_len": cache.len, "_dead_generated": dead_generated[i],
-                        "gate_skips": gate_skips[i]})
-        return out, streams
-
-    # ------------------------------------------------------------------ public modes
-    def _account(self, result, live_streams):
-        gen = sum(s.generated for s in live_streams) + result.get("_dead_generated", 0)
-        result["tokens_generated"] = gen + result["chk_forwards"]
-        result["tokens_sampled"] = gen
-        return result
-
-    @torch.no_grad()
-    def greedy1(self, prompts, strat=STRAT_IDS[0], max_new=512):
-        res, streams = self.run_streams(prompts, [[strat]] * len(prompts), temperature=0.0, max_new=max_new, use_judge=False)
-        out = []
-        for i, r in enumerate(res):
-            r = self._account(r, streams[i]); r["answer"] = extract_answer(r["streams"][0]["text"]); out.append(r)
-        return out
-
-    @torch.no_grad()
-    def bestofK(self, prompts, K=3, temperature=0.7, max_new=512, seed=1337,
-                sample_strats: Optional[List[int]] = None, sample_temps: Optional[List[float]] = None):
-        """K independent sampled streams (strategies round-robin, or sample_strats[j] per sample; temperature, or
-        sample_temps[j] per sample, 0 = greedy), majority vote on the final answer.
-        Implemented as K single-stream rows sharing nothing but the code path (prompt prefill is cheap)."""
-        rows = [p for p in prompts for _ in range(K)]
-        strats = [[sample_strats[j] if sample_strats else STRAT_IDS[j % 3]] for _ in prompts for j in range(K)]
-        row_temps = [[sample_temps[j]] for _ in prompts for j in range(K)] if sample_temps else None
-        res, streams = self.run_streams(rows, strats, temperature=temperature, max_new=max_new, use_judge=False, sample_seed=seed,
-                                        row_temps=row_temps)
-        out = []
-        for i in range(len(prompts)):
-            grp = res[i * K:(i + 1) * K]; gs = streams[i * K:(i + 1) * K]
-            answers = [extract_answer(g["streams"][0]["text"]) for g in grp]
-            ans, ans_str = majority_vote(answers)     # value vote (expressions evaluated); string vote kept for reference
-            tg = sum(g["streams"][0]["generated_own"] for g in grp)
-            out.append({"answer": ans, "answer_stringvote": ans_str, "answers": answers, "tokens_generated": tg, "tokens_sampled": tg,
-                        "chk_forwards": 0, "streams": [g["streams"][0] for g in grp], "kills": 0, "forks": 0})
-        return out
-
-    @torch.no_grad()
-    def native(self, prompts, K=3, N=None, temperature=0.7, max_new=512, respawn=True, seed=1337,
-               slot_strats: Optional[List[int]] = None, slot_temps: Optional[List[float]] = None, strat_pool: Optional[List[int]] = None,
-               fork_temp: Optional[float] = None, prune: str = "always", tau: float = 0.0, delta: float = 0.0,
-               immune_slots: Optional[List[int]] = None, dump_tokens: bool = False):
-        """slot_strats: strategy token per slot (default round-robin over all three); slot_temps / strat_pool / fork_temp /
-        prune / tau / delta / immune_slots / dump_tokens as in run_streams."""
-        base = list(slot_strats) if slot_strats else [STRAT_IDS[j % 3] for j in range(K)]
-        strats = [list(base) for _ in prompts]
-        self._all_streams = []
-        res, streams = self._run_native(prompts, strats, temperature, max_new, respawn, N, seed, slot_temps, strat_pool, fork_temp,
-                                        prune, tau, delta, immune_slots, dump_tokens)
-        return res
-
-    def _run_native(self, prompts, strats, temperature, max_new, respawn, N, seed, slot_temps=None, strat_pool=None, fork_temp=None,
-                    prune="always", tau=0.0, delta=0.0, immune_slots=None, dump_tokens=False):
-        # run_streams replaces killed streams in place; we need the killed ones' generated counts, so wrap Stream creation
-        res, streams = self.run_streams(prompts, strats, temperature=temperature, max_new=max_new, use_judge=True,
-                                        respawn=respawn, N=N, sample_seed=seed, slot_temps=slot_temps, strat_pool=strat_pool,
-                                        fork_temp=fork_temp, prune=prune, tau=tau, delta=delta, immune_slots=immune_slots,
-                                        dump_tokens=dump_tokens)
-        out = []
-        for i, r in enumerate(res):
-            # generated by killed streams is recorded in judge_log lens at kill time; recompute from stream history
-            gen_live = sum(s.generated for s in streams[i])
-            gen_dead = r.get("_dead_generated", 0)
-            r["tokens_sampled"] = gen_live + gen_dead
-            r["tokens_generated"] = r["tokens_sampled"] + r["chk_forwards"]
-            fin = [s for s in r["streams"] if s["finished"] and not s["dead"]]
-            pick = max(fin, key=lambda s: (s["score"] if s["score"] == s["score"] else -1e9)) if fin else \
-                   max([s for s in r["streams"] if not s["dead"]], key=lambda s: (s["score"] if s["score"] == s["score"] else -1e9))
-            r["answer"] = extract_answer(pick["text"]); r["picked_slot"] = pick["slot"]
-            # value vote over the surviving streams' final answers (same samples as the judge pick): isolates the judge's
-            # selection value from the pruning; identical to best-of-K's vote when prune == "none"
-            live_answers = [extract_answer(s["text"]) for s in r["streams"] if not s["dead"]]
-            r["answer_vote"], r["answer_stringvote"] = majority_vote(live_answers); r["answers"] = live_answers
-            r["anchor_killed"] = any(d["slot"] == 0 and d.get("gen", 0) == 0 for d in r["dead_streams"])
-            out.append(r)
-        return out, streams
-
-
-def extract_answer(text: str):
-    ms = ANSWER_RE.findall(text or "")
-    return ms[-1].strip() if ms else None
-
-
-def answer_key(ans):
-    """Canonical vote key for majority voting. An arithmetic expression (Countdown-style answer) is keyed by its
-    evaluated value — computed from the expression alone, without the gold target and without the uses-each-number-
-    once check — so distinct valid expressions for the same target agree. A number is keyed by its normalised
-    value; anything else by its stripped string. None stays None."""
-    from fractions import Fraction
-    from maram.data.label import normalize_expr, _fractionize, norm_number
-    if ans is None:
-        return None
-    e = normalize_expr(ans)
-    if e and re.fullmatch(r"[\d\s\+\-\*/\(\)]+", e) and re.search(r"[\+\-\*/]", e) and re.search(r"\d", e):
-        try:
-            v = eval(_fractionize(e), {"Fraction": Fraction, "__builtins__": {}})
-            return f"v:{v}"
-        except Exception:
-            pass
-    n = norm_number(ans)
-    return f"n:{n}" if n is not None else f"s:{ans.strip()}"
-
-
-def majority_vote(answers):
-    """Returns (winning raw answer, string-vote answer). Ties: first seen."""
-    keys = [answer_key(a) for a in answers]
-    counts = {}
-    for k in keys:
-        if k is not None:
-            counts[k] = counts.get(k, 0) + 1
-    win = max(counts, key=counts.get) if counts else None
-    ans = next((a for a, k in zip(answers, keys) if k == win), None) if win is not None else None
-    sc = {}
-    for a in answers:
-        if a is not None:
-            sc[a.strip()] = sc.get(a.strip(), 0) + 1
-    return ans, (max(sc, key=sc.get) if sc else None)
-
-
-def load_model(ckpt_path, device="cuda", judge_mode="full"):
-    sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-    cfg = ModelConfig.from_dict(sd["model_cfg"])
-    cfg.judge_mode = judge_mode
-    m = MaramModel(cfg); m.load_state_dict(sd["model"])
-    return m.to(device), sd