bench/bench_ledger_needle.py

"""Ledger needle bench: fact recall across a conversation far longer than the context.

The long-context analogue for this architecture. A needle-in-a-haystack test measures whether a
fact survives distance inside a huge prompt; here the prompt never exceeds the training window
(last ``keep`` messages) and facts survive only through the ledger of the model's own notes, so
the same question — does distance kill recall? — is asked of the memory mechanism instead of the
attention span.

``run`` drives a deterministic scripted conversation through the engine on the GPU box and
records everything per turn. ``render`` turns that record into a self-contained HTML viewer:
timeline, per-turn notes, the exact ledger state, and every probe scored.
"""

from __future__ import annotations

import argparse
import html as html_lib
import json
import random
import re
import time
from pathlib import Path

FACT_POOL = [
    ('sister flight', 'flight {code}', 'my sister lands on {value}, write that down'),
    ('hotel room', 'room {num3}', 'we got {value} at the hotel'),
    ('locker code', 'code {num4}', 'the gym locker is {value}'),
    ('dentist slot', '{clock}', 'dentist moved me to {value}'),
    ('plumber quote', '${money}', 'the plumber quoted {value} for the bathroom'),
    ('car plate', '{plate}', 'the rental has plate {value}'),
    ('wifi password', '{word}{num3}', 'cabin wifi password is {value}'),
    ('train platform', 'platform {num2}', 'our train leaves from {value}'),
    ('order number', 'order {code}', 'the couch is {value}, keep it handy'),
    ('rent due', 'day {num2}', 'landlord wants rent by {value} each month'),
    ('kid teacher', 'ms. {name}', "tomas' new teacher is {value}"),
    ('parking spot', 'spot {letter}{num2}', 'we always park at {value}'),
    ('meds dose', '{num2}mg', 'doctor changed the dose to {value}'),
    ('deposit', '${money}', 'the deposit for the venue was {value}'),
    ('boarding gate', 'gate {letter}{num2}', 'boarding is at {value}'),
]
DISTRACTORS = [
    'what a week, honestly', 'did you catch the game last night?', 'i love this weather lately',
    'work was chaos today', 'thinking of making pasta tonight', 'my back hurts from the gym',
    'the neighbor is renovating again, so loud', 'saw a great documentary yesterday',
    'coffee here is getting expensive', 'might go for a walk later', 'the cat knocked over a plant',
    'traffic was unreal this morning', 'finally finished that book', 'craving something sweet',
    'they repaved our street', 'my phone battery dies so fast now',
]
PROBES = ['wait, what was the {key} again?', 'remind me of the {key}?',
          'i forgot the {key}, what was it?', 'quick — the {key}?']
WORDS = ('maple', 'harbor', 'cactus', 'violet', 'ember', 'quartz')
NAMES = ('ferro', 'silva', 'duarte', 'campos', 'rojas', 'ibanez')


def _value(template: str, rng: random.Random) -> str:
    return (template
            .replace('{code}', f'{rng.choice("ABKQZ")}{rng.choice("RLMT")}{rng.randint(1000, 9999)}')
            .replace('{num2}', str(rng.randint(10, 99)))
            .replace('{num3}', str(rng.randint(100, 999)))
            .replace('{num4}', str(rng.randint(1000, 9999)))
            .replace('{clock}', f'{rng.randint(1, 12)}:{rng.choice(("15", "30", "45"))}pm')
            .replace('{money}', f'{rng.randint(2, 90) * 100:,}')
            .replace('{plate}', f'{"".join(rng.choice("BCDFGHJK") for _ in range(3))}-{rng.randint(100, 999)}')
            .replace('{word}', rng.choice(WORDS))
            .replace('{letter}', rng.choice('ABCDE'))
            .replace('{name}', rng.choice(NAMES)))


def build_script(turns: int, seed: int) -> list[dict]:
    """Deterministic conversation plan: facts early and throughout, probes at all distances."""

    rng = random.Random(seed)
    facts = []
    for key, template, phrasing in rng.sample(FACT_POOL, len(FACT_POOL)):
        facts.append({'key': key, 'value': _value(template, rng), 'phrasing': phrasing})

    plan: list[dict] = []
    stated: dict[str, dict] = {}
    fact_iter = iter(facts)
    for index in range(turns):
        remaining = turns - index
        can_probe = [f for f in stated.values() if index - f['turn'] >= 3]
        if index >= turns - 3 and can_probe:
            kind = 'probe'
        elif index % 7 in (0, 3) and (fact := next(fact_iter, None)) is not None:
            plan.append({'kind': 'fact', **fact})
            stated[fact['key']] = {**fact, 'turn': index}
            continue
        elif index % 11 == 5 and can_probe:
            kind = 'probe'
        elif index % 13 == 8 and stated and remaining > 5:
            kind = 'correction'
        else:
            kind = 'distractor'

        if kind == 'probe':
            target = rng.choice(can_probe)
            plan.append({'kind': 'probe', 'key': target['key'], 'value': target['value'],
                         'distance': index - target['turn'],
                         'text': rng.choice(PROBES).format(key=target['key'])})
        elif kind == 'correction':
            target = rng.choice(list(stated.values()))
            new = _value(next(t for k, t, _ in FACT_POOL if k == target['key']), rng)
            plan.append({'kind': 'correction', 'key': target['key'], 'old': target['value'],
                         'value': new,
                         'text': f'actually scratch that, the {target["key"]} is {new} now'})
            stated[target['key']] = {**target, 'value': new, 'turn': index}
        else:
            plan.append({'kind': 'distractor', 'text': rng.choice(DISTRACTORS)})
    return plan


def _hit(value: str, answer: str) -> bool:
    """Whether the answer states the value, matched on its distinctive core.

    Demanding the full phrase penalised correct answers on both systems ("the locker code is
    4805" failed against 'code 4805'), so the core — the last token, which carries the
    identifier — is what must appear. Values are generated with 2+ digit cores, so incidental
    collisions stay unlikely.
    """

    canon = lambda t: re.sub(r'[^a-z0-9]', '', t.lower())  # noqa: E731
    core = value.split()[-1]
    return canon(core) in canon(answer)


def run(args: argparse.Namespace) -> None:
    from diffusion_lm.claims import SYSTEM, chat_prefix, ledger_line, ledger_notes, merge_notes
    from diffusion_lm.reasoning_playground import ReasoningEngine
    from diffusion_lm.train import resolve_device

    engine = ReasoningEngine(args.checkpoint, args.tokenizer, resolve_device('auto'))
    plan = build_script(args.turns, args.seed)
    messages: list[dict] = []
    records = []
    # Appended per turn and flushed, so a tail -f (or a crash) sees every completed turn.
    sink = args.output.open('w', encoding='utf-8')
    for index, turn in enumerate(plan):
        user_text = turn.get('text') or turn['phrasing'].format(value=turn['value'])
        messages.append({'role': 'user', 'content': user_text})
        older = merge_notes(ledger_notes(messages, args.keep))
        window = messages[max(0, len(messages) - args.keep):]
        prefix = chat_prefix(
            [{'role': m['role'], 'content': m['content']} for m in window],
            system=SYSTEM, extra=ledger_line(older),
        )
        blocks: list[tuple[int, str]] = []
        answer = ''
        begin = time.perf_counter()
        for _, answer, _ in engine.stream_chat(
            prefix, temperature=0.7, steps_per_block=args.steps,
            max_answer_tokens=120, seed=args.seed * 1000 + index, blocks_out=blocks,
        ):
            pass
        note = '; '.join(text for _, text in blocks if text)
        messages.append({'role': 'assistant', 'content': answer.strip(), 'note': note})
        record = {
            'turn': index, 'kind': turn['kind'], 'user': user_text,
            'notes': [text for _, text in blocks if text], 'ledger': older,
            'prefix_tokens': len(engine.tokenizer.encode(prefix, add_special_tokens=False).ids),
            'answer': answer.strip(), 'seconds': round(time.perf_counter() - begin, 2),
        }
        if turn['kind'] == 'probe':
            record.update({'key': turn['key'], 'expected': turn['value'],
                           'distance': turn['distance'],
                           'hit': _hit(turn['value'], answer)})
        if turn['kind'] in ('fact', 'correction'):
            record.update({'key': turn['key'], 'value': turn['value']})
        records.append(record)
        sink.write(json.dumps(record, ensure_ascii=False) + '\n')
        sink.flush()
        flag = '' if turn['kind'] != 'probe' else (' HIT' if record['hit'] else ' MISS')
        print(f'[{index + 1}/{len(plan)}] {turn["kind"]}{flag} {record["seconds"]}s '
              f'ledger={len(older)}', flush=True)
        print(f'  U: {user_text}', flush=True)
        if note:
            print(f'  T: {note}', flush=True)
        print(f'  A: {answer.strip()}', flush=True)
        if turn['kind'] == 'probe':
            print(f'  >> esperado {turn["value"]} (distancia {turn["distance"]}) -> '
                  f'{"HIT" if record["hit"] else "MISS"}', flush=True)

    sink.close()
    probes = [r for r in records if r['kind'] == 'probe']
    hits = sum(r['hit'] for r in probes)
    print(f'\nrecall: {hits}/{len(probes)} ({hits / max(1, len(probes)):.0%})')


def run_ar(args: argparse.Namespace) -> None:
    """Same scripted conversation through a plain AR baseline of the same size.

    ``--context full`` gives the baseline the whole history (its native memory);
    ``--context budget`` truncates the rendered prompt to the newest messages that fit the
    same token budget the ledger system runs under, so both systems answer from equally many
    prefix tokens and only the memory MECHANISM differs.
    """

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path, torch_dtype=torch.bfloat16, device_map='cuda',
    ).eval()
    plan = build_script(args.turns, args.seed)
    # Deliberately generous: the baselines get an explicit memory instruction, greedy decoding,
    # and (with --thinking) their native reasoning mode. The comparison should show the ledger
    # beating the baselines at their best, or it shows nothing.
    system = ('You are a meticulous assistant. Remember every concrete detail the user '
              'mentions — codes, amounts, times, names, numbers. When asked to recall one, '
              'answer with the exact value.')
    messages: list[dict] = []
    records = []
    try:
        tokenizer.apply_chat_template(
            [{'role': 'system', 'content': 'x'}, {'role': 'user', 'content': 'y'}],
            tokenize=True, add_generation_prompt=True)
        args._no_system = False
    except Exception:
        args._no_system = True
    sink = args.output.open('w', encoding='utf-8')
    for index, turn in enumerate(plan):
        user_text = turn.get('text') or turn['phrasing'].format(value=turn['value'])
        messages.append({'role': 'user', 'content': user_text})

        def _assemble(tail: list[dict]) -> list[dict]:
            if getattr(args, '_no_system', False):
                head = dict(tail[0]) if tail else {'role': 'user', 'content': ''}
                head['content'] = f'{system}\n\n{head["content"]}'
                return [head] + [dict(m) for m in tail[1:]]
            return [{'role': 'system', 'content': system}] + [dict(m) for m in tail]

        chat = _assemble(messages)
        if args.context == 'budget':
            # Newest-first packing under the same prefix budget the ledger system uses.
            kept: list[dict] = []
            for message in reversed(messages):
                candidate = _assemble([message, *kept])
                rendered = tokenizer.apply_chat_template(
                    candidate, tokenize=True, add_generation_prompt=True,
                    enable_thinking=False,
                )
                # BatchEncoding in newer transformers: len() counts KEYS, which silently
                # disabled the cap; count tokens explicitly.
                ids = rendered if isinstance(rendered, list) else rendered['input_ids']
                if len(ids) > args.budget:
                    break
                kept = [message, *kept]
            chat = _assemble(kept)
        encoded = tokenizer.apply_chat_template(
            chat, tokenize=True, add_generation_prompt=True, return_tensors='pt',
            enable_thinking=args.thinking,
        )
        # Newer transformers return a BatchEncoding here rather than the bare tensor.
        prompt_ids = (encoded if torch.is_tensor(encoded) else encoded['input_ids']).to(
            model.device)
        torch.manual_seed(args.seed * 1000 + index)
        begin = time.perf_counter()
        with torch.no_grad():
            out = model.generate(
                prompt_ids, max_new_tokens=640 if args.thinking else 120, do_sample=False,
                pad_token_id=tokenizer.eos_token_id or tokenizer.pad_token_id,
            )
        answer = tokenizer.decode(out[0][prompt_ids.shape[1]:], skip_special_tokens=True)
        answer = re.sub(r'(?s)<think>.*?(</think>|$)', '', answer).strip()
        messages.append({'role': 'assistant', 'content': answer})
        record = {
            'turn': index, 'kind': turn['kind'], 'user': user_text, 'notes': [], 'ledger': [],
            'prefix_tokens': int(prompt_ids.shape[1]), 'answer': answer,
            'seconds': round(time.perf_counter() - begin, 2),
        }
        if turn['kind'] == 'probe':
            record.update({'key': turn['key'], 'expected': turn['value'],
                           'distance': turn['distance'], 'hit': _hit(turn['value'], answer)})
        if turn['kind'] in ('fact', 'correction'):
            record.update({'key': turn['key'], 'value': turn['value']})
        records.append(record)
        sink.write(json.dumps(record, ensure_ascii=False) + '\n')
        sink.flush()
        flag = '' if turn['kind'] != 'probe' else (' HIT' if record['hit'] else ' MISS')
        print(f'[{index + 1}/{len(plan)}] {turn["kind"]}{flag} {record["seconds"]}s '
              f'prefix={record["prefix_tokens"]}', flush=True)
    sink.close()
    probes = [r for r in records if r['kind'] == 'probe']
    hits = sum(r['hit'] for r in probes)
    print(f'\nrecall: {hits}/{len(probes)} ({hits / max(1, len(probes)):.0%})')


TEMPLATE = '''<!doctype html><html><head><meta charset="utf-8">
<title>Ledger needle — __MODEL__</title><style>
:root { --bg:#11141a; --panel:#1a1f29; --text:#dde3ee; --dim:#8a93a6; --hit:#3fb96b;
  --miss:#e05555; --fact:#4d8fd1; --corr:#d9a23c; --line:#2a3140; }
* { box-sizing:border-box; margin:0 } body { background:var(--bg); color:var(--text);
  font:14px/1.5 ui-monospace,Menlo,monospace; padding:24px; }
h1 { font-size:18px; margin-bottom:4px } .sub { color:var(--dim); margin-bottom:18px }
.metrics { display:flex; gap:14px; flex-wrap:wrap; margin-bottom:18px }
.metric { background:var(--panel); border:1px solid var(--line); border-radius:8px;
  padding:10px 16px } .metric b { font-size:20px; display:block }
.dist { margin-bottom:18px } .dist .row { display:flex; align-items:center; gap:8px;
  margin:3px 0 } .dist .bar { height:14px; background:var(--hit); border-radius:3px }
.dist .bar.m { background:var(--miss) } .dist span { color:var(--dim); font-size:12px;
  min-width:110px }
.wrap { display:grid; grid-template-columns:340px 1fr; gap:16px; align-items:start }
.timeline { max-height:75vh; overflow-y:auto; background:var(--panel);
  border:1px solid var(--line); border-radius:8px }
.t { padding:7px 10px; border-bottom:1px solid var(--line); cursor:pointer;
  display:flex; gap:8px; align-items:center } .t:hover,.t.sel { background:#242b38 }
.t .n { color:var(--dim); min-width:30px } .dot { width:9px; height:9px; border-radius:50%;
  flex:none } .dot.distractor { background:var(--dim) } .dot.fact { background:var(--fact) }
.dot.correction { background:var(--corr) } .dot.probe.hit { background:var(--hit) }
.dot.probe.miss { background:var(--miss) } .t .txt { white-space:nowrap; overflow:hidden;
  text-overflow:ellipsis; font-size:12px }
.detail { background:var(--panel); border:1px solid var(--line); border-radius:8px;
  padding:16px; max-height:75vh; overflow-y:auto }
.detail h3 { font-size:13px; color:var(--dim); margin:14px 0 6px; text-transform:uppercase }
.detail h3:first-child { margin-top:0 }
.bubble { background:#242b38; border-radius:8px; padding:10px 12px; margin:4px 0 }
.note { color:#9fc4ea } .ledger-entry { display:inline-block; background:#242b38;
  border:1px solid var(--line); border-radius:5px; padding:2px 8px; margin:2px;
  font-size:12px } .verdict { padding:10px 12px; border-radius:8px; margin-top:6px }
.verdict.hit { background:#173523; border:1px solid var(--hit) }
.verdict.miss { background:#3a1d1d; border:1px solid var(--miss) }
.meta { color:var(--dim); font-size:12px; margin-top:12px }
</style></head><body>
<h1>Ledger needle</h1>
<div class="sub">__MODEL__ · __TURNS__ turnos · ventana de __KEEP__ mensajes ·
steps_per_block __STEPS__ · el prefijo nunca excede __MAXPREFIX__ tokens</div>
<div class="metrics" id="metrics"></div>
<div class="dist" id="dist"></div>
<div class="wrap"><div class="timeline" id="timeline"></div>
<div class="detail" id="detail">elegí un turno</div></div>
<script>
const DATA = __DATA__;
const probes = DATA.filter(r => r.kind === 'probe');
const hits = probes.filter(r => r.hit).length;
const corr = DATA.filter(r => r.kind === 'correction').length;
const maxTok = Math.max(...DATA.map(r => r.prefix_tokens));
const secs = DATA.reduce((a, r) => a + r.seconds, 0) / DATA.length;
document.getElementById('metrics').innerHTML = [
  ['recall', hits + '/' + probes.length + ' (' + Math.round(100 * hits / probes.length) + '%)'],
  ['hechos plantados', DATA.filter(r => r.kind === 'fact').length],
  ['correcciones', corr], ['prefijo máx', maxTok + ' tok'],
  ['media', secs.toFixed(1) + ' s/turno'],
].map(([k, v]) => '<div class="metric"><b>' + v + '</b>' + k + '</div>').join('');
const buckets = [[3, 10], [11, 25], [26, 50], [51, 999]];
document.getElementById('dist').innerHTML = '<span style="color:var(--dim)">recall por distancia (turnos desde que el hecho se dijo):</span>' +
  buckets.map(([lo, hi]) => {
    const set = probes.filter(p => p.distance >= lo && p.distance <= hi);
    if (!set.length) return '';
    const h = set.filter(p => p.hit).length;
    return '<div class="row"><span>' + lo + '–' + (hi > 100 ? '∞' : hi) + ' (' + h + '/' + set.length +
      ')</span><div class="bar" style="width:' + (300 * h / set.length) + 'px"></div>' +
      '<div class="bar m" style="width:' + (300 * (set.length - h) / set.length) + 'px"></div></div>';
  }).join('');
const tl = document.getElementById('timeline');
DATA.forEach(r => {
  const div = document.createElement('div');
  div.className = 't'; div.dataset.turn = r.turn;
  const cls = r.kind + (r.kind === 'probe' ? (r.hit ? ' hit' : ' miss') : '');
  div.innerHTML = '<span class="n">' + (r.turn + 1) + '</span><span class="dot ' + cls +
    '"></span><span class="txt">' + r.user.replace(/</g, '&lt;') + '</span>';
  div.onclick = () => show(r, div); tl.appendChild(div);
});
function esc(t) { return String(t).replace(/</g, '&lt;'); }
function show(r, el) {
  document.querySelectorAll('.t.sel').forEach(x => x.classList.remove('sel'));
  el.classList.add('sel');
  let h = '<h3>usuario (turno ' + (r.turn + 1) + ' · ' + r.kind + ')</h3><div class="bubble">' +
    esc(r.user) + '</div>';
  h += '<h3>ledger que vio el modelo (' + r.ledger.length + ' entradas, historia fuera de ventana)</h3>' +
    (r.ledger.length ? r.ledger.map(e => '<span class="ledger-entry">' + esc(e) + '</span>').join('')
      : '<span style="color:var(--dim)">vacío</span>');
  h += '<h3>notas de este turno</h3>' + (r.notes.length
    ? r.notes.map(n => '<div class="bubble note">' + esc(n) + '</div>').join('')
    : '<span style="color:var(--dim)">sin pensamiento</span>');
  h += '<h3>respuesta</h3><div class="bubble">' + esc(r.answer) + '</div>';
  if (r.kind === 'probe') h += '<div class="verdict ' + (r.hit ? 'hit' : 'miss') + '">esperado: <b>' +
    esc(r.expected) + '</b> · distancia ' + r.distance + ' turnos · ' + (r.hit ? 'RECUPERADO' : 'PERDIDO') + '</div>';
  if (r.kind === 'correction') h += '<div class="verdict" style="border:1px solid var(--corr)">corrige ' +
    esc(r.key) + ': ahora <b>' + esc(r.value) + '</b></div>';
  h += '<div class="meta">prefijo ' + r.prefix_tokens + ' tokens · ' + r.seconds + ' s</div>';
  document.getElementById('detail').innerHTML = h;
}
</script></body></html>'''


def _rescore(records: list[dict]) -> list[dict]:
    """Verdicts are re-derived from the CURRENT rule, never trusted from the run.

    The stored ``hit`` reflects whatever rule was live at run time; a scoring fix must reach
    old runs, or the viewers keep showing correct answers marked as losses.
    """

    for record in records:
        if 'expected' in record:
            record['hit'] = _hit(record['expected'], record['answer'])
    return records


def render(args: argparse.Namespace) -> None:
    records = _rescore(
        [json.loads(line) for line in args.input.open(encoding='utf-8') if line.strip()])
    page = (TEMPLATE
            .replace('__DATA__', json.dumps(records, ensure_ascii=False))
            .replace('__MODEL__', html_lib.escape(args.model_name))
            .replace('__TURNS__', str(len(records)))
            .replace('__KEEP__', str(args.keep))
            .replace('__STEPS__', str(args.steps))
            .replace('__MAXPREFIX__', str(max(r['prefix_tokens'] for r in records))))
    args.output.write_text(page, encoding='utf-8')
    print(f'viewer -> {args.output}')


COMPARE_TEMPLATE = '''<!doctype html><html><head><meta charset="utf-8">
<title>Ledger needle — comparación</title><style>
:root { --bg:#11141a; --panel:#1a1f29; --text:#dde3ee; --dim:#8a93a6; --hit:#3fb96b;
  --miss:#e05555; --line:#2a3140; --us:#16233a; }
* { box-sizing:border-box; margin:0 } body { background:var(--bg); color:var(--text);
  font:14px/1.5 ui-monospace,Menlo,monospace; padding:24px; max-width:1240px; margin:0 auto }
h1 { font-size:18px } .sub { color:var(--dim); margin:4px 0 20px }
table { border-collapse:collapse; width:100%; margin-bottom:26px }
th, td { border:1px solid var(--line); padding:7px 10px; text-align:left; font-size:13px }
th { background:var(--panel); color:var(--dim); font-weight:normal; cursor:pointer;
  user-select:none; white-space:nowrap }
th.sorted { color:var(--text) } th.sorted::after { content:' ↓' }
th.sorted.asc::after { content:' ↑' }
tr.us td { background:var(--us) }
td.hitcell { background:#173523; color:var(--hit); text-align:center; font-weight:bold }
td.misscell { background:#3a1d1d; color:var(--miss); text-align:center }
#matrix { width:auto } #matrix th, #matrix td { padding:5px 8px }
#matrix td.hitcell, #matrix td.misscell { width:36px; min-width:36px; padding:5px 0 }
#matrix th.sys { writing-mode:vertical-rl; transform:rotate(180deg); text-align:left;
  vertical-align:bottom; max-height:190px; font-size:12px; cursor:default; padding:8px 4px }
#matrix th.sys.usc { color:#9fc4ea }
#matrix td.exp { white-space:nowrap }
.bar { display:inline-block; height:11px; background:var(--hit); border-radius:2px;
  vertical-align:middle } .bar.m { background:var(--miss) }
h2 { font-size:14px; color:var(--dim); margin:22px 0 10px; text-transform:uppercase }
.caveats { background:var(--panel); border:1px solid var(--line); border-radius:8px;
  padding:14px 18px; color:var(--dim); font-size:13px } .caveats li { margin:6px 0 0 18px }
.hint { color:var(--dim); font-size:12px; margin:-18px 0 10px }
</style></head><body>
<h1>Ledger needle — comparación hasta 1.5B</h1>
<div class="sub">100 turnos guionados (seed 21) · 15 hechos, 5 correcciones, 10 sondas a
distancia 3–96 · scoring por núcleo del valor, idéntico para todos · baselines con historia
completa, decoding greedy e instrucción explícita de memoria</div>
<table id="summary"><thead><tr>
<th data-k="label" data-t="s">sistema</th>
<th data-k="params" data-t="n">tamaño (B)</th>
<th data-k="hits" data-t="n">recall</th>
<th data-k="d1" data-t="n">d 3–10</th>
<th data-k="d2" data-t="n">d 11–50</th>
<th data-k="d3" data-t="n">d 51+</th>
<th data-k="total" data-t="n">test total (s)</th>
<th data-k="per" data-t="n">s/turno</th>
<th data-k="prefix" data-t="n">prefijo máx</th>
</tr></thead><tbody></tbody></table>
<div class="hint">click en una cabecera reordena; por defecto: más aciertos primero</div>
<h2>Sonda por sonda</h2>
__MATRIX__
<h2>Advertencias de lectura</h2>
<div class="caveats"><ul>
<li>Un solo seed y 10 sondas por sistema: las diferencias chicas son ruido; las grandes
sobreviven a ese margen.</li>
<li>La comparación mide el paquete mecanismo+entrenamiento: el sistema de ledger fue entrenado
en este registro conversacional y los baselines no. Muestra que el paquete funciona, no que el
mecanismo solo cause la diferencia.</li>
<li>Los baselines reciben la mayor ventaja razonable: historia completa en contexto, greedy,
instrucción de memoria, y el modo thinking nativo donde existe. El sistema de ledger corre con
sampling a temperatura 0.7 y un prefijo constante de ~500 tokens.</li>
<li>El guion es sintético y sus valores son arbitrarios por diseño; un valor sin sentido
("day 97") puede penalizar a un modelo que se resiste a repetir datos absurdos.</li>
</ul></div>
<script>
const ROWS = __ROWS__;
const tbody = document.querySelector('#summary tbody');
let sortKey = 'hits', asc = false;
function fmt(r) {
  return '<td><b>' + r.label + '</b></td><td>' + r.params.toFixed(2) + '</td>' +
    '<td><b>' + r.hits + '/' + r.probes + '</b> <span class="bar" style="width:' + 6 * r.hits +
    'px"></span><span class="bar m" style="width:' + 6 * (r.probes - r.hits) + 'px"></span></td>' +
    ['d1', 'd2', 'd3'].map(k => '<td>' + r[k] + '/' + r[k + 'n'] + '</td>').join('') +
    '<td>' + r.total.toFixed(0) + '</td><td>' + r.per.toFixed(1) + '</td><td>' + r.prefix + '</td>';
}
function draw() {
  const rows = [...ROWS].sort((a, b) => {
    const va = a[sortKey], vb = b[sortKey];
    const c = typeof va === 'string' ? va.localeCompare(vb) : va - vb;
    return asc ? c : -c;
  });
  tbody.innerHTML = rows.map(r =>
    '<tr' + (r.us ? ' class="us"' : '') + '>' + fmt(r) + '</tr>').join('');
  document.querySelectorAll('#summary th').forEach(th => {
    th.classList.toggle('sorted', th.dataset.k === sortKey);
    th.classList.toggle('asc', th.dataset.k === sortKey && asc);
  });
}
document.querySelectorAll('#summary th').forEach(th => th.onclick = () => {
  if (sortKey === th.dataset.k) asc = !asc;
  else { sortKey = th.dataset.k; asc = th.dataset.t === 's'; }
  draw();
});
draw();
</script></body></html>'''


def compare(args: argparse.Namespace) -> None:
    systems = []
    for spec in args.inputs:
        label, _, rest = spec.partition('=')
        path, _, params = rest.partition('=')
        records = _rescore(
            [json.loads(line) for line in Path(path).open(encoding='utf-8') if line.strip()])
        probes = [r for r in records if r['kind'] == 'probe']
        row = {'label': html_lib.escape(label), 'params': float(params or 0),
               'probes': len(probes), 'hits': sum(r['hit'] for r in probes),
               'total': sum(r['seconds'] for r in records),
               'per': sum(r['seconds'] for r in records) / max(1, len(records)),
               'prefix': max(r['prefix_tokens'] for r in records),
               'us': bool(args.ours and label == args.ours)}
        for name, (lo, hi) in (('d1', (3, 10)), ('d2', (11, 50)), ('d3', (51, 999))):
            subset = [r for r in probes if lo <= r['distance'] <= hi]
            row[name] = sum(r['hit'] for r in subset)
            row[name + 'n'] = len(subset)
        systems.append((label, records, row))

    first_probes = [r for r in systems[0][1] if r['kind'] == 'probe']
    head = ['<tr><th>turno</th><th>dist</th><th>esperado</th>']
    head += [
        f'<th class="sys{" usc" if args.ours and label == args.ours else ""}">'
        f'{html_lib.escape(label)}</th>'
        for label, _, _ in systems
    ]
    rows_html = []
    for probe in first_probes:
        row = [f'<td>{probe["turn"] + 1}</td><td>{probe["distance"]}</td>'
               f'<td class="exp">{html_lib.escape(probe["expected"])}</td>']
        for _, records, _ in systems:
            match = next((r for r in records if r['kind'] == 'probe'
                          and r['turn'] == probe['turn']), None)
            if match is None:
                row.append('<td>—</td>')
            else:
                row.append('<td class="hitcell">✓</td>' if match['hit']
                           else '<td class="misscell">✗</td>')
        rows_html.append(f'<tr>{"".join(row)}</tr>')
    matrix = f'<table id="matrix">{"".join(head)}</tr>{"".join(rows_html)}</table>'

    page = (COMPARE_TEMPLATE
            .replace('__ROWS__', json.dumps([row for _, _, row in systems], ensure_ascii=False))
            .replace('__MATRIX__', matrix))
    args.output.write_text(page, encoding='utf-8')
    print(f'comparison -> {args.output}')


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='command', required=True)
    runner = sub.add_parser('run', help='drive the scripted conversation through the engine')
    runner.add_argument('--checkpoint', type=Path,
                        default=Path('outputs/qwen06b-genchat-sft/inference-latest.pt'))
    runner.add_argument('--tokenizer', type=Path,
                        default=Path('artifacts/tokenizer-qwen3-adaptive.json'))
    runner.add_argument('--turns', type=int, default=100)
    runner.add_argument('--keep', type=int, default=4)
    runner.add_argument('--steps', type=int, default=16)
    runner.add_argument('--seed', type=int, default=7)
    runner.add_argument('--output', type=Path, default=Path('bench-ledger-needle.jsonl'))
    ar = sub.add_parser('run-ar', help='same script through a plain AR baseline')
    ar.add_argument('--model-path', type=Path, default=Path('/root/models/Qwen3-0.6B'))
    ar.add_argument('--turns', type=int, default=100)
    ar.add_argument('--seed', type=int, default=7)
    ar.add_argument('--context', choices=('full', 'budget'), default='full')
    ar.add_argument('--budget', type=int, default=512)
    ar.add_argument('--thinking', action='store_true',
                    help='enable the native reasoning mode where the chat template supports it')
    ar.add_argument('--output', type=Path, default=Path('bench-needle-ar.jsonl'))

    cmp_ = sub.add_parser('compare', help='comparison table across systems as HTML')
    cmp_.add_argument('--inputs', nargs='+', required=True, metavar='LABEL=PATH[=PARAMS_B]')
    cmp_.add_argument('--ours', default=None, help='label to highlight as our system')
    cmp_.add_argument('--output', type=Path, default=Path('docs/ledger-needle-compare.html'))

    view = sub.add_parser('render', help='write the self-contained HTML viewer')
    view.add_argument('--input', type=Path, default=Path('bench-ledger-needle.jsonl'))
    view.add_argument('--output', type=Path, default=Path('docs/ledger-needle.html'))
    view.add_argument('--model-name', default='qwen06b-genchat-sft')
    view.add_argument('--keep', type=int, default=4)
    view.add_argument('--steps', type=int, default=16)
    args = parser.parse_args()
    if args.command == 'run':
        run(args)
    elif args.command == 'run-ar':
        run_ar(args)
    elif args.command == 'compare':
        compare(args)
    else:
        render(args)


if __name__ == '__main__':
    main()