LFM3-1.7B inference modules: experimental release (part 2)

research/scripts/rate.py

@@ -0,0 +1,7 @@
+import json, time, urllib.request
+def snap():
+    s = json.loads(urllib.request.urlopen("http://localhost:8080/slots").read())
+    return {x["id"]: (x.get("id_task"), x.get("next_token", [{}])[0].get("n_decoded", 0) if isinstance(x.get("next_token"), list) else x.get("next_token", {}).get("n_decoded", 0)) for x in s}
+a = snap(); time.sleep(20); b = snap()
+d = sum(b[k][1] - a[k][1] for k in b if a[k][0] == b[k][0] and b[k][1] >= a[k][1])
+print("approx tok/s:", d / 20, "busy:", sum(1 for k in b if b[k][1] > 0))

research/scripts/rbud_test.py

@@ -0,0 +1,32 @@
+"""Paired check: bare llama-server vs llama-server --reasoning-budget (no modules), math180, same seeds."""
+import json, sys, os, urllib.request
+from concurrent.futures import ThreadPoolExecutor
+sys.path.insert(0, os.path.dirname(__file__))
+from grade import grade_math, last_boxed
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+items = [json.loads(l) for l in open(os.path.join(ROOT, "data/coc/math180.jsonl"), encoding="utf-8")]
+MAXT = int(sys.argv[1]) if len(sys.argv) > 1 else 6144
+
+def ask(port, it):
+    body = {"messages": [{"role": "user", "content": it["prompt"]}], "max_tokens": MAXT, "temperature": 0.6,
+            "top_p": 0.95, "top_k": 50, "seed": 0}
+    r = json.loads(urllib.request.urlopen(urllib.request.Request(f"http://localhost:{port}/v1/chat/completions",
+        data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}), timeout=7200).read())
+    c = r["choices"][0]["message"].get("content") or ""
+    return {"id": it["id"], "ok": bool(grade_math(c, it["answer"])), "fin": r["choices"][0]["finish_reason"],
+            "tok": r["usage"]["completion_tokens"], "content": c[-400:]}
+
+def run(port, tag):
+    with ThreadPoolExecutor(8) as ex:
+        out = list(ex.map(lambda it: ask(port, it), items))
+    json.dump(out, open(os.path.join(ROOT, f"results/rbud/{tag}.json"), "w"), indent=1)
+    return out
+
+with ThreadPoolExecutor(2) as ex:
+    a, b = ex.map(lambda pt: run(*pt), [(8081, "bare"), (8082, "budget")])
+import random
+d = [y["ok"] - x["ok"] for x, y in zip(a, b)]
+bs = sorted(sum(random.choice(d) for _ in d) / len(d) * 100 for _ in range(5000))
+print(f"bare {100*sum(x['ok'] for x in a)/len(a):.1f}  budget {100*sum(y['ok'] for y in b)/len(b):.1f}  "
+      f"delta {100*sum(d)/len(d):+.1f} [{bs[125]:+.1f}, {bs[4875]:+.1f}]  no-answer bare {sum(not x['content'].strip() for x in a)} budget {sum(not y['content'].strip() for y in b)}")
+print("RBUD_DONE")

research/scripts/ref_check.py

@@ -0,0 +1,40 @@
+"""Reference solutions written by Claude for a LiveCodeBench sample (data/claude_ref/sample.jsonl, hidden tests
+removed), under the same conditions as the model: the statement and its public examples only.
+  ref_check.py show <n>      print the statement of sample problem n (0-based)
+  ref_check.py check <n>     run data/claude_ref/<id>.py on the PUBLIC tests of problem n
+  ref_check.py grade         grade every written solution on the HIDDEN tests (run once, at the end)
+"""
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from lcb_tests import run_lcb  # noqa: E402
+
+ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
+REF = os.path.join(ROOT, "data", "claude_ref")
+sample = [json.loads(l) for l in open(os.path.join(REF, "sample.jsonl"), encoding="utf-8")]
+path = lambda it: os.path.join(REF, it["id"].replace("/", "_") + ".py")
+
+cmd = sys.argv[1]
+if cmd == "show":
+    it = sample[int(sys.argv[2])]
+    print(f"[{sys.argv[2]}] {it['id']} {it['difficulty']} {it['testtype']} func={it.get('func_name')}\n")
+    print(it["prompt"])
+elif cmd == "check":
+    it = sample[int(sys.argv[2])]
+    k, n, fail = run_lcb(open(path(it), encoding="utf-8").read(), it, it["public"])
+    print(f"{it['id']}: public {k}/{n}" + (f"\n{fail}" if fail else ""))
+elif cmd == "grade":
+    full = {json.loads(l)["id"]: json.loads(l) for l in open(os.path.join(ROOT, "data", "lcb", "lcb6.jsonl"),
+                                                              encoding="utf-8")}
+    out = []
+    for it in sample:
+        if not os.path.exists(path(it)):
+            continue
+        k, n, _ = run_lcb(open(path(it), encoding="utf-8").read(), full[it["id"]], full[it["id"]]["hidden"])
+        out.append({"id": it["id"], "difficulty": it["difficulty"], "passed": k, "total": n, "correct": k == n})
+        print(f"{it['id']:16s} {it['difficulty']:7s} {k}/{n} {'OK' if k == n else 'FAIL'}", flush=True)
+    with open(os.path.join(REF, "graded.jsonl"), "w", encoding="utf-8") as f:
+        for g in out:
+            f.write(json.dumps(g) + "\n")

research/scripts/render.py

@@ -0,0 +1,122 @@
+"""Headless renderer for visual answers (SVG / HTML / canvas / three.js) with objective checks.
+
+render(kind, source, png_path) opens the page in headless Chrome (the installed browser, via playwright) at 1024x768,
+waits for it to settle and returns a report:
+  errors      JavaScript exceptions and console errors (messages)
+  failed      network requests that failed (e.g. a wrong library URL)
+  blank       the screenshot is (almost) a single colour
+  overflow    the page is wider than the viewport (horizontal scroll)
+  animated    two screenshots 0.7 s apart differ (for animation tasks)
+  coverage    share of the viewport that is not background colour
+  parse_error for SVG: the XML does not parse
+The screenshot is saved to png_path. Used by the visual benchmark and by the render-check-repair module.
+"""
+import io
+import os
+import re
+import tempfile
+import xml.dom.minidom
+
+import numpy as np
+from PIL import Image
+from playwright.sync_api import sync_playwright
+
+W, H = 1024, 768
+TMP = os.path.join(tempfile.gettempdir(), "lfm3_render")
+os.makedirs(TMP, exist_ok=True)
+
+SVG_PAGE = """<!doctype html><html><head><meta charset="utf-8"><style>html,body{margin:0;background:#fff}
+#c{width:1024px;height:768px;display:flex;align-items:center;justify-content:center}
+#c>svg{max-width:1024px;max-height:768px;width:100%;height:100%}</style></head><body><div id="c">__SVG__</div></body></html>"""
+
+
+def _stats(png):
+    a = np.asarray(Image.open(io.BytesIO(png)).convert("RGB")).astype(np.int16)
+    flat = a.reshape(-1, 3)
+    vals, counts = np.unique(flat[::7], axis=0, return_counts=True)
+    bg = vals[counts.argmax()]
+    diff = np.abs(a - bg).sum(axis=2) > 24
+    return a, float(diff.mean())
+
+
+def extract(text, kind):
+    if text is None:
+        return None
+    blocks = re.findall(r"```(\w*)\s*\n(.*?)```", text, re.S)
+    for lang, body in reversed(blocks):
+        if kind == "svg" and ("<svg" in body):
+            return body[body.index("<svg"):].strip()
+        if kind == "html" and ("<html" in body.lower() or "<body" in body.lower() or "<canvas" in body or "<script" in body):
+            return body.strip()
+    if kind == "svg":
+        m = re.search(r"<svg.*</svg>", text, re.S)
+        return m.group(0) if m else None
+    m = re.search(r"<!doctype html.*</html>|<html.*</html>", text, re.S | re.I)
+    return m.group(0) if m else None
+
+
+class Renderer:
+    def __init__(self, channel="chrome"):
+        self.pw = sync_playwright().start()
+        self.browser = self.pw.chromium.launch(channel=channel, headless=True,
+                                               args=["--use-angle=swiftshader", "--enable-unsafe-swiftshader"])
+
+    def close(self):
+        self.browser.close()
+        self.pw.stop()
+
+    def render(self, kind, source, png_path, wait_ms=1500):
+        rep = {"errors": [], "failed": [], "parse_error": None}
+        if kind == "svg":
+            try:
+                xml.dom.minidom.parseString(source.encode("utf-8"))
+            except Exception as e:
+                rep["parse_error"] = str(e)[:200]
+            html = SVG_PAGE.replace("__SVG__", source)
+        else:
+            html = source
+        path = os.path.join(TMP, f"page_{os.getpid()}_{abs(hash(png_path)) % 10**8}.html")
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(html)
+        page = self.browser.new_page(viewport={"width": W, "height": H})
+        page.on("pageerror", lambda e: rep["errors"].append(f"Uncaught {e.name}: {e.message}"[:300]))
+        page.on("console", lambda m: rep["errors"].append(("console.error: " + m.text)[:300]) if m.type == "error" else None)
+        page.on("requestfailed", lambda r: rep["failed"].append(r.url[:200]))
+        try:
+            page.goto("file:///" + path.replace("\\", "/"), wait_until="load", timeout=20000)
+            page.wait_for_timeout(wait_ms)
+            png1 = page.screenshot()
+            page.wait_for_timeout(700)
+            png2 = page.screenshot()
+            rep["overflow"] = bool(page.evaluate("document.documentElement.scrollWidth > window.innerWidth + 2"))
+            rep["page_height"] = int(page.evaluate("document.documentElement.scrollHeight"))
+        except Exception as e:
+            rep["errors"].append(f"page failed to load: {str(e)[:200]}")
+            png1 = png2 = page.screenshot() if not page.is_closed() else None
+            rep["overflow"], rep["page_height"] = False, 0
+        finally:
+            page.close()
+            try:
+                os.remove(path)
+            except OSError:
+                pass
+        if png1 is None:
+            rep.update(blank=True, coverage=0.0, animated=False)
+            return rep
+        with open(png_path, "wb") as f:
+            f.write(png2)
+        a1, cov = _stats(png1)
+        a2, _ = _stats(png2)
+        rep["coverage"] = round(cov, 4)
+        rep["blank"] = cov < 0.003
+        rep["animated"] = bool((np.abs(a1 - a2).sum(axis=2) > 24).mean() > 0.0005)
+        return rep
+
+
+if __name__ == "__main__":
+    import json
+    import sys
+    r = Renderer()
+    kind, src, out = sys.argv[1], open(sys.argv[2], encoding="utf-8").read(), sys.argv[3]
+    print(json.dumps(r.render(kind, src, out), ensure_ascii=False))
+    r.close()

research/scripts/repair.py

@@ -0,0 +1,162 @@
+"""Test-and-repair module for code answers (applied after the model has answered).
+
+The final program is run on the public tests of the statement only (examples / doctests / quoted asserts, see
+public_tests.py; the hidden grading tests are never used). If it fails, the failure is shown to the model in a new
+user turn (expected vs actual output or the error) and the model thinks again and returns a corrected program.
+Up to --rounds repairs; the kept answer is the latest attempt with the most public tests passed.
+
+Records whose first answer already passes (or that have no public tests / no answer) are copied unchanged, so the
+comparison with the source run is exactly paired.
+
+usage: repair.py --src results/dev_dm_code4096/dev.jsonl --items data/train/dev.jsonl --out results/rep1/dev.jsonl
+"""
+import argparse
+import json
+import os
+import sys
+import threading
+import time
+import urllib.request
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from code_rescue import PREFILL, close_block, runnable  # noqa: E402
+from grade import extract_code  # noqa: E402
+from public_tests import run_public  # noqa: E402
+from kit_inline import inline_cplib  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+CODE_TASKS = {"code_io", "humaneval", "mbpp", "code_fn", "lcb"}
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--src", required=True)
+ap.add_argument("--items", nargs="+", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--rounds", type=int, default=2)
+ap.add_argument("--think", type=int, default=6144, help="thinking budget of one repair turn")
+ap.add_argument("--reserve", type=int, default=2048, help="answer budget after the thinking of a repair turn")
+ap.add_argument("--workers", type=int, default=32)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+ap.add_argument("--port", type=int, default=8080)
+ap.add_argument("--inline-cplib", action="store_true",
+                help="answers may use `from cplib import *` (kits/cplib.py): test and store them with the library pasted in")
+ap.add_argument("--code-rescue", action="store_true",
+                help="if the answer at the thinking budget has no program that compiles, answer again with ```python opened")
+args = ap.parse_args()
+
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+PRE_IDS = tok(PREFILL, add_special_tokens=False)["input_ids"]
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
+            "samplers": ["top_k", "top_p", "temperature"]}
+FEEDBACK = ("I ran your program on {what} from the problem statement and it does not give the expected result.\n\n"
+            "{fail}\n\nFind the mistake (re-read the problem statement if your understanding may be wrong) and give "
+            "the complete corrected program in a single ```python code block.")
+
+
+def post(body):
+    for attempt in range(5):
+        try:
+            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 4:
+                raise
+            time.sleep(5)
+
+
+def generate(ids, seed):
+    """One thinking turn with an answer-at-budget fallback. Returns (answer text, generated token count)."""
+    r = post({"prompt": ids, "n_predict": args.think + args.reserve, "seed": seed, "cache_prompt": False,
+              "return_tokens": True, **SAMPLING})
+    gen = [t for t in r["tokens"] if t != EOS]
+    if THINK_END in gen:
+        return r["content"].split("</think>", 1)[1], r["tokens_predicted"]
+    lim = min(args.think, len(gen))
+    cut = next((j for j in range(lim, 0, -1) if gen[j - 1] in ENDS), lim)
+    r2 = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.reserve, "seed": seed + 17,
+               "cache_prompt": False, **SAMPLING})
+    text, n = r2["content"], cut + r2["tokens_predicted"]
+    if args.code_rescue and not runnable(text):
+        r3 = post({"prompt": ids + gen[:cut] + [THINK_END] + PRE_IDS, "n_predict": args.reserve, "seed": seed + 29,
+                   "cache_prompt": False, **SAMPLING})
+        text, n = close_block(r3["content"]), n + len(PRE_IDS) + r3["tokens_predicted"]
+    return text, n
+
+
+def task_of(item):
+    if item["task"] == "lcb":
+        return "lcb"
+    return "humaneval" if item["task"] == "humaneval" else "mbpp" if item["task"] in ("mbpp", "code_fn") else "code_io"
+
+
+def process(rec):
+    try:
+        return _process(rec)
+    except Exception as e:  # a failed repair leaves the record as it was
+        return {**rec, "repair_error": repr(e)[:200]}
+
+
+def prep(text):
+    return inline_cplib(text) if args.inline_cplib else text
+
+
+def _process(rec):
+    item = items[rec["id"]]
+    if item["task"] not in CODE_TASKS or "</think>" not in rec["completion"]:
+        return rec
+    t = task_of(item)
+    head, final = rec["completion"].split("</think>", 1)
+    code = extract_code(prep(final))
+    if code is None:  # answered without a program: ask for it like any other failure
+        ok, total, _ = run_public(t, "raise SystemExit(1)", item["prompt"], item)
+        fail = "Your answer does not contain the complete program in a ```python code block."
+    else:
+        ok, total, fail = run_public(t, code, item["prompt"], item)
+    if total == 0 or ok == total:
+        return {**rec, "completion": head + "</think>" + prep(final), "public": [ok, total]}
+    attempts, used, msgs = [(ok, final)], 0, [{"role": "user", "content": item["prompt"]}]
+    seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec["sample"] + 5000
+    what = "the examples" if t in ("code_io", "lcb") else "the tests"
+    for rnd in range(args.rounds):
+        msgs = msgs + [{"role": "assistant", "content": final},
+                       {"role": "user", "content": FEEDBACK.format(what=what, fail=fail)}]
+        ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_dict=False)
+        final, n = generate(ids, seed + rnd)
+        used += n
+        code = extract_code(prep(final))
+        if code is None:
+            attempts.append((-1, final))
+            continue
+        ok, total, fail = run_public(t, code, item["prompt"], item)
+        attempts.append((ok, final))
+        if ok == total:
+            break
+    best = max(range(len(attempts)), key=lambda k: (attempts[k][0], -k))  # a repair must pass strictly more
+    return {**rec, "completion": "<think>(repaired)</think>" + prep(attempts[best][1]), "n_tokens": rec["n_tokens"] + used,
+            "repair_rounds": len(attempts) - 1, "repair_kept": best, "public": [attempts[best][0], total],
+            "public_first": attempts[0][0]}
+
+
+items = {}
+for path in args.items:
+    with open(path, encoding="utf-8") as f:
+        for l in f:
+            r = json.loads(l)
+            items[r["id"]] = r
+with open(args.src, encoding="utf-8") as f:
+    recs = [json.loads(l) for l in f]
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+t0 = time.time()
+with ThreadPoolExecutor(args.workers) as ex:
+    out = list(ex.map(process, recs))
+with open(args.out, "w", encoding="utf-8") as f:
+    for r in out:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+rep = [r for r in out if "repair_rounds" in r]
+print(f"{len(recs)} records, {len(rep)} repaired attempts, "
+      f"{sum(r['public'][0] == r['public'][1] for r in rep)} now pass public tests, {(time.time() - t0) / 60:.1f} min")

research/scripts/rl_iter.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+# usage: rl_iter.sh <iter_name> <init_hf_dir> <init_q5_gguf> <n_prompts> [lr] [epochs]
+# rollouts on the Q5 quantized policy (llama.cpp), update in bf16 with LoRA (merged), re-quantize for the next round
+set -e
+cd /c/Users/user/lfm3proj
+IT=$1; INIT=$2; GG=$3; N=$4; LR=${5:-2e-5}; EP=${6:-1}
+PY=.venv/Scripts/python.exe
+D=results/rl/$IT; P=data/rl/$IT; mkdir -p $D $P
+[ -f $P/prompts.jsonl ] || $PY scripts/sample_prompts.py $N 0.25 $P/prompts.jsonl $RANDOM
+taskkill //F //IM llama-server.exe >/dev/null 2>&1 || true; sleep 2
+(/c/Users/user/llamabin/llama-server.exe -m $GG -ngl 99 -np 32 -c 327680 -bs -fa on -ctxcp 0 -cram 0 --port 8080 --no-webui > $D/server.log 2>&1 &)
+for i in $(seq 1 100); do curl -s -m 2 localhost:8080/health | grep -q '"ok"' && break; sleep 2; done
+$PY scripts/gen.py --inp $P/prompts.jsonl --out $D/rollouts.jsonl --n 8 --temperature 1.0 --top-p 1.0 --top-k 50 \
+    --max-tokens 8192 --save-tokens --workers 32 2>&1 | grep -viE "warn|instantiate" | tail -1
+$PY scripts/grade.py rl/$IT "$P/prompts.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+taskkill //F //IM llama-server.exe >/dev/null 2>&1 || true; sleep 3
+$PY scripts/rl_train.py --init $INIT --rollouts $D/rollouts.jsonl --graded $D/graded.jsonl --out models/$IT --lr $LR \
+    --epochs $EP --mb 32 --max-len 9300 2>&1 | grep -viE "warn|instantiate|Loading weights" | tee $D/train.log
+$PY tools/llama-src/convert_hf_to_gguf.py models/$IT --outtype bf16 --outfile gguf/$IT-bf16.gguf 2>&1 | grep -E "rror" || true
+/c/Users/user/llamabin/llama-quantize.exe gguf/$IT-bf16.gguf gguf/$IT-q5km.gguf Q5_K_M 8 >/dev/null 2>&1
+rm -f gguf/$IT-bf16.gguf
+echo ITER_DONE $IT

research/scripts/rl_iter_bc.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# RL iteration with adaptive sampling + commitment bootstrap.
+# usage: rl_iter_bc.sh <iter_name> <init_hf_dir> <init_q5_gguf> <n_prompts> [lr]
+#  stage 1: 4 samples/prompt; prompts solved 4/4 are dropped (no signal); stage 2: the rest get 8 samples
+#  bc_probe: truncated math rollouts -> forced-commit samples (see bc_probe.py); then clipped PG on LoRA, merge, Q5
+set -e
+cd /c/Users/user/lfm3proj
+IT=$1; INIT=$2; GG=$3; N=$4; LR=${5:-3e-5}
+PY=.venv/Scripts/python.exe
+D=results/rl/$IT; P=data/rl/$IT; mkdir -p $D $P
+[ -f $P/prompts.jsonl ] || $PY scripts/sample_prompts.py $N 0.2 $P/prompts.jsonl $RANDOM
+taskkill //F //IM llama-server.exe >/dev/null 2>&1 || true; sleep 2
+(/c/Users/user/llamabin/llama-server.exe -m $GG -ngl 99 -np 32 -c 327680 -bs -fa on -ctxcp 0 -cram 0 --port 8080 --no-webui > $D/server.log 2>&1 &)
+for i in $(seq 1 100); do curl -s -m 2 localhost:8080/health | grep -q '"ok"' && break; sleep 2; done
+G="--temperature 1.0 --top-p 1.0 --top-k 50 --max-tokens 8192 --save-tokens --workers 32"
+$PY scripts/gen.py --inp $P/prompts.jsonl --out $D/rollouts.jsonl --n 4 $G 2>&1 | grep -viE "warn|instantiate" | tail -1
+$PY scripts/grade.py rl/$IT "$P/prompts.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+$PY - <<PYEOF
+import json, collections
+ok = collections.defaultdict(list)
+for l in open("$D/graded.jsonl", encoding="utf-8"):
+    g = json.loads(l); ok[g["id"]].append(g["correct"])
+keep = [l for l in open("$P/prompts.jsonl", encoding="utf-8") if not all(ok[json.loads(l)["id"]])]
+open("$P/prompts_s2.jsonl", "w", encoding="utf-8").writelines(keep)
+print("stage 2 prompts:", len(keep), "of", len(ok))
+PYEOF
+$PY scripts/gen.py --inp $P/prompts_s2.jsonl --out $D/rollouts.jsonl --n 8 $G 2>&1 | grep -viE "warn|instantiate" | tail -1
+$PY scripts/grade.py rl/$IT "$P/prompts.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+$PY scripts/bc_probe.py --rollouts $D/rollouts.jsonl --graded $D/graded.jsonl --items $P/prompts.jsonl --out-dir $D \
+   2>&1 | grep -viE "warn|instantiate"
+mkdir -p $D/bc && mv $D/forced.jsonl $D/forced_graded.jsonl $D/bc/
+taskkill //F //IM llama-server.exe >/dev/null 2>&1 || true; sleep 3
+$PY scripts/rl_train.py --init $INIT --rollouts $D/rollouts.jsonl $D/bc/forced.jsonl \
+   --graded $D/graded.jsonl $D/bc/forced_graded.jsonl --out models/$IT --lr $LR --mb 32 --max-len 9300 2>&1 \
+   | grep -viE "warn|instantiate|Loading weights|fast path|Redirect|Writing" | tee $D/train.log
+$PY tools/llama-src/convert_hf_to_gguf.py models/$IT --outtype bf16 --outfile gguf/$IT-bf16.gguf 2>&1 | grep -E "rror" || true
+/c/Users/user/llamabin/llama-quantize.exe gguf/$IT-bf16.gguf gguf/$IT-q5km.gguf Q5_K_M 8 >/dev/null 2>&1
+rm -f gguf/$IT-bf16.gguf
+echo ITER_DONE $IT

research/scripts/rl_train.py

@@ -0,0 +1,217 @@
+"""One policy-improvement iteration on self-generated, automatically graded rollouts.
+
+Rollouts come from llama.cpp running the current policy (exact token ids are stored), rewards from grade.py.
+For every prompt the G samples form a group; advantage = reward - group mean (groups where all samples agree carry
+no signal and are skipped). The update is a clipped policy-gradient objective (PPO/GRPO-style, asymmetric
+"clip-higher", token-level normalisation) on a fresh LoRA adapter, which is merged into the weights at the end,
+so the next iteration samples from the improved policy.
+
+Memory: B=1 sequences, gradient checkpointing, and log-probs computed in vocabulary chunks under checkpoint so
+the 128k-vocab logits of a 12k-token sequence are never materialised.
+"""
+import argparse
+import json
+import math
+import os
+import random
+import shutil
+import time
+from collections import defaultdict
+
+import torch
+from peft import LoraConfig, get_peft_model
+from torch.utils.checkpoint import checkpoint
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import transformers.integrations.sdpa_attention as _sdpa
+
+# Windows torch has no flash kernel: masked/GQA SDPA would fall back to the O(L^2)-memory math kernel.
+# Expanding K/V instead keeps it on the memory-efficient kernel.
+_sdpa.use_gqa_in_sdpa = lambda *a, **k: False
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--init", required=True, help="HF dir of the policy that generated the rollouts")
+ap.add_argument("--rollouts", nargs="+", required=True)
+ap.add_argument("--graded", nargs="+", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--lr", type=float, default=2e-5)
+ap.add_argument("--rank", type=int, default=32)
+ap.add_argument("--epochs", type=int, default=1)
+ap.add_argument("--mb", type=int, default=64, help="sequences per optimizer step")
+ap.add_argument("--max-len", type=int, default=12800)
+ap.add_argument("--clip-lo", type=float, default=0.2)
+ap.add_argument("--clip-hi", type=float, default=0.28)
+ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--items", nargs="+", default=["data/train/pool_math.jsonl", "data/train/pool_code.jsonl"])
+ap.add_argument("--max-seqs", type=int, default=0, help="debug: use only this many sequences")
+ap.add_argument("--no-save", action="store_true")
+ap.add_argument("--mode", choices=["lora", "bias"], default="lora",
+                help="bias: train only one residual-stream bias vector per layer (layers 1..L-1), exported as a "
+                     "llama.cpp control vector; weights stay frozen")
+args = ap.parse_args()
+random.seed(args.seed)
+torch.manual_seed(args.seed)
+EOS = 124900  # <|im_end|>
+
+tok = AutoTokenizer.from_pretrained(args.init)
+reward = {}
+for path in args.graded:
+    for l in open(path, encoding="utf-8"):
+        g = json.loads(l)
+        reward[(g["id"], g["sample"])] = float(g["correct"])
+prompts = {}
+for path in args.rollouts:
+    for l in open(path, encoding="utf-8"):
+        r = json.loads(l)
+        if (r["id"], r["sample"]) in reward:
+            prompts.setdefault(r["id"], []).append(r)
+
+items = {}
+for path in args.items:
+    for l in open(path, encoding="utf-8"):
+        it = json.loads(l)
+        items[it["id"]] = it
+
+seqs, n_groups, n_flat = [], 0, 0
+for pid, group in prompts.items():
+    rs = [reward[(r["id"], r["sample"])] for r in group]
+    mean = sum(rs) / len(rs)
+    if max(rs) == min(rs):
+        n_flat += 1
+        continue
+    n_groups += 1
+    p_ids = tok.apply_chat_template([{"role": "user", "content": items[pid]["prompt"]}], tokenize=True,
+                                    add_generation_prompt=True, return_dict=False)
+    for r, rw in zip(group, rs):
+        resp = list(r["tokens"])
+        if r.get("stop_type") == "eos":
+            resp.append(EOS)
+        ids = (p_ids + resp)[: args.max_len]
+        # forced-commit samples share their prefix with an on-policy rollout: learn only from the decision onwards
+        seqs.append({"ids": ids, "start": len(p_ids), "adv": rw - mean, "from": r.get("loss_from", 0)})
+if args.max_seqs:
+    seqs = seqs[: args.max_seqs]
+print(f"{len(prompts)} prompts, {n_groups} informative groups, {n_flat} flat groups skipped, {len(seqs)} sequences",
+      flush=True)
+
+model = AutoModelForCausalLM.from_pretrained(args.init, dtype=torch.bfloat16, attn_implementation="sdpa", trust_remote_code=True).cuda()
+model.config.use_cache = False
+model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+model.enable_input_require_grads()
+if args.mode == "lora":
+    lora = LoraConfig(r=args.rank, lora_alpha=2 * args.rank, lora_dropout=0.0, bias="none",
+                      target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "in_proj", "out_proj",
+                                      "w1", "w2", "w3"])
+    model = get_peft_model(model, lora)
+    model.print_trainable_parameters()
+    base = model.get_base_model()
+else:
+    for p_ in model.parameters():
+        p_.requires_grad_(False)
+    n_l, hid = model.config.num_hidden_layers, model.config.hidden_size
+    resid_bias = torch.nn.Parameter(torch.zeros(n_l, hid, device="cuda", dtype=torch.float32))
+    model.register_parameter("resid_bias", resid_bias)
+    for k, layer in enumerate(model.model.layers):
+        if k == 0:
+            continue  # llama.cpp control vectors cannot address layer 0
+
+        def hook(m, i, o, k=k):
+            if isinstance(o, tuple):
+                return (o[0] + resid_bias[k].to(o[0].dtype),) + o[1:]
+            return o + resid_bias[k].to(o.dtype)
+        layer.register_forward_hook(hook)
+    print("trainable residual bias params:", (n_l - 1) * hid)
+    base = model
+
+
+def _chunk_logp(h, w, t):
+    return torch.log_softmax((h @ w.t()).float(), -1).gather(-1, t[:, None]).squeeze(-1)
+
+
+def token_logps(ids, start):
+    x = torch.tensor([ids], device="cuda")
+    h = base.model(input_ids=x, use_cache=False).last_hidden_state[0, start - 1: -1]
+    t = x[0, start:]
+    w = base.lm_head.weight
+    out = []
+    for i in range(0, h.shape[0], 512):
+        out.append(checkpoint(_chunk_logp, h[i: i + 512], w, t[i: i + 512], use_reentrant=False))
+    return torch.cat(out)
+
+
+# log-probs of the policy that generated the data (adapter / bias are zero-initialised, so this is the init policy)
+t0 = time.time()
+model.eval()
+with torch.no_grad():
+    for k, s in enumerate(seqs):
+        s["old"] = token_logps(s["ids"], s["start"]).cpu()
+        if k % 200 == 0:
+            print(f"old logp {k}/{len(seqs)}  {time.time() - t0:.0f}s", flush=True)
+model.train()
+
+opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr, betas=(0.9, 0.99),
+                        weight_decay=0.0)
+steps_total = args.epochs * math.ceil(len(seqs) / args.mb)
+sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / 3))
+step = 0
+log = []
+for ep in range(args.epochs):
+    order = list(range(len(seqs)))
+    random.shuffle(order)
+    for b in range(0, len(order), args.mb):
+        batch = [seqs[i] for i in order[b: b + args.mb]]
+        n_tok = sum(len(s["ids"]) - s["start"] - s["from"] for s in batch)
+        stats = defaultdict(float)
+        for s in batch:
+            lp = token_logps(s["ids"], s["start"])
+            old = s["old"].to(lp.device)
+            ratio = torch.exp(lp - old)
+            adv = s["adv"]
+            unclipped = ratio * adv
+            clipped = torch.clamp(ratio, 1 - args.clip_lo, 1 + args.clip_hi) * adv
+            loss = -torch.minimum(unclipped, clipped)[s["from"]:].sum() / n_tok
+            loss.backward()
+            if args.max_seqs:
+                print(f"  seq len {len(s['ids'])}  t {time.time() - t0:.0f}s  peak {torch.cuda.max_memory_allocated() / 1e9:.1f}GB",
+                      flush=True)
+            with torch.no_grad():
+                stats["loss"] += loss.item()
+                stats["kl"] += (old - lp)[s["from"]:].sum().item() / n_tok
+                stats["clipfrac"] += ((ratio - 1).abs() > args.clip_lo)[s["from"]:].float().sum().item() / n_tok
+        gn = torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
+        opt.step()
+        sched.step()
+        opt.zero_grad(set_to_none=True)
+        step += 1
+        rec = {"step": step, "of": steps_total, "gnorm": gn.item(), **stats, "t": round(time.time() - t0)}
+        log.append(rec)
+        print(json.dumps(rec), flush=True)
+
+print("peak mem GB", torch.cuda.max_memory_allocated() / 1e9)
+if args.no_save:
+    raise SystemExit
+if args.mode == "bias":
+    import gguf
+    os.makedirs(args.out, exist_ok=True)
+    w = gguf.GGUFWriter(os.path.join(args.out, "resid_bias.gguf"), "controlvector")
+    w.add_string("controlvector.model_hint", "lfm3")
+    w.add_uint32("controlvector.layer_count", n_l - 1)
+    for k in range(1, n_l):
+        w.add_tensor(f"direction.{k}", resid_bias[k].detach().float().cpu().numpy())
+    w.write_header_to_file(); w.write_kv_data_to_file(); w.write_tensors_to_file(); w.close()
+    torch.save(resid_bias.detach().cpu(), os.path.join(args.out, "resid_bias.pt"))
+    json.dump({"args": vars(args), "log": log, "n_groups": n_groups, "n_seqs": len(seqs)},
+              open(os.path.join(args.out, "train_log.json"), "w"), indent=1)
+    print("saved", args.out, "bias norms", [round(x, 3) for x in resid_bias.norm(dim=-1).tolist()])
+    raise SystemExit
+model = model.merge_and_unload()
+os.makedirs(args.out, exist_ok=True)
+model.save_pretrained(args.out, safe_serialization=True)
+tok.save_pretrained(args.out)
+for f in ["configuration_lfm3.py", "modeling_lfm3.py", "lfm3_conv.py", "chat_template.jinja", "generation_config.json",
+          "config.json"]:
+    src = os.path.join(args.init, f)
+    if os.path.exists(src):
+        shutil.copy(src, os.path.join(args.out, f))
+json.dump({"args": vars(args), "log": log, "n_groups": n_groups, "n_seqs": len(seqs)},
+          open(os.path.join(args.out, "train_log.json"), "w"), indent=1)
+print("saved", args.out)

research/scripts/run_dev.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+# usage: run_dev.sh <tag> [general]   (server for the model must already be running)
+set -e
+cd /c/Users/user/lfm3proj
+T=$1; PY=.venv/Scripts/python.exe
+$PY scripts/gen.py --out results/$T/dev.jsonl --workers 32 --n 2 --save-tokens --inp data/train/dev.jsonl 2>&1 | grep -viE "warn|instantiate" | tail -1
+if [ "$2" == "general" ]; then
+  $PY scripts/gen.py --out results/$T/devgen.jsonl --workers 32 --inp data/train/dev_general.jsonl 2>&1 | grep -viE "warn|instantiate" | tail -1
+fi
+$PY scripts/grade.py $T "data/train/dev*.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+echo DEV_DONE $T

research/scripts/run_eval.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# usage: run_eval.sh <tag>   (server for the model must already be running on :8080)
+set -e
+cd /c/Users/user/lfm3proj
+T=$1
+PY=.venv/Scripts/python.exe
+$PY scripts/gen.py --out results/$T/main.jsonl --workers 32 --inp data/eval/gsm8k.jsonl data/eval/math500.jsonl data/eval/humaneval.jsonl data/eval/mbpp.jsonl data/eval/ifeval.jsonl data/eval/mmlupro.jsonl data/eval/bbh.jsonl
+$PY scripts/gen.py --out results/$T/aime.jsonl --workers 32 --n 4 --inp data/eval/aime.jsonl
+echo EVAL_DONE

research/scripts/run_mini.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+# usage: run_mini.sh <tag> <gguf> [extra server args]  -- ~8 min screening: 140 dev items, n=1, max 6144 tokens
+set -e
+cd /c/Users/user/lfm3proj
+T=$1; PY=.venv/Scripts/python.exe
+taskkill //F //IM llama-server.exe >/dev/null 2>&1 || true; sleep 2
+(/c/Users/user/llamabin/llama-server.exe -m "$2" "${@:3}" -ngl 99 -np 32 -c 212992 -bs -fa on -ctxcp 0 -cram 0 --port 8080 --no-webui > results/server_$T.log 2>&1 &)
+for i in $(seq 1 100); do curl -s -m 2 localhost:8080/health | grep -q '"ok"' && break; sleep 2; done
+$PY scripts/gen.py --out results/$T/mini.jsonl --workers 32 --max-tokens 6144 --inp data/quick/mini.jsonl 2>&1 | grep -viE "warn|instantiate" | tail -1
+$PY scripts/grade.py $T "data/quick/mini.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+echo MINI_DONE $T

research/scripts/run_mini2.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# usage: [ENV=..] run_mini2.sh <tag> <gguf>   -- same as run_mini.sh, served by the custom build (LLAMA_BIN), env passes through
+cd /c/Users/user/lfm3proj
+T=$1; PY=.venv/Scripts/python.exe
+LLAMA_BIN=${LLAMA_BIN:-/c/Users/user/llama-build/bin} bash scripts/serve.sh "$2" results/server_$T.log -ngl 99 -np 32 -c 212992 -bs -fa on -ctxcp 0 -cram 0
+$PY scripts/gen.py --out results/$T/mini.jsonl --workers 32 --max-tokens 6144 --inp data/quick/mini.jsonl 2>&1 | grep -viE "warn|instantiate" | tail -1
+$PY scripts/grade.py $T "data/quick/mini.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+echo MINI_DONE $T

research/scripts/run_quick.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# usage: run_quick.sh <tag> <gguf> [extra server args]  -- quick screening run (274 dev items, n=1, same sampling as eval)
+set -e
+cd /c/Users/user/lfm3proj
+T=$1; PY=.venv/Scripts/python.exe
+bash scripts/start_server.sh "${@:2}"
+$PY scripts/gen.py --out results/$T/quick.jsonl --workers 32 --save-tokens --inp data/quick/quick.jsonl 2>&1 | grep -viE "warn|instantiate" | tail -1
+$PY scripts/grade.py $T "data/quick/quick.jsonl" 2>>results/grade_err.log | grep -vE "^\s*$"
+echo QUICK_DONE $T

research/scripts/sample_prompts.py

@@ -0,0 +1,17 @@
+"""Draws N unused training prompts (math/code mix) for one RL iteration; remembers what was used."""
+import json, os, random, sys
+n, frac_code, out, seed = int(sys.argv[1]), float(sys.argv[2]), sys.argv[3], int(sys.argv[4])
+used_path = "data/train/used_ids.txt"
+used = set(open(used_path).read().split()) if os.path.exists(used_path) else set()
+rng = random.Random(seed)
+pick = []
+for path, k in [("data/train/pool_math.jsonl", n - int(n * frac_code)), ("data/train/pool_code.jsonl", int(n * frac_code))]:
+    rows = [json.loads(l) for l in open(path, encoding="utf-8")]
+    rows = [r for r in rows if r["id"] not in used and r.get("difficulty") not in ("HARD", "VERY_HARD", "MEDIUM_HARD")]
+    pick += rng.sample(rows, k)
+with open(out, "w", encoding="utf-8") as f:
+    for r in pick:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+with open(used_path, "a") as f:
+    f.write("\n".join(r["id"] for r in pick) + "\n")
+print(len(pick))

research/scripts/sandbox.py

@@ -0,0 +1,42 @@
+"""subprocess.run replacement for running untrusted model programs on Windows.
+
+Two Windows traps make subprocess.run(..., timeout=) hang forever on a stuck program:
+  * in a venv, python.exe is a launcher that starts the real interpreter as a child; the timeout kills only the
+    launcher and the child keeps the pipes open;
+  * Popen.communicate writes `input` synchronously before it starts waiting, so a program that stops reading a large
+    input blocks the write and the timeout is never checked.
+run() feeds stdin from a temporary file and kills the whole process tree on timeout, then raises
+subprocess.TimeoutExpired like subprocess.run."""
+import os
+import subprocess
+import tempfile
+
+
+def run(cmd, input=None, timeout=None, capture_output=True, text=False, encoding=None, errors=None, **kw):
+    kw.setdefault("creationflags", getattr(subprocess, "CREATE_NO_WINDOW", 0))
+    fin, path = subprocess.DEVNULL, None
+    if input is not None:
+        fd, path = tempfile.mkstemp(suffix=".in")
+        with os.fdopen(fd, "wb") as f:
+            f.write(input if isinstance(input, bytes) else input.encode(encoding or "utf-8"))
+        fin = open(path, "rb")
+    try:
+        p = subprocess.Popen(cmd, stdin=fin, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=text,
+                             encoding=encoding, errors=errors, **kw)
+        try:
+            out, err = p.communicate(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            subprocess.run(["taskkill", "/F", "/T", "/PID", str(p.pid)], capture_output=True)
+            try:
+                p.communicate(timeout=10)
+            except subprocess.TimeoutExpired:
+                pass
+            raise subprocess.TimeoutExpired(cmd, timeout)
+        return subprocess.CompletedProcess(cmd, p.returncode, out, err)
+    finally:
+        if path:
+            fin.close()
+            try:
+                os.remove(path)
+            except OSError:
+                pass

research/scripts/serve.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+# usage: serve.sh <gguf (path relative to project)> <log> [server args...]
+# llama-server must start with an ASCII working directory (it scans the cwd for backends; the project path is not ASCII)
+# A start right after killing the previous server sometimes fails to load the CPU backend, so it retries.
+P=/c/Users/user/lfm3proj; BIN=${LLAMA_BIN:-/c/Users/user/llamabin}
+for attempt in 1 2 3; do
+  taskkill //F //IM llama-server.exe >/dev/null 2>&1; sleep $((3 * attempt))
+  (cd $BIN && ./llama-server.exe -m "$P/$1" "${@:3}" --port 8080 --no-webui > "$P/$2" 2>&1 &)
+  for i in $(seq 1 120); do curl -s -m 2 localhost:8080/health | grep -q '"ok"' && break; grep -q "exiting due" "$P/$2" 2>/dev/null && break; sleep 2; done
+  curl -s -m 2 localhost:8080/health | grep -q '"ok"' && break
+done
+curl -s localhost:8080/health; echo

research/scripts/sft_train.py

@@ -0,0 +1,115 @@
+"""Supervised fine-tuning (LoRA, merged at the end) on token-level examples from build_sft.py.
+
+Loss = mean NLL over completion tokens. Same memory tricks as rl_train.py: B=1 sequences, gradient checkpointing,
+log-probs in vocabulary chunks under checkpoint, expanded K/V for the memory-efficient SDPA kernel on Windows.
+usage: sft_train.py --data data/teach/sft_pilot.jsonl --init models/LFM3-1.7B --out models/teach1
+"""
+import argparse
+import json
+import math
+import os
+import random
+import shutil
+import time
+
+import torch
+from peft import LoraConfig, get_peft_model
+from torch.utils.checkpoint import checkpoint
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import transformers.integrations.sdpa_attention as _sdpa
+
+_sdpa.use_gqa_in_sdpa = lambda *a, **k: False
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--data", nargs="+", required=True)
+ap.add_argument("--init", default="models/LFM3-1.7B")
+ap.add_argument("--out", required=True)
+ap.add_argument("--lr", type=float, default=1e-4)
+ap.add_argument("--rank", type=int, default=64)
+ap.add_argument("--epochs", type=int, default=2)
+ap.add_argument("--mb", type=int, default=16, help="sequences per optimizer step")
+ap.add_argument("--max-len", type=int, default=7000)
+ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--max-seqs", type=int, default=0)
+args = ap.parse_args()
+random.seed(args.seed)
+torch.manual_seed(args.seed)
+
+tok = AutoTokenizer.from_pretrained(args.init)
+seqs = []
+for p in args.data:
+    for l in open(p, encoding="utf-8"):
+        r = json.loads(l)
+        ids = (r["prompt_ids"] + r["completion_ids"])[: args.max_len]
+        seqs.append({"ids": ids, "start": len(r["prompt_ids"])})
+if args.max_seqs:
+    seqs = seqs[: args.max_seqs]
+n_tok = sum(len(s["ids"]) - s["start"] for s in seqs)
+print(f"{len(seqs)} sequences, {n_tok} completion tokens", flush=True)
+
+model = AutoModelForCausalLM.from_pretrained(args.init, dtype=torch.bfloat16, attn_implementation="sdpa",
+                                             trust_remote_code=True).cuda()
+model.config.use_cache = False
+model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+model.enable_input_require_grads()
+lora = LoraConfig(r=args.rank, lora_alpha=2 * args.rank, lora_dropout=0.0, bias="none",
+                  target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "in_proj", "out_proj",
+                                  "w1", "w2", "w3"])
+model = get_peft_model(model, lora)
+model.print_trainable_parameters()
+base = model.get_base_model()
+
+
+def _chunk_nll(h, w, t):
+    return -torch.log_softmax((h @ w.t()).float(), -1).gather(-1, t[:, None]).squeeze(-1)
+
+
+def nll(ids, start):
+    x = torch.tensor([ids], device="cuda")
+    h = base.model(input_ids=x, use_cache=False).last_hidden_state[0, start - 1: -1]
+    t = x[0, start:]
+    w = base.lm_head.weight
+    return torch.cat([checkpoint(_chunk_nll, h[i: i + 512], w, t[i: i + 512], use_reentrant=False)
+                      for i in range(0, h.shape[0], 512)])
+
+
+params = [p for p in model.parameters() if p.requires_grad]
+opt = torch.optim.AdamW(params, lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0)
+steps_total = args.epochs * math.ceil(len(seqs) / args.mb)
+sched = torch.optim.lr_scheduler.LambdaLR(
+    opt, lambda s: min(1.0, (s + 1) / 5) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / steps_total))))
+t0, step, log = time.time(), 0, []
+model.train()
+for ep in range(args.epochs):
+    order = list(range(len(seqs)))
+    random.shuffle(order)
+    for b in range(0, len(order), args.mb):
+        batch = [seqs[i] for i in order[b: b + args.mb]]
+        nt = sum(len(s["ids"]) - s["start"] for s in batch)
+        tot = 0.0
+        for s in batch:
+            loss = nll(s["ids"], s["start"]).sum() / nt
+            loss.backward()
+            tot += loss.item()
+        gn = torch.nn.utils.clip_grad_norm_(params, 1.0)
+        opt.step()
+        sched.step()
+        opt.zero_grad(set_to_none=True)
+        step += 1
+        rec = {"step": step, "of": steps_total, "ep": ep, "nll": round(tot, 4), "gnorm": round(gn.item(), 3),
+               "lr": sched.get_last_lr()[0], "t": round(time.time() - t0)}
+        log.append(rec)
+        print(json.dumps(rec), flush=True)
+print("peak mem GB", torch.cuda.max_memory_allocated() / 1e9)
+model = model.merge_and_unload()
+os.makedirs(args.out, exist_ok=True)
+model.save_pretrained(args.out, safe_serialization=True)
+tok.save_pretrained(args.out)
+for f in ["configuration_lfm3.py", "modeling_lfm3.py", "lfm3_conv.py", "chat_template.jinja", "generation_config.json",
+          "config.json"]:
+    src = os.path.join(args.init, f)
+    if os.path.exists(src):
+        shutil.copy(src, os.path.join(args.out, f))
+json.dump({"args": vars(args), "log": log, "n_seqs": len(seqs)}, open(os.path.join(args.out, "train_log.json"), "w"),
+          indent=1)
+print("saved", args.out)

research/scripts/shadow_split.py

@@ -0,0 +1,16 @@
+"""Split a shadow CoC run into two gradeable result dirs on the SAME trajectories:
+results/<tag>_full/ (generation run to the end) and results/<tag>_coc/ (committed answer where consensus was reached,
+otherwise identical to full). usage: shadow_split.py results/shadow/raw.jsonl shadow"""
+import json, os, sys
+src, tag = sys.argv[1], sys.argv[2]
+recs = [json.loads(l) for l in open(src, encoding="utf-8")]
+for kind in ("full", "coc"):
+    d = f"results/{tag}_{kind}"
+    os.makedirs(d, exist_ok=True)
+    with open(f"{d}/gen.jsonl", "w", encoding="utf-8") as f:
+        for r in recs:
+            r2 = {k: v for k, v in r.items() if k not in ("coc_completion", "coc_tokens")}
+            if kind == "coc" and "coc_completion" in r:
+                r2.update(completion=r["coc_completion"], n_tokens=r["coc_n_tokens"], truncated=False)
+            f.write(json.dumps(r2, ensure_ascii=False) + "\n")
+print(len(recs), "records,", sum("coc_completion" in r for r in recs), "committed")

research/scripts/sim_coc.py

@@ -0,0 +1,50 @@
+"""Offline simulation of commit-on-consensus settings from a --shadow --record-all run (probes every 512 tokens).
+For each policy (probe every `chunk` tokens, commit after `agree` identical consecutive answers) and task group it
+reports accuracy of the full generation vs. the CoC output on the SAME trajectories, flips, and average tokens.
+usage: sim_coc.py results/rec/raw.jsonl data/coc/*.jsonl"""
+import glob, json, sys, collections
+sys.path.insert(0, "scripts")
+from grade import grade_one  # noqa: E402
+recs = [json.loads(l) for l in open(sys.argv[1], encoding="utf-8")]
+for r in recs:  # older runs kept the end-of-turn token in the text
+    r["completion"] = r["completion"].replace("<|im_end|>", "")
+items = {}
+for pat in sys.argv[2:]:
+    for f in glob.glob(pat):
+        for l in open(f, encoding="utf-8"):
+            r = json.loads(l); items[r["id"]] = r
+group = lambda r: ("math_dev" if r["id"].startswith(("math_train", "deepscaler")) and r["id"] in MINI else
+                   "math_pool" if r["task"] == "math" else r["task"])
+MINI = {json.loads(l)["id"] for l in open("data/coc/mini_math.jsonl", encoding="utf-8")}
+full_ok, probe_ok = {}, {}
+for r in recs:
+    it = items[r["id"]]
+    full_ok[r["id"]] = grade_one(r, it)["correct"]
+    probe_ok[r["id"]] = [grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"]
+                         for p in r.get("probe_log", [])]
+print(f"{'policy':14s} {'group':10s} {'n':>4s} {'commits':>7s} {'full':>6s} {'coc':>6s} {'+':>3s} {'-':>3s} {'tok_full':>8s} {'tok_coc':>8s}")
+for chunk in (512, 1024, 2048):
+    for agree in (2, 3):
+        st = collections.defaultdict(lambda: [0, 0, 0, 0, 0, 0, 0, 0])
+        for r in recs:
+            g = group(r)
+            log = r.get("probe_log", [])
+            step = chunk // 512
+            idx = [i for i in range(len(log)) if (i + 1) % step == 0]
+            commit = None
+            for j in range(agree - 1, len(idx)):
+                a = [log[idx[t]]["ans"] for t in range(j - agree + 1, j + 1)]
+                if a[0] is not None and all(x == a[0] for x in a):
+                    commit = idx[j]
+                    break
+            fo = full_ok[r["id"]]
+            co = probe_ok[r["id"]][commit] if commit is not None else fo
+            tf = r["n_tokens"]
+            tc = log[commit]["cut"] + log[commit]["n"] if commit is not None else tf
+            for key in (g, "ALL"):
+                s = st[key]
+                s[0] += 1; s[1] += commit is not None; s[2] += fo; s[3] += co
+                s[4] += (not fo) and co; s[5] += fo and not co; s[6] += tf; s[7] += tc
+        for g, s in sorted(st.items()):
+            print(f"c{chunk:<5d}a{agree:<6d} {g:10s} {s[0]:4d} {s[1]:7d} {100*s[2]/s[0]:6.1f} {100*s[3]/s[0]:6.1f} "
+                  f"{s[4]:3d} {s[5]:3d} {s[6]/s[0]:8.0f} {s[7]/s[0]:8.0f}")

research/scripts/sim_multi.py

@@ -0,0 +1,165 @@
+"""Adaptive deliberation: offline evaluation of controllers over recorded parallel reasoning streams.
+
+Input: a gen_coc2.py --shadow --record-all --n N run (N streams per problem, a forced-answer probe every ~1024
+tokens of each stream, full completions). Streams are assumed to advance at the same speed (they run in parallel
+slots). Every controller is evaluated on the same recorded trajectories; cost = total generated tokens (all streams,
+probes' own tokens excluded; probes cost prefill only).
+
+Controllers
+  single        stream 0 to the end (the baseline generation)
+  sc@N          majority vote over the N full answers (self-consistency)
+  coc1          stream 0 with temporal consensus (commit when `agree` successive probes match)
+  xs-K          N parallel streams; stop all as soon as K streams currently hold the same answer
+                (a stream holds its final answer once finished, otherwise its latest probe answer)
+  adapt-K       stream 0 alone; if it has not finished by `t0` tokens, streams 1..N-1 start; then as xs-K over the
+                streams that are running, and stream 0's natural answer is accepted if it finishes first
+usage: sim_multi.py results/multi/raw.jsonl "data/coc/mini_nocode.jsonl"
+"""
+import collections
+import glob
+import json
+import re
+import sys
+
+sys.path.insert(0, "scripts")
+from grade import grade_one, last_boxed  # noqa: E402
+
+recs = [json.loads(l) for l in open(sys.argv[1], encoding="utf-8")]
+items = {}
+for pat in sys.argv[2:]:
+    for f in glob.glob(pat):
+        for l in open(f, encoding="utf-8"):
+            r = json.loads(l)
+            items[r["id"]] = r
+
+
+def extract(task, text):
+    if task == "mmlupro":
+        m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", text) or re.findall(r"\\boxed\{\(?([A-J])\)?\}", text)
+        a = m[-1] if m else None
+    elif task == "bbh":
+        m = re.findall(r"answer is:?\s*(.+)", text, re.I)
+        a = m[-1] if m else (last_boxed(text) or "")
+        a = re.sub(r"\\boxed\{(.*)\}", r"\1", a).strip().strip("*$ .").lower() or None
+    else:
+        b = last_boxed(text)
+        a = re.sub(r"\s+", "", b) if b else None
+    if a is None or "..." in a or a in ("?", "") or len(a) > 80:
+        return None
+    return a
+
+
+by = collections.defaultdict(dict)
+for r in recs:
+    r["completion"] = r["completion"].replace("<|im_end|>", "")
+    it = items[r["id"]]
+    task = it["task"]
+    fin = r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else None
+    r["final_ans"] = extract(task, fin) if fin is not None else None
+    r["final_ok"] = grade_one(r, it)["correct"]
+    for p in r.get("probe_log", []):
+        p["ok"] = grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"]
+    by[r["id"]][r["sample"]] = r
+ids = [i for i in by if len(by[i]) >= 4]
+N = 4
+
+
+def ans_ok(i, a):
+    """is answer string a correct for item i (look it up among recorded answers)"""
+    for r in by[i].values():
+        if r["final_ans"] == a:
+            return r["final_ok"]
+        for p in r.get("probe_log", []):
+            if p["ans"] == a:
+                return p["ok"]
+    return False
+
+
+def state(r, t):
+    """answer held by stream r after t generated tokens; (answer, finished)"""
+    if r["n_tokens"] <= t and not r.get("truncated"):
+        return r["final_ans"], True
+    a = None
+    for p in r.get("probe_log", []):
+        if p["cut"] <= t:
+            a = p["ans"]
+    return a, False
+
+
+def run_xs(i, K, streams, start=None, step=256):
+    start = start or {s: 0 for s in streams}
+    rs = {s: by[i][s] for s in streams}
+    horizon = max(start[s] + rs[s]["n_tokens"] for s in streams)
+    t = 0
+    while True:
+        t += step
+        held = {}
+        for s in streams:
+            if t < start[s]:
+                continue
+            held[s] = state(rs[s], t - start[s])
+        cnt = collections.Counter(a for a, _ in held.values() if a is not None)
+        best = cnt.most_common(1)
+        done_all = all(t - start[s] >= rs[s]["n_tokens"] for s in streams)
+        if best and best[0][1] >= K:
+            a = best[0][0]
+            cost = sum(min(rs[s]["n_tokens"], max(0, t - start[s])) for s in streams)
+            return ans_ok(i, a), cost
+        if done_all or t > horizon + step:
+            # no consensus: majority over what the streams ended with (ties -> stream 0's answer)
+            fin = [rs[s]["final_ans"] for s in streams if rs[s]["final_ans"] is not None]
+            c = collections.Counter(fin).most_common()
+            a = c[0][0] if c and (len(c) == 1 or c[0][1] > c[1][1]) else rs[streams[0]]["final_ans"]
+            cost = sum(rs[s]["n_tokens"] for s in streams)
+            return (ans_ok(i, a) if a is not None else rs[streams[0]]["final_ok"]), cost
+
+
+def run_adapt(i, K, t0):
+    r0 = by[i][0]
+    if (r0["n_tokens"] <= t0 and not r0.get("truncated")):
+        return r0["final_ok"], r0["n_tokens"]
+    return run_xs(i, K, list(range(N)), start={0: 0, 1: t0, 2: t0, 3: t0})
+
+
+def run_coc1(i, agree, every):
+    r = by[i][0]
+    log = r.get("probe_log", [])
+    sub = [p for j, p in enumerate(log) if (j + 1) % every == 0]
+    for j in range(agree - 1, len(sub)):
+        a = [sub[t]["ans"] for t in range(j - agree + 1, j + 1)]
+        if a[0] is not None and all(x == a[0] for x in a):
+            return sub[j]["ok"], sub[j]["cut"]
+    return r["final_ok"], r["n_tokens"]
+
+
+def run_sc(i):
+    fin = [by[i][s]["final_ans"] for s in range(N) if by[i][s]["final_ans"] is not None]
+    c = collections.Counter(fin).most_common()
+    cost = sum(by[i][s]["n_tokens"] for s in range(N))
+    if not c:
+        return by[i][0]["final_ok"], cost
+    a = c[0][0] if len(c) == 1 or c[0][1] > c[1][1] else by[i][0]["final_ans"] or c[0][0]
+    return ans_ok(i, a), cost
+
+
+policies = {"single": lambda i: (by[i][0]["final_ok"], by[i][0]["n_tokens"]),
+            "sc@4": run_sc,
+            "coc1 2048x3": lambda i: run_coc1(i, 3, 2),
+            "coc1 1024x2": lambda i: run_coc1(i, 2, 1)}
+for K in (2, 3):
+    policies[f"xs-{K}"] = lambda i, K=K: run_xs(i, K, list(range(N)))
+    for t0 in (1024, 2048, 3072):
+        policies[f"adapt-{K} t0={t0}"] = lambda i, K=K, t0=t0: run_adapt(i, K, t0)
+
+tasks = sorted({items[i]["task"] for i in ids})
+print(f"{len(ids)} problems x {N} streams")
+hdr = f"{'controller':20s} {'ALL acc':>8s} {'tok/prob':>9s} {'x base':>7s} " + " ".join(f"{t:>8s}" for t in tasks)
+print(hdr)
+base_tok = None
+for name, f in policies.items():
+    res = {i: f(i) for i in ids}
+    acc = 100 * sum(ok for ok, _ in res.values()) / len(ids)
+    tok = sum(c for _, c in res.values()) / len(ids)
+    base_tok = base_tok or tok
+    per = " ".join(f"{100 * sum(res[i][0] for i in ids if items[i]['task'] == t) / sum(1 for i in ids if items[i]['task'] == t):8.1f}" for t in tasks)
+    print(f"{name:20s} {acc:8.1f} {tok:9.0f} {tok / base_tok:7.2f} {per}")

research/scripts/sim_single.py

@@ -0,0 +1,82 @@
+"""Single-stream deliberation controllers on recorded trajectories (stream 0 of a --shadow --record-all run).
+
+  base            the ordinary generation
+  fallback        ordinary generation, but a stream that hits the token limit answers with its last probe instead
+                  of being cut off (truncation scores 0 anyway)
+  coc E x A       commit when A successive probes (one every E*1024 tokens) give the same answer
+  coc E x A + fb  the same plus the limit fallback
+  budget          budget-aware consensus: strict (3 agreeing probes) in the first half of the budget, 2 agreeing probes
+                  in the second half, plus the limit fallback
+usage: sim_single.py RAW.jsonl "ITEMS_GLOB" [...]
+"""
+import collections
+import glob
+import json
+import sys
+
+sys.path.insert(0, "scripts")
+from grade import grade_one  # noqa: E402
+
+recs = [json.loads(l) for l in open(sys.argv[1], encoding="utf-8")]
+items = {}
+for pat in sys.argv[2:]:
+    for f in glob.glob(pat):
+        for l in open(f, encoding="utf-8"):
+            r = json.loads(l)
+            items[r["id"]] = r
+recs = [r for r in recs if r.get("sample", 0) == 0 and r["id"] in items and r.get("probe_log") is not None]
+limit = max(r["n_tokens"] for r in recs)
+for r in recs:
+    r["completion"] = r["completion"].replace("<|im_end|>", "")
+    it = items[r["id"]]
+    r["ok"] = grade_one(r, it)["correct"]
+    for p in r["probe_log"]:
+        p["ok"] = grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"]
+
+
+def coc(r, every, agree, fb, budget_aware=False):
+    log = r["probe_log"]
+    sub = [p for j, p in enumerate(log) if (j + 1) % every == 0]
+    for j in range(len(sub)):
+        need = agree
+        if budget_aware:
+            need = 3 if sub[j]["cut"] < limit / 2 else 2
+        if j + 1 < need:
+            continue
+        a = [sub[t]["ans"] for t in range(j - need + 1, j + 1)]
+        if a[0] is not None and all(x == a[0] for x in a):
+            return sub[j]["ok"], sub[j]["cut"] + sub[j]["n"]
+    if fb and r.get("truncated") and log:
+        last = [p for p in log if p["ans"] is not None]
+        if fb == "maj":  # most frequent answer among the probes of the second half (ties -> latest)
+            last = [p for p in last if p["cut"] >= limit / 2] or last
+            if last:
+                c = collections.Counter(p["ans"] for p in last)
+                top = max(c.values())
+                pick = [p for p in last if c[p["ans"]] == top][-1]
+                return pick["ok"], r["n_tokens"]
+        if last:
+            return last[-1]["ok"], r["n_tokens"]
+    return r["ok"], r["n_tokens"]
+
+
+pols = {"base": lambda r: (r["ok"], r["n_tokens"]), "fallback": lambda r: coc(r, 1, 99, True)}
+for e, a in [(1, 2), (1, 3), (2, 2), (2, 3)]:
+    pols[f"coc {e * 1024}x{a}"] = lambda r, e=e, a=a: coc(r, e, a, False)
+    pols[f"coc {e * 1024}x{a} + fb"] = lambda r, e=e, a=a: coc(r, e, a, True)
+pols["budget-aware + fb"] = lambda r: coc(r, 1, 0, True, budget_aware=True)
+pols["fallback-maj"] = lambda r: coc(r, 1, 99, "maj")
+pols["coc 2048x3 + fb-maj"] = lambda r: coc(r, 2, 3, "maj")
+
+tasks = sorted({r["task"] for r in recs})
+print(f"{len(recs)} problems, token limit {limit}")
+print(f"{'controller':20s} {'ALL':>6s} {'+':>3s} {'-':>3s} {'tokens':>7s} " + " ".join(f"{t:>9s}" for t in tasks))
+for name, f in pols.items():
+    res = [(r, *f(r)) for r in recs]
+    acc = 100 * sum(ok for _, ok, _ in res) / len(res)
+    gain = sum((not r["ok"]) and ok for r, ok, _ in res)
+    loss = sum(r["ok"] and not ok for r, ok, _ in res)
+    tok = sum(t for _, _, t in res) / len(res)
+    per = " ".join(f"{100 * sum(ok for r, ok, _ in res if r['task'] == t) / sum(1 for r in recs if r['task'] == t):9.1f}"
+                   for t in tasks)
+    print(f"{name:20s} {acc:6.1f} {gain:3d} {loss:3d} {tok:7.0f} {per}")

research/scripts/soft_post.py

@@ -0,0 +1,36 @@
+"""Turn the output of lfm3-soft-think.exe (cpp/soft_think.cpp) into records for grade.py / cmp.py.
+The visible thinking is the chosen token of every step (for soft/gumbel the argmax of the mixture weights); n_tokens
+counts thinking steps + answer tokens, i.e. the tokens the user would see.
+usage: soft_post.py --txt results/soft/m180_soft_s0.txt --out results/st2_soft/mini.jsonl --items data/coc/math180.jsonl
+"""
+import argparse
+import json
+import os
+
+from transformers import AutoTokenizer
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--txt", nargs="+", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--items", required=True)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+args = ap.parse_args()
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+task = {json.loads(l)["id"]: json.loads(l)["task"] for l in open(args.items, encoding="utf-8")}
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+n = 0
+with open(args.out, "w", encoding="utf-8") as f:
+    for s, path in enumerate(args.txt):
+        for line in open(path, encoding="utf-8"):
+            parts = line.rstrip("\n").split("\t")
+            if len(parts) != 6:
+                continue
+            id_, steps, forced, cold, think, answer = parts
+            think = [int(x) for x in think.split()]
+            answer = [int(x) for x in answer.split()]
+            f.write(json.dumps({"id": id_, "task": task[id_], "sample": s,
+                                "completion": "<think>" + tok.decode(think) + "</think>" + tok.decode(answer),
+                                "n_tokens": int(steps) + len(answer), "truncated": False, "think_steps": int(steps),
+                                "forced": forced == "1", "cold": cold == "1"}, ensure_ascii=False) + "\n")
+            n += 1
+print(f"{n} records -> {args.out}")

research/scripts/soft_stats.py

@@ -0,0 +1,18 @@
+"""Summary of the latent-thinking runs (results/st2_<mode>): accuracy, visible thinking length, how thinking ended.
+usage: soft_stats.py discrete soft gumbel moi mti"""
+import json
+import statistics
+import sys
+
+for mode in sys.argv[1:]:
+    recs = {json.loads(l)["id"]: json.loads(l) for l in open(f"results/st2_{mode}/mini.jsonl", encoding="utf-8")}
+    graded = {json.loads(l)["id"]: json.loads(l)["correct"] for l in open(f"results/st2_{mode}/graded.jsonl", encoding="utf-8")}
+    n = len(recs)
+    steps = [r["think_steps"] for r in recs.values()]
+    forced = sum(r["forced"] for r in recs.values())
+    cold = sum(r.get("cold", False) for r in recs.values())
+    natural = n - forced - cold
+    acc_nat = [graded[i] for i, r in recs.items() if not r["forced"] and not r.get("cold")]
+    print(f"{mode:9s} n={n}  acc {100 * sum(graded.values()) / n:5.1f}  thinking mean {statistics.mean(steps):6.0f} "
+          f"median {statistics.median(steps):6.0f}  ended: natural {natural} (acc {100 * sum(acc_nat) / max(1, len(acc_nat)):.0f}%)"
+          f", budget {forced}, cold stop {cold}")

research/scripts/soft_think.py

@@ -0,0 +1,177 @@
+"""Latent ("soft") thinking for LFM3-1.7B, with a discrete control in the same implementation.
+
+During the thinking phase the next input is not one sampled token but a continuous "concept token": the embedding
+average of the top candidate tokens weighted by their probabilities (Soft Thinking, Zhang et al. 2025). Later work
+found that plain soft thinking collapses towards greedy decoding and that randomness restores it (Gumbel-softmax
+weights), so three modes are compared on the same problems with the same budget and answer read-out:
+  discrete     ordinary sampling (temperature 0.6, top-k 50, top-p 0.95)          -- control
+  soft         probability-weighted top-k embedding mixture                         -- Soft Thinking
+  soft_gumbel  weights softmax((log p + Gumbel)/tau) over the same candidates       -- randomized soft thinking
+Thinking ends when the most probable token is </think> (soft modes) / </think> is sampled (discrete), or at the
+budget, where </think> is forced; the answer is then sampled discretely.
+usage: soft_think.py --items data/coc/math180.jsonl --mode soft --out results/st_soft/mini.jsonl [--limit 120]
+       soft_think.py --selftest      (batched left-padded decoding must match single-sequence decoding)
+"""
+import argparse
+import json
+import math
+import os
+import time
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+THINK_END, EOS = 124902, 124900
+ap = argparse.ArgumentParser()
+ap.add_argument("--items")
+ap.add_argument("--mode", choices=["discrete", "soft", "soft_gumbel"], default="soft")
+ap.add_argument("--out")
+ap.add_argument("--limit", type=int, default=0)
+ap.add_argument("--budget", type=int, default=6144)
+ap.add_argument("--answer-tokens", type=int, default=768)
+ap.add_argument("--batch", type=int, default=16)
+ap.add_argument("--topk-soft", type=int, default=15)
+ap.add_argument("--tau", type=float, default=0.5)
+ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--model", default="models/LFM3-1.7B")
+ap.add_argument("--selftest", action="store_true")
+ap.add_argument("--cpu", action="store_true")
+args = ap.parse_args()
+
+dev = "cpu" if args.cpu else "cuda"
+tok = AutoTokenizer.from_pretrained(args.model)
+model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.float32 if args.cpu else torch.bfloat16,
+                                             trust_remote_code=True).to(dev).eval()
+E = model.get_input_embeddings().weight  # (V, D)
+PAD = tok.pad_token_id if tok.pad_token_id is not None else 0
+
+
+def left_pad(seqs):
+    L = max(len(s) for s in seqs)
+    ids = torch.full((len(seqs), L), PAD, dtype=torch.long)
+    mask = torch.zeros((len(seqs), L), dtype=torch.long)
+    for i, s in enumerate(seqs):
+        ids[i, L - len(s):] = torch.tensor(s)
+        mask[i, L - len(s):] = 1
+    return ids.to(dev), mask.to(dev)
+
+
+def filtered_probs(logits, temperature=0.6, top_k=50, top_p=0.95):
+    logits = logits.float() / temperature
+    v, ix = logits.topk(top_k, dim=-1)
+    p = torch.softmax(v, dim=-1)
+    keep = (p.cumsum(-1) - p) < top_p
+    p = p * keep
+    return p / p.sum(-1, keepdim=True), ix
+
+
+@torch.no_grad()
+def run_batch(prompts, mode, seed):
+    g = torch.Generator(device=dev).manual_seed(seed)
+    ids, mask = left_pad(prompts)
+    out = model(input_ids=ids, attention_mask=mask, use_cache=True)
+    cache, logits = out.past_key_values, out.logits[:, -1]
+    B = len(prompts)
+    phase = ["think"] * B
+    think_steps, transcript, answer = [0] * B, [[] for _ in range(B)], [[] for _ in range(B)]
+    forced = [False] * B
+    for step in range(args.budget + args.answer_tokens + 2):
+        if all(ph == "done" for ph in phase):
+            break
+        p, ix = filtered_probs(logits)
+        sampled = ix.gather(1, torch.multinomial(p, 1, generator=g)).squeeze(1)
+        top1 = ix[:, 0]
+        emb = E[sampled].clone()
+        if mode != "discrete":
+            k = args.topk_soft
+            ps, ixs = p[:, :k], ix[:, :k]
+            if mode == "soft_gumbel":
+                gum = -torch.log(-torch.log(torch.rand(ps.shape, generator=g, device=dev).clamp_min(1e-9)))
+                w = torch.softmax((torch.log(ps.clamp_min(1e-12)) + gum) / args.tau, dim=-1) * (ps > 0)
+            else:
+                w = ps
+            w = w / w.sum(-1, keepdim=True)
+            soft = torch.einsum("bk,bkd->bd", w.to(E.dtype), E[ixs])
+        for i in range(B):
+            if phase[i] == "think":
+                end = (top1[i].item() == THINK_END) if mode != "discrete" else (sampled[i].item() == THINK_END)
+                if think_steps[i] >= args.budget:
+                    end, forced[i] = True, True
+                if end:
+                    emb[i] = E[THINK_END]
+                    phase[i] = "answer"
+                    continue
+                think_steps[i] += 1
+                transcript[i].append(top1[i].item() if mode != "discrete" else sampled[i].item())
+                if mode != "discrete":
+                    emb[i] = soft[i]
+            elif phase[i] == "answer":
+                t = sampled[i].item()
+                if t == EOS or len(answer[i]) >= args.answer_tokens:
+                    phase[i] = "done"
+                else:
+                    answer[i].append(t)
+        mask = torch.cat([mask, torch.ones((B, 1), dtype=mask.dtype, device=dev)], dim=1)
+        out = model(inputs_embeds=emb[:, None, :], attention_mask=mask, past_key_values=cache, use_cache=True)
+        cache, logits = out.past_key_values, out.logits[:, -1]
+    return [{"think_steps": think_steps[i], "forced": forced[i], "thinking": tok.decode(transcript[i]),
+             "answer": tok.decode(answer[i])} for i in range(B)]
+
+
+def chat_ids(prompt):
+    ids = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=True, add_generation_prompt=True,
+                                  return_dict=False)
+    return ids["input_ids"] if isinstance(ids, dict) else ids
+
+
+if args.selftest:
+    a = chat_ids("What is 2+2?")
+    b = chat_ids("Explain in one sentence why the sky is blue, and mention Rayleigh scattering.")
+    with torch.no_grad():
+        single = model(input_ids=torch.tensor([b], device=dev)).logits[0, -1].float()
+        ids, mask = left_pad([a, b])
+        batched = model(input_ids=ids, attention_mask=mask).logits[1, -1].float()
+        sa = model(input_ids=torch.tensor([a], device=dev)).logits[0, -1].float()
+        ba = model(input_ids=ids, attention_mask=mask).logits[0, -1].float()
+    print("prefill max|diff| long row", (single - batched).abs().max().item(), "short (padded) row",
+          (sa - ba).abs().max().item(), "logit scale", single.abs().max().item())
+    # decode 8 forced tokens both ways
+    forced_toks = tok("<think>We add the numbers", add_special_tokens=False)["input_ids"][:6]
+    with torch.no_grad():
+        o1 = model(input_ids=torch.tensor([a], device=dev), use_cache=True)
+        c1 = o1.past_key_values
+        o2 = model(input_ids=ids, attention_mask=mask, use_cache=True)
+        c2, m2 = o2.past_key_values, mask
+        for t in forced_toks:
+            o1 = model(input_ids=torch.tensor([[t]], device=dev), past_key_values=c1, use_cache=True)
+            c1 = o1.past_key_values
+            m2 = torch.cat([m2, torch.ones((2, 1), dtype=m2.dtype, device=dev)], 1)
+            o2 = model(inputs_embeds=E[torch.tensor([t, t], device=dev)][:, None, :], attention_mask=m2,
+                       past_key_values=c2, use_cache=True)
+            c2 = o2.past_key_values
+    print("decode max|diff| padded row after 6 steps", (o1.logits[0, -1].float() - o2.logits[0, -1].float()).abs().max().item())
+    raise SystemExit
+
+items = [json.loads(l) for l in open(args.items, encoding="utf-8")]
+items = items[: args.limit] if args.limit else items
+done = set()
+if os.path.exists(args.out):
+    done = {json.loads(l)["id"] for l in open(args.out, encoding="utf-8")}
+items = [it for it in items if it["id"] not in done]
+items.sort(key=lambda it: len(it["prompt"]))
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+t0, n_tok = time.time(), 0
+with open(args.out, "a", encoding="utf-8") as f:
+    for s in range(0, len(items), args.batch):
+        chunk = items[s:s + args.batch]
+        res = run_batch([chat_ids(it["prompt"]) for it in chunk], args.mode, args.seed * 100003 + s)
+        for it, r in zip(chunk, res):
+            n = r["think_steps"] + len(tok(r["answer"], add_special_tokens=False)["input_ids"])
+            n_tok += n
+            f.write(json.dumps({"id": it["id"], "task": it["task"], "sample": 0, "mode": args.mode,
+                                "completion": "<think>" + r["thinking"] + "</think>" + r["answer"], "n_tokens": n,
+                                "truncated": False, "think_steps": r["think_steps"], "forced": r["forced"]},
+                               ensure_ascii=False) + "\n")
+        f.flush()
+        el = time.time() - t0
+        print(f"{s + len(chunk)}/{len(items)}  {n_tok / el:.0f} tok/s  {el / 60:.1f} min", flush=True)

research/scripts/start_server.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+# usage: start_server.sh <gguf> [extra llama-server args, e.g. --control-vector f.gguf]   -- identical server configuration for every model under test
+cd /c/Users/user/lfm3proj
+taskkill //F //IM llama-server.exe >/dev/null 2>&1; sleep 2
+(/c/Users/user/llamabin/llama-server.exe -m "$1" "${@:2}" -ngl 99 -np 32 -c 557056 -bs -fa on -ctxcp 0 -cram 0 --metrics --port 8080 --no-webui > results/server.log 2>&1 &)
+for i in $(seq 1 100); do curl -s -m 2 localhost:8080/health | grep -q '"ok"' && break; grep -q "exiting" results/server.log && break; sleep 3; done
+curl -s localhost:8080/health; echo

research/scripts/stress.py

@@ -0,0 +1,340 @@
+"""Stress-test module for code answers: the way a strong competitive programmer checks a solution.
+
+For a program that already passes the public examples (after test-and-repair), the model writes, in separate turns:
+  * two brute-force solutions (simplest obviously-correct method, any complexity), each kept only if it passes the
+    public examples;
+  * a generator of small random valid inputs.
+Random inputs are run through the brute forces and the program. An input is a counterexample only when the brute
+forces agree with each other and the program disagrees (or crashes / hangs on a tiny input). The model is then
+shown the smallest counterexample and fixes the program; a fix is kept only if it still passes the public examples
+and agrees with the brute forces on strictly more random inputs. Hidden tests are never used.
+
+Records that are not candidates (no program, failing public examples) are copied unchanged, so the comparison with
+the source run is exactly paired.
+usage: stress.py --src results/lcb_rep/lcb.jsonl --items data/lcb/lcb6.jsonl --out results/lcb_stress/lcb.jsonl
+"""
+import argparse
+import json
+import os
+import re
+import sys
+import threading
+import time
+import urllib.request
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from grade import extract_code  # noqa: E402
+from lcb_tests import HARNESS, _run, run_lcb  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--src", required=True)
+ap.add_argument("--items", nargs="+", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--rounds", type=int, default=2)
+ap.add_argument("--think", type=int, default=4096)
+ap.add_argument("--reserve", type=int, default=1536)
+ap.add_argument("--n-inputs", type=int, default=40)
+ap.add_argument("--workers", type=int, default=24)
+ap.add_argument("--tle", action="store_true", help="also check the speed on maximum-size inputs")
+ap.add_argument("--time-limit", type=float, default=6.0)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+ap.add_argument("--port", type=int, default=8080)
+args = ap.parse_args()
+
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
+            "samplers": ["top_k", "top_p", "temperature"]}
+
+STATEMENT = re.compile(r"\n\n(Write a Python 3 program|Complete the following starter code|You will use the following "
+                       r"starter code).*", re.S)
+
+BRUTE = ("{statement}\n\n---\n\nWrite a BRUTE-FORCE solution to this problem: the simplest program that is obviously "
+         "correct, for example by trying every possibility. Ignore efficiency completely; it will only be run on tiny "
+         "inputs. {form}")
+FORM_STDIN = ("It must read the input from standard input in the format above and print the answer. Put the complete "
+              "program in a single ```python code block.")
+FORM_FUNC = ("Complete the same starter code (the class Solution with the method {func}) and put the full class in a "
+             "single ```python code block.")
+GEN = ("{statement}\n\n---\n\nWrite a Python program that prints ONE random valid test input for this problem. "
+       "{fmt} The input will be used to compare a solution with a slow brute force, so make it tiny and tricky: sizes "
+       "from the smallest allowed up to about 6, and values drawn from a very small range (for example 1 to 3, or "
+       "only two different letters) so that equal values, repeats, ties and edge cases appear often. Every constraint "
+       "of the problem must still hold. Use the random module seeded with int(sys.argv[1]). Put the complete program "
+       "in a single ```python code block.")
+FMT_STDIN = "Print the input exactly in the input format of the problem (as it would be given on standard input)."
+FMT_FUNC = ("The input is the list of arguments of the method {func}: print each argument as one JSON value "
+            "(json.dumps) on its own line, in the order of the parameters.")
+MAXGEN = ("{statement}\n\n---\n\nWrite a Python program that prints ONE valid test input of the MAXIMUM size allowed "
+          "by the constraints (the largest allowed sizes and values), for checking the running time of a solution. "
+          "{fmt} Use random content, seeded with int(sys.argv[1]), and make sure every constraint holds. Put the "
+          "complete program in a single ```python code block.")
+SLOW = ("Your program gives correct answers on small inputs, but it is too slow: on a valid input of the maximum size "
+        "allowed by the constraints it did not finish within the time limit. The input starts like this:\n{head}\n\n"
+        "Find an asymptotically faster algorithm (check the constraints to see which complexity is needed) and give "
+        "the complete program in a single ```python code block.")
+FEEDBACK = ("I tested your program against a slow brute-force solution on small random inputs, and they disagree.\n\n"
+            "{fail}\n\nFind the mistake (re-read the problem statement if your understanding may be wrong) and give the "
+            "complete corrected program in a single ```python code block.")
+
+
+def post(body):
+    for attempt in range(5):
+        try:
+            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 4:
+                raise
+            time.sleep(5)
+
+
+def generate(msgs, seed):
+    ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_dict=False)
+    r = post({"prompt": ids, "n_predict": args.think + args.reserve, "seed": seed, "cache_prompt": False,
+              "return_tokens": True, **SAMPLING})
+    gen = [t for t in r["tokens"] if t != EOS]
+    if THINK_END in gen:
+        return r["content"].split("</think>", 1)[1], r["tokens_predicted"]
+    lim = min(args.think, len(gen))
+    cut = next((j for j in range(lim, 0, -1) if gen[j - 1] in ENDS), lim)
+    r2 = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.reserve, "seed": seed + 17,
+               "cache_prompt": False, **SAMPLING})
+    return r2["content"], cut + r2["tokens_predicted"]
+
+
+def outputs(code, item, inputs):
+    """Outputs of a candidate on the inputs (None where it crashes or times out)."""
+    if item["testtype"] == "functional":
+        src = HARNESS.split("_tests = json.loads")[0].replace("__CODE__", code) + (
+            "\nfor _inp in json.loads(sys.stdin.read()):\n"
+            "    try:\n"
+            "        _r = getattr(Solution(), %r)(*[json.loads(_l) for _l in _inp.strip().split('\\n')])\n"
+            "        print(json.dumps(list(_r) if isinstance(_r, tuple) else _r))\n"
+            "    except Exception:\n"
+            "        print('null')\n"
+            "    sys.stdout.flush()\n" % item["func_name"])
+        rc, so, _ = _run(src, json.dumps(inputs), 3 + 0.5 * len(inputs))
+        lines = so.strip().split("\n") if so.strip() else []
+        return [(l if l != "null" else None) for l in lines] + [None] * (len(inputs) - len(lines))
+    res = []
+    for inp in inputs:
+        rc, so, _ = _run(code, inp, 3)
+        res.append(" ".join(so.split()) if rc == 0 and so.strip() else None)
+    return res
+
+
+def make_inputs(gen_code, item):
+    seen, out = set(), []
+    for k in range(args.n_inputs):
+        rc, so, _ = _run(f"import sys\nsys.argv = ['gen', '{k}']\n" + gen_code, "", 3)  # _run passes no argv
+        inp = so.strip()
+        if rc != 0 or not inp or len(inp) > 2000 or inp in seen:
+            continue
+        if item["testtype"] == "functional":
+            try:
+                [json.loads(l) for l in inp.split("\n")]
+            except Exception:
+                continue
+        seen.add(inp)
+        out.append(inp + "\n")
+    return out
+
+
+def process(rec):
+    try:
+        return _process(rec)
+    except Exception as e:
+        return {**rec, "stress_error": repr(e)[:200]}
+
+
+def _process(rec):
+    item = items[rec["id"]]
+    if "</think>" not in rec["completion"]:
+        return rec
+    final = rec["completion"].split("</think>", 1)[1]
+    code = extract_code(final)
+    if code is None:
+        return rec
+    k, n, _ = run_lcb(code, item, item["public"])
+    if k < n:
+        return rec
+    statement = STATEMENT.sub("", item["prompt"])
+    func = item.get("func_name")
+    seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + 9000
+    used = [0]
+
+    def gen(msgs, sd):
+        text, t = generate(msgs, sd)
+        used[0] += t
+        return text
+
+    status, dbg, tests = correctness_inputs(item, statement, func, seed, gen)
+    score = lambda c: stress_score(c, item, tests)
+    ok0, bad = score(code) if tests else (0, [])
+    msgs = [{"role": "user", "content": item["prompt"]}]
+    kept = 0
+    if bad:  # phase 1: disagreement with the brute forces on small inputs
+        attempts = [(ok0, final, code)]
+        cur = final
+        for rnd in range(args.rounds):
+            fail = "\n\n".join(
+                f"Input:\n{i.strip()}\n\nBrute-force output:\n{e}\n\nYour program's output:\n"
+                f"{g if g is not None else '(it crashed or ran too long)'}"
+                for i, e, g in sorted(bad, key=lambda x: len(x[0]))[:2])
+            msgs = msgs + [{"role": "assistant", "content": cur}, {"role": "user", "content": FEEDBACK.format(fail=fail)}]
+            cur = gen(msgs, seed + 100 + rnd)
+            c = extract_code(cur)
+            if c is None or run_lcb(c, item, item["public"])[0] < len(item["public"]):
+                attempts.append((-1, cur, c))
+                continue
+            okr, bad = score(c)
+            attempts.append((okr, cur, c))
+            if not bad:
+                break
+        kept = max(range(len(attempts)), key=lambda k: (attempts[k][0], -k))  # must agree on strictly more inputs
+        ok_best, final, code = attempts[kept]
+        status = f"fixed {ok_best}/{len(tests)} (was {ok0})" if kept else f"unfixed {ok0}/{len(tests)}"
+        msgs = [{"role": "user", "content": item["prompt"]}]
+    elif tests:
+        status = f"agree {len(tests)}"
+    tle = ""
+    if args.tle:  # phase 2: speed on inputs of the maximum size allowed by the constraints
+        fmt = FMT_FUNC.format(func=func) if item["testtype"] == "functional" else FMT_STDIN
+        gtext = gen([{"role": "user", "content": MAXGEN.format(statement=statement, fmt=fmt)}], seed + 300)
+        big = make_big(extract_code(gtext), item)
+        dbg["big_inputs"] = len(big)
+        if big and not finishes(code, item, big):
+            tle = "slow"
+            ok_ref = score(code)[0] if tests else 0
+            cur = final
+            for rnd in range(args.rounds):
+                msgs = msgs + [{"role": "assistant", "content": cur},
+                               {"role": "user", "content": SLOW.format(head=big[0][:200].strip())}]
+                cur = gen(msgs, seed + 400 + rnd)
+                c = extract_code(cur)
+                if (c is None or run_lcb(c, item, item["public"])[0] < len(item["public"])
+                        or (tests and score(c)[0] < ok_ref) or not finishes(c, item, big)):
+                    continue
+                final, code, tle = cur, c, "sped_up"
+                break
+    out = {**rec, "n_tokens": rec["n_tokens"] + used[0], "stress": status, "stress_dbg": dbg, "tle": tle}
+    if kept or tle == "sped_up":
+        out["completion"] = "<think>(stress-tested)</think>" + final
+    return out
+
+
+def correctness_inputs(item, statement, func, seed, gen):
+    """Brute forces (kept only if they pass the public examples) + small random inputs -> tests they agree on."""
+    form = FORM_FUNC.format(func=func) if item["testtype"] == "functional" else FORM_STDIN
+    starter = re.search(r"```python\n.*?```", item["prompt"], re.S)
+    brute_prompt = BRUTE.format(statement=statement, form=form) + (f"\n\n{starter.group(0)}" if starter else "")
+    brutes = []
+    for b in range(2):
+        bc = extract_code(gen([{"role": "user", "content": brute_prompt}], seed + b))
+        if bc and run_lcb(bc, item, item["public"])[0] == len(item["public"]):
+            brutes.append(bc)
+    dbg = {"brutes": len(brutes)}
+    if not brutes:
+        return "no_brute", dbg, []
+    fmt = FMT_FUNC.format(func=func) if item["testtype"] == "functional" else FMT_STDIN
+    text = gen([{"role": "user", "content": GEN.format(statement=statement, fmt=fmt)}], seed + 7)
+    gc = extract_code(text)
+    inputs = make_inputs(gc, item) if gc else []
+    dbg.update(inputs=len(inputs), gen=(gc or text)[-1500:])
+    if not inputs:
+        return "no_inputs", dbg, []
+    outs = [outputs(bc, item, inputs) for bc in brutes]
+    tests = [(inp, outs[0][j]) for j, inp in enumerate(inputs)
+             if outs[0][j] is not None and all(o[j] == outs[0][j] for o in outs)]
+    dbg.update(tests=len(tests), sample_input=inputs[0][:300])
+    return ("no_agreement" if not tests else ""), dbg, tests
+
+
+def stress_score(c, item, tests):
+    got = outputs(c, item, [i for i, _ in tests])
+    bad = [(i, e, g) for (i, e), g in zip(tests, got) if not same(g, e)]
+    return len(tests) - len(bad), bad
+
+
+def make_big(code, item):
+    if not code:
+        return []
+    out = []
+    for k in range(2):
+        rc, so, _ = _run(f"import sys\nsys.argv = ['gen', '{k}']\n" + code, "", 15)
+        if rc == 0 and so.strip() and len(so) < 30_000_000:
+            if item["testtype"] == "functional":
+                try:
+                    [json.loads(l) for l in so.strip().split("\n")]
+                except Exception:
+                    continue
+            out.append(so.strip() + "\n")
+    return out
+
+
+def finishes(code, item, big):
+    if item["testtype"] == "functional":
+        src = HARNESS.split("_tests = json.loads")[0].replace("__CODE__", code) + (
+            "\nfor _inp in json.loads(sys.stdin.read()):\n"
+            "    getattr(Solution(), %r)(*[json.loads(_l) for _l in _inp.strip().split('\\n')])\n" % item["func_name"])
+        return all(_run(src, json.dumps([b]), args.time_limit)[0] == 0 for b in big)
+    return all(_run(code, b, args.time_limit)[0] == 0 for b in big)
+
+
+def same(got, exp):
+    if got is None:
+        return False
+    a, b = got.split(), exp.split()
+    if len(a) != len(b):
+        return False
+    for x, y in zip(a, b):
+        if x != y:
+            try:
+                if abs(float(x) - float(y)) > 1e-6 * max(1.0, abs(float(y))):
+                    return False
+            except ValueError:
+                return False
+    return True
+
+
+items = {}
+for path in args.items:
+    with open(path, encoding="utf-8") as f:
+        for l in f:
+            r = json.loads(l)
+            items[r["id"]] = r
+with open(args.src, encoding="utf-8") as f:
+    recs = [json.loads(l) for l in f]
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+t0 = time.time()
+lock = threading.Lock()
+done = [0]
+
+
+def run(rec):
+    r = process(rec)
+    with lock:
+        done[0] += 1
+        if done[0] % 20 == 0:
+            print(f"{done[0]}/{len(recs)}  {(time.time() - t0) / 60:.1f} min", flush=True)
+    return r
+
+
+with ThreadPoolExecutor(args.workers) as ex:
+    out = list(ex.map(run, recs))
+with open(args.out, "w", encoding="utf-8") as f:
+    for r in out:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+st = [r.get("stress", "") for r in out]
+print(f"{len(out)} records: " + ", ".join(f"{k} {sum(s.startswith(k) for s in st)}" for k in
+                                          ("agree", "fixed", "unfixed", "no_brute", "no_inputs", "no_agreement")) +
+      f", slow {sum(r.get('tle') in ('slow', 'sped_up') for r in out)}, sped up {sum(r.get('tle') == 'sped_up' for r in out)}"
+      f", errors {sum('stress_error' in r for r in out)}, {(time.time() - t0) / 60:.1f} min")

research/scripts/svg_program.py

@@ -0,0 +1,156 @@
+"""Program-aided drawing for SVG tasks: instead of writing coordinates by hand, the model writes a Python program that
+computes the geometry (math.cos / math.sin, loops) and prints the SVG. The program is run in the sandbox; if it fails
+or prints no valid SVG, the model gets the error and fixes the program (up to --rounds times). The printed SVG becomes
+the answer, so the result is scored and rendered exactly like a direct SVG answer.
+usage: svg_program.py --tasks data/visual/tasks.jsonl --out results/vis_svgprog/vis.jsonl --n 2
+"""
+import argparse
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import time
+import urllib.request
+import xml.dom.minidom
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import sandbox  # noqa: E402
+from grade import extract_code  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+ap = argparse.ArgumentParser()
+ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+ap.add_argument("--out", required=True)
+ap.add_argument("--n", type=int, default=2)
+ap.add_argument("--rounds", type=int, default=2)
+ap.add_argument("--think", type=int, default=8192)
+ap.add_argument("--reserve", type=int, default=4096)
+ap.add_argument("--workers", type=int, default=16)
+ap.add_argument("--port", type=int, default=8080)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+args = ap.parse_args()
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
+            "samplers": ["top_k", "top_p", "temperature"]}
+PROMPT = ("{task}\n\nDo not write the SVG by hand. Write a Python 3 program that builds the SVG and prints it to "
+          "standard output. Compute every position with arithmetic and the math module (math.cos, math.sin, "
+          "math.radians, loops over repeated parts) instead of guessing numbers, and remember that in SVG the y axis "
+          "points down. Use only the standard library. Give the program in a single ```python code block.")
+FIX = ("I ran your program. {problem}\n\nFix the program and give the complete corrected program in a single "
+       "```python code block.")
+TMP = os.path.join(tempfile.gettempdir(), "lfm3_svgprog")
+os.makedirs(TMP, exist_ok=True)
+
+
+def wait_server():
+    """The night chain restarts the shared server with other models; generate only on the original model."""
+    while True:
+        try:
+            p = json.loads(urllib.request.urlopen(f"http://localhost:{args.port}/props", timeout=10).read())
+            if "lfm3-orig-q5km" in p.get("model_path", ""):
+                return p["default_generation_settings"]["n_ctx"]
+        except Exception:
+            pass
+        time.sleep(15)
+
+
+def post(body):
+    for attempt in range(40):
+        n_ctx = wait_server()
+        body = {**body, "n_predict": min(body["n_predict"], n_ctx - len(body["prompt"]) - 16)}  # fit the slot
+        try:
+            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 39:
+                raise
+            time.sleep(10)
+
+
+def generate(msgs, seed):
+    ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_dict=False)
+    r = post({"prompt": ids, "n_predict": args.think + args.reserve, "seed": seed, "cache_prompt": False,
+              "return_tokens": True, **SAMPLING})
+    gen = [t for t in r["tokens"] if t != EOS]
+    if THINK_END in gen:
+        return r["content"].split("</think>", 1)[1], r["tokens_predicted"]
+    lim = min(args.think, len(gen))
+    cut = next((j for j in range(lim, 0, -1) if gen[j - 1] in ENDS), lim)
+    pre = tok("```python\n", add_special_tokens=False)["input_ids"]
+    r2 = post({"prompt": ids + gen[:cut] + [THINK_END] + pre, "n_predict": args.reserve, "seed": seed + 17,
+               "cache_prompt": False, **SAMPLING})
+    c = r2["content"]
+    return "```python\n" + (c if "```" in c else c.rstrip() + "\n```"), cut + len(pre) + r2["tokens_predicted"]
+
+
+def run_program(code):
+    """Runs the program; returns (svg or None, problem description or None)."""
+    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(code)
+    try:
+        r = sandbox.run([sys.executable, "-I", path], input=b"", capture_output=True, timeout=20, cwd=TMP,
+                        creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
+        out, err = r.stdout.decode(errors="replace"), r.stderr.decode(errors="replace")
+        if r.returncode != 0:
+            return None, "It failed:\n" + "\n".join(err.strip().splitlines()[-6:])
+    except subprocess.TimeoutExpired:
+        return None, "It did not finish within 20 seconds."
+    finally:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+    m = re.search(r"<svg.*</svg>", out, re.S)
+    if not m:
+        return None, "It printed no <svg>...</svg> document. Its output began with:\n" + out[:300]
+    try:
+        xml.dom.minidom.parseString(m.group(0).encode("utf-8"))
+    except Exception as e:
+        return None, f"The printed SVG is not valid XML: {str(e)[:200]}"
+    return m.group(0), None
+
+
+def process(job):
+    task, s = job
+    seed = zlib.crc32(task["id"].encode()) % 1000003 * 7 + s + 11000
+    msgs = [{"role": "user", "content": PROMPT.format(task=task["prompt"].split("Answer with")[0].strip())}]
+    used, svg, final, rounds = 0, None, "", 0
+    for rnd in range(args.rounds + 1):
+        final, n = generate(msgs, seed + rnd)
+        used += n
+        code = extract_code(final)
+        if code is None:
+            problem = "Your answer does not contain the program in a ```python code block."
+        else:
+            svg, problem = run_program(code)
+        if svg:
+            break
+        if rnd < args.rounds:
+            rounds += 1
+            msgs = msgs + [{"role": "assistant", "content": final}, {"role": "user", "content": FIX.format(problem=problem)}]
+    answer = f"```svg\n{svg}\n```" if svg else final
+    print(f"{task['id']} s{s}: {'svg' if svg else 'no svg'} after {rounds} repairs, {used} tokens", flush=True)
+    return {"id": task["id"], "task": "visual", "sample": s, "completion": "<think>(program)</think>" + answer,
+            "n_tokens": used, "truncated": False, "program_rounds": rounds}
+
+
+tasks = [json.loads(l) for l in open(args.tasks, encoding="utf-8")]
+jobs = [(t, s) for t in tasks if t["kind"] == "svg" for s in range(args.n)]
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+t0 = time.time()
+with ThreadPoolExecutor(args.workers) as ex:
+    out = list(ex.map(process, jobs))
+with open(args.out, "w", encoding="utf-8") as f:
+    for r in out:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+print(f"{len(out)} answers, {sum('```svg' in r['completion'] for r in out)} with a valid SVG, "
+      f"{(time.time() - t0) / 60:.1f} min")

research/scripts/teach_gen.py

@@ -0,0 +1,149 @@
+"""Generates teaching data with the system "model + deliberation probes + answer checker" on training problems.
+
+For every problem: N sampled solutions (exact token ids kept) graded against the known answer. Solutions that end
+wrong or are cut off by the budget are read out every --every thinking tokens: the thinking so far is closed at the
+last sentence end (+ </think>) and the model gives its answer greedily (the deliberation module's probe). If the
+read-out answer is correct from some cut to the end of the solution, that is a moment where the model already knew
+the answer and should have stopped: the module then writes the model's own closing answer after that cut (sampled,
+checked) -> a "stop" example.
+
+out record: id, task, sample, tokens, n_tokens, truncated, correct, probes[{cut, ans, ok}], stop{cut, tokens, ok}
+usage: teach_gen.py --inp data/teach/math_pilot.jsonl --out results/teach/pilot.jsonl --n 4 --max-tokens 6144
+"""
+import argparse
+import json
+import os
+import re
+import sys
+import threading
+import time
+import urllib.request
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from grade import grade_math, grade_mmlupro, last_boxed  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+ap = argparse.ArgumentParser()
+ap.add_argument("--inp", nargs="+", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--n", type=int, default=4)
+ap.add_argument("--max-tokens", type=int, default=6144)
+ap.add_argument("--every", type=int, default=1024)
+ap.add_argument("--workers", type=int, default=40)
+ap.add_argument("--limit", type=int, default=0)
+ap.add_argument("--port", type=int, default=8080)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+ap.add_argument("--system-file", default="")
+ap.add_argument("--no-probes", action="store_true", help="skip read-out probes of failed samples (no stop examples)")
+args = ap.parse_args()
+
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+PROBE = {"math": ("The final answer is $\\boxed{", 48, ["}$", "}\n"]), "mmlupro": ("The answer is (", 3, [")"])}
+PROBE_IDS = {k: tok(v[0], add_special_tokens=False)["input_ids"] for k, v in PROBE.items()}
+SAMP = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
+        "samplers": ["top_k", "top_p", "temperature"]}
+SYSTEM = open(args.system_file, encoding="utf-8").read() if args.system_file else ""
+
+
+def post(body):
+    for attempt in range(6):
+        try:
+            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 5:
+                raise
+            time.sleep(5)
+
+
+def correct(task, text, gold):
+    return grade_math(text, gold) if task == "math" else grade_mmlupro(text, gold)
+
+
+def kind(it):
+    return "math" if it["task"] in ("math", "math500", "aime", "gsm8k") else "mmlupro"
+
+
+def probe(ids, th, cut, k, gold):
+    suf, n_new, stops = PROBE[k]
+    r = post({"prompt": ids + th[:cut] + [THINK_END] + PROBE_IDS[k], "n_predict": n_new, "temperature": 0.0,
+              "cache_prompt": False, "stop": stops})
+    c = r["content"].strip()
+    text = suf + c + ("}$" if k == "math" else ").")
+    return {"cut": cut, "ans": c[:60], "ok": bool(correct(k, text, gold))}
+
+
+def run(job):
+    it, s = job
+    k = kind(it)
+    msgs = ([{"role": "system", "content": SYSTEM}] if SYSTEM else []) + [{"role": "user", "content": it["prompt"]}]
+    ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_dict=False)
+    seed = zlib.crc32(it["id"].encode()) % 1000003 * 7 + s + 900001
+    r = post({"prompt": ids, "n_predict": args.max_tokens, "seed": seed, "cache_prompt": False, "return_tokens": True,
+              **SAMP})
+    gen = [t for t in r["tokens"] if t != EOS]
+    text = r["content"]
+    trunc = r.get("stop_type") == "limit"
+    ok = "</think>" in text and correct(k, text.split("</think>", 1)[1], it["answer"])
+    rec = {"id": it["id"], "task": it["task"], "sample": s, "tokens": r["tokens"], "n_tokens": r["tokens_predicted"],
+           "truncated": trunc, "correct": bool(ok), "prompt_len": len(ids)}
+    if not ok and not args.no_probes:
+        th = gen[: gen.index(THINK_END)] if THINK_END in gen else gen
+        probes = []
+        for c in range(args.every, len(th) + 1, args.every):
+            cut = next((j for j in range(c, c - 400, -1) if j <= len(th) and th[j - 1] in ENDS), None)
+            if cut and (not probes or cut > probes[-1]["cut"]):
+                probes.append(probe(ids, th, cut, k, it["answer"]))
+        rec["probes"] = probes
+        # earliest cut from which every later read-out is correct (at least 2 read-outs)
+        start = None
+        for i in range(len(probes) - 1, -1, -1):
+            if probes[i]["ok"]:
+                start = i
+            else:
+                break
+        if start is not None and len(probes) - start >= 2:
+            cut = probes[start]["cut"]
+            r2 = post({"prompt": ids + th[:cut] + [THINK_END], "n_predict": 1536, "seed": seed + 31,
+                       "cache_prompt": False, "return_tokens": True, **SAMP})
+            ans_ok = correct(k, r2["content"], it["answer"]) and r2.get("stop_type") != "limit"
+            rec["stop"] = {"cut": cut, "tokens": th[:cut] + [THINK_END] + r2["tokens"], "ok": bool(ans_ok)}
+    return rec
+
+
+items = []
+for path in args.inp:
+    rows = [json.loads(l) for l in open(path, encoding="utf-8")]
+    items += rows[: args.limit] if args.limit else rows
+done = set()
+if os.path.exists(args.out):
+    done = {(json.loads(l)["id"], json.loads(l)["sample"]) for l in open(args.out, encoding="utf-8")}
+os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+jobs = [(it, s) for it in items for s in range(args.n) if (it["id"], s) not in done]
+print(f"{len(jobs)} jobs", flush=True)
+lock, f_out, t0, cnt = threading.Lock(), open(args.out, "a", encoding="utf-8"), time.time(), [0]
+
+
+def work(job):
+    try:
+        rec = run(job)
+    except Exception as e:
+        print("ERR", job[0]["id"], repr(e)[:200], flush=True)
+        return
+    with lock:
+        f_out.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        f_out.flush()
+        cnt[0] += 1
+        if cnt[0] % 100 == 0:
+            print(f"{cnt[0]}/{len(jobs)}  {(time.time() - t0) / 60:.1f} min", flush=True)
+
+
+with ThreadPoolExecutor(args.workers) as ex:
+    list(ex.map(work, jobs))
+print("done", cnt[0], f"{(time.time() - t0) / 60:.1f} min")

research/scripts/teach_probe.py

@@ -0,0 +1,110 @@
+"""Adds "stop" examples for solutions that ended CORRECT but long: the thinking is read out every --every tokens
+(deliberation probe); from the earliest cut after which every read-out is the correct answer, the model's own
+closing answer is generated and checked. The tail after that cut was re-verification the model did not need.
+Runs on teach_gen.py output in place (adds rec["stop"] where it is shorter than the natural solution).
+usage: teach_probe.py --inp results/teach/pilot.jsonl --items data/teach/math_pilot.jsonl [--min-len 1500]
+"""
+import argparse
+import json
+import os
+import sys
+import time
+import urllib.request
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from grade import grade_math, grade_mmlupro  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+ap = argparse.ArgumentParser()
+ap.add_argument("--inp", required=True)
+ap.add_argument("--items", nargs="+", required=True)
+ap.add_argument("--every", type=int, default=1024)
+ap.add_argument("--min-len", type=int, default=1500)
+ap.add_argument("--workers", type=int, default=48)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+args = ap.parse_args()
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+PROBE = {"math": ("The final answer is $\\boxed{", 48, ["}$", "}\n"]), "mmlupro": ("The answer is (", 3, [")"])}
+PROBE_IDS = {k: tok(v[0], add_special_tokens=False)["input_ids"] for k, v in PROBE.items()}
+SAMP = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "samplers": ["top_k", "top_p", "temperature"]}
+items = {}
+for p in args.items:
+    for l in open(p, encoding="utf-8"):
+        r = json.loads(l)
+        items[r["id"]] = r
+
+
+def post(body):
+    for attempt in range(6):
+        try:
+            req = urllib.request.Request("http://localhost:8080/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 5:
+                raise
+            time.sleep(5)
+
+
+def ok(k, text, gold):
+    return grade_math(text, gold) if k == "math" else grade_mmlupro(text, gold)
+
+
+def work(rec):
+    if not rec["correct"] or rec["n_tokens"] < args.min_len or rec.get("stop") or rec["id"] not in items:
+        return rec
+    it = items[rec["id"]]
+    k = "math" if it["task"] == "math" else "mmlupro"
+    gen = [t for t in rec["tokens"] if t != EOS]
+    if THINK_END not in gen:
+        return rec
+    th = gen[: gen.index(THINK_END)]
+    ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
+                                  add_generation_prompt=True, return_dict=False)
+    probes = []
+    for c in range(args.every, len(th) + 1, args.every):
+        cut = next((j for j in range(c, c - 400, -1) if th[j - 1] in ENDS), None)
+        if not cut or (probes and cut <= probes[-1]["cut"]):
+            continue
+        suf, n_new, stops = PROBE[k]
+        r = post({"prompt": ids + th[:cut] + [THINK_END] + PROBE_IDS[k], "n_predict": n_new, "temperature": 0.0,
+                  "cache_prompt": False, "stop": stops})
+        c_ = r["content"].strip()
+        probes.append({"cut": cut, "ans": c_[:60],
+                       "ok": bool(ok(k, suf + c_ + ("}$" if k == "math" else ")."), it["answer"]))})
+    rec["probes_correct"] = probes
+    start = None
+    for i in range(len(probes) - 1, -1, -1):
+        if probes[i]["ok"]:
+            start = i
+        else:
+            break
+    if start is None:
+        return rec
+    cut = probes[start]["cut"]
+    seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec["sample"] + 77
+    r2 = post({"prompt": ids + th[:cut] + [THINK_END], "n_predict": 1536, "seed": seed, "cache_prompt": False,
+               "return_tokens": True, **SAMP})
+    good = ok(k, r2["content"], it["answer"]) and r2.get("stop_type") != "limit"
+    toks = th[:cut] + [THINK_END] + r2["tokens"]
+    if good and len(toks) < rec["n_tokens"]:
+        rec["stop"] = {"cut": cut, "tokens": toks, "ok": True, "from_correct": True}
+    return rec
+
+
+recs = [json.loads(l) for l in open(args.inp, encoding="utf-8")]
+t0 = time.time()
+with ThreadPoolExecutor(args.workers) as ex:
+    out = list(ex.map(work, recs))
+with open(args.inp, "w", encoding="utf-8") as f:
+    for r in out:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+n = sum(1 for r in out if r.get("stop", {}).get("from_correct"))
+saved = [r["n_tokens"] - len(r["stop"]["tokens"]) for r in out if r.get("stop", {}).get("from_correct")]
+print(f"{n} early-stop examples from correct solutions, mean tokens saved {sum(saved) / max(1, n):.0f}, "
+      f"{(time.time() - t0) / 60:.1f} min")

research/scripts/teach_to_rl.py

@@ -0,0 +1,36 @@
+"""Converts teach_gen.py records into rl_train.py inputs (rollouts with exact token ids + graded rewards), keeping
+only informative groups (some samples right, some wrong) — the problems where the model is unstable and a policy
+update can move pass@1 towards pass@k.
+usage: teach_to_rl.py --inp results/teach/rft1.jsonl ... --out-dir results/rl/rft1
+"""
+import argparse
+import collections
+import json
+import os
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--inp", nargs="+", required=True)
+ap.add_argument("--out-dir", required=True)
+a = ap.parse_args()
+by = collections.defaultdict(list)
+for p in a.inp:
+    for l in open(p, encoding="utf-8"):
+        r = json.loads(l)
+        by[r["id"]].append(r)
+os.makedirs(a.out_dir, exist_ok=True)
+n_inf = n_seq = 0
+with open(os.path.join(a.out_dir, "rollouts.jsonl"), "w", encoding="utf-8") as fr, \
+        open(os.path.join(a.out_dir, "graded.jsonl"), "w", encoding="utf-8") as fg:
+    for pid, g in by.items():
+        ks = {r["sample"]: r for r in g}
+        rs = [r["correct"] for r in ks.values()]
+        if all(rs) or not any(rs):
+            continue
+        n_inf += 1
+        for s, r in ks.items():
+            toks = [t for t in r["tokens"] if t != 124900]
+            fr.write(json.dumps({"id": pid, "sample": s, "tokens": toks,
+                                 "stop_type": "limit" if r["truncated"] else "eos"}) + "\n")
+            fg.write(json.dumps({"id": pid, "sample": s, "correct": bool(r["correct"])}) + "\n")
+            n_seq += 1
+print(f"{len(by)} problems, {n_inf} informative groups, {n_seq} sequences")

research/scripts/tir.py

@@ -0,0 +1,197 @@
+"""Tool-integrated reasoning client: the model may call a `python` tool while it reasons; the module executes the code
+and feeds the output back as a tool turn (the model's native tool format), then the model continues thinking.
+
+The whole multi-turn exchange is flattened into the gen.py record format, so grade.py grades it unchanged:
+completion = "<think>" + all turns (thinking, calls, outputs) + "</think>" + final answer of the last turn.
+The token budget (--max-tokens) counts every generated token over all turns.
+
+usage: tir.py --inp data/quick/mini.jsonl --out results/tir/mini.jsonl [--max-tokens 6144] [--system default]
+"""
+import argparse
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+import urllib.request
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import sandbox  # noqa: E402
+
+EOS, THINK_END = 124900, 124902
+FORCE_PREFIX = "<tool_call>\n<function=python>\n<parameter=code>\n"
+SYSTEMS = {
+    "none": "",
+    "default": "You have a Python interpreter (the `python` tool). Use it whenever a calculation, an enumeration, "
+               "a simulation or a check of your answer can be done by code: write a short program, print the result, "
+               "and use the printed output in your reasoning.",
+    "verify": "You have a Python interpreter (the `python` tool). Plan: think the problem through, then before giving "
+              "the final answer call the `python` tool to compute or check the result by code (brute force, "
+              "enumeration, direct computation or substitution). If the code disagrees with your reasoning, find "
+              "the mistake before answering.",
+}
+TOOL = {"type": "function", "function": {
+    "name": "python",
+    "description": "Execute Python 3 code and return everything it prints (stdout and errors). Available: math, "
+                   "itertools, fractions, collections, sympy, numpy. Always print() the values you need. "
+                   "Time limit 10 seconds.",
+    "parameters": {"type": "object", "properties": {"code": {"type": "string", "description": "Python source"}},
+                   "required": ["code"]}}}
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--inp", nargs="+", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--n", type=int, default=1)
+ap.add_argument("--max-tokens", type=int, default=6144)
+ap.add_argument("--max-calls", type=int, default=8)
+ap.add_argument("--system", default="default")
+ap.add_argument("--force-check", default="", help="if set (e.g. '# Independent check\\n'), a model that ends its "
+                "thinking without running code is made to write a check program starting with this line")
+ap.add_argument("--tasks", default="", help="comma list: only these tasks get the tool (others: plain generation)")
+ap.add_argument("--temperature", type=float, default=0.6)
+ap.add_argument("--top-p", type=float, default=0.95)
+ap.add_argument("--top-k", type=int, default=50)
+ap.add_argument("--workers", type=int, default=24)
+ap.add_argument("--limit", type=int, default=0)
+ap.add_argument("--port", type=int, default=8080)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+ap.add_argument("--seed", type=int, default=0)
+args = ap.parse_args()
+
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+SANDBOX = os.path.join(tempfile.gettempdir(), "lfm3_tir")
+os.makedirs(SANDBOX, exist_ok=True)
+CALL_RE = re.compile(r"<tool_call>\s*<function=python>\s*<parameter=code>\n?(.*?)\n?</parameter>", re.S)
+
+
+def salvage(code):
+    """Small models often drift back into prose inside a code block: keep the longest prefix that is valid Python."""
+    lines = code.split("\n")
+    for n in range(len(lines), 0, -1):
+        try:
+            compile("\n".join(lines[:n]), "<code>", "exec")
+            return "\n".join(lines[:n]), n < len(lines)
+        except SyntaxError:
+            continue
+    return code, False
+
+
+def run_code(code, timeout=10):
+    code, cut = salvage(code)
+    note = "(note: the end of the code was not valid Python and was skipped)\n" if cut else ""
+    try:
+        p = sandbox.run([sys.executable, "-I", "-c", code], capture_output=True, text=True, timeout=timeout,
+                           cwd=SANDBOX, encoding="utf-8", errors="replace",
+                           env={**os.environ, "PYTHONIOENCODING": "utf-8"})
+        out = p.stdout + (("\n" + p.stderr.strip().splitlines()[-1]) if p.returncode and p.stderr.strip() else "")
+    except subprocess.TimeoutExpired:
+        out = "TimeoutError: the code ran longer than 10 seconds"
+    out = note + (out.strip() or "(no output; use print())")
+    return out if len(out) <= 1500 else out[:1500] + "\n... (output truncated)"
+
+
+def post(body):
+    for attempt in range(5):
+        try:
+            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 4:
+                raise
+            time.sleep(5)
+
+
+def solve(it, s):
+    use_tool = not args.tasks or it["task"] in args.tasks.split(",")
+    msgs = ([{"role": "system", "content": SYSTEMS[args.system]}] if use_tool and SYSTEMS[args.system] else []) + \
+        [{"role": "user", "content": it["prompt"]}]
+    ids = tok.apply_chat_template(msgs, tools=[TOOL] if use_tool else None, tokenize=True,
+                                  add_generation_prompt=True, return_dict=False)
+    seed = args.seed * 1000003 + zlib.crc32(it["id"].encode()) % 1000003 * 7 + s
+    used, calls, parts, stop_type = 0, [], [], None
+    while True:
+        res = post({"prompt": ids, "n_predict": args.max_tokens - used, "temperature": args.temperature,
+                    "top_p": args.top_p, "top_k": args.top_k, "min_p": 0.0, "repeat_penalty": 1.0,
+                    "seed": seed + 101 * len(calls), "cache_prompt": False, "return_tokens": True,
+                    "samplers": ["top_k", "top_p", "temperature"]})
+        gen = [t for t in res["tokens"] if t != EOS]
+        used += res["tokens_predicted"]
+        text = res["content"]
+        stop_type = res.get("stop_type")
+        m = CALL_RE.search(text.split("</think>", 1)[1]) if "</think>" in text else None
+        if use_tool and not m and args.force_check and not calls and THINK_END in gen \
+                and used < args.max_tokens - 1024:
+            # the model finished thinking without running anything: make its next move a code check of the answer
+            k = gen.index(THINK_END) + 1
+            used -= len(gen) - k  # the answer text after </think> is replaced by the check
+            pre = tok(FORCE_PREFIX + args.force_check, add_special_tokens=False)["input_ids"]
+            r2 = post({"prompt": ids + gen[:k] + pre, "n_predict": min(1536, args.max_tokens - used),
+                       "temperature": args.temperature, "top_p": args.top_p, "top_k": args.top_k, "min_p": 0.0,
+                       "seed": seed + 7, "cache_prompt": False, "stop": ["</parameter>"],
+                       "samplers": ["top_k", "top_p", "temperature"]})
+            used += r2["tokens_predicted"]
+            code = args.force_check + r2["content"]
+            out = run_code(code)
+            calls.append({"code": code, "out": out, "forced": True})
+            parts.append(tok.decode(gen[:k]).replace("</think>", "\n") + FORCE_PREFIX + code +
+                         "\n</parameter>\n</function>\n</tool_call>\n<tool_response>\n" + out + "\n</tool_response>\n")
+            ids = ids + gen[:k] + pre + tok(r2["content"].rstrip("\n") + "\n</parameter>\n</function>\n</tool_call>"
+                                            "<|im_end|>\n<|im_start|>tool\n<tool_response>\n" + out +
+                                            "\n</tool_response><|im_end|>\n<|im_start|>assistant\n<think>",
+                                            add_special_tokens=False)["input_ids"]
+            continue
+        if not (use_tool and m and stop_type != "limit" and len(calls) < args.max_calls and used < args.max_tokens):
+            parts.append(text)
+            break
+        out = run_code(m.group(1))
+        calls.append({"code": m.group(1), "out": out})
+        turn = text.split("</tool_call>", 1)[0] + "</tool_call>"
+        parts.append(turn.replace("</think>", "\n", 1) + "\n<tool_response>\n" + out + "\n</tool_response>\n")
+        ids = ids + gen + tok("<|im_end|>\n<|im_start|>tool\n<tool_response>\n" + out +
+                              "\n</tool_response><|im_end|>\n<|im_start|>assistant\n<think>",
+                              add_special_tokens=False)["input_ids"]
+    last = parts[-1]
+    if "</think>" in last:
+        think, final = last.split("</think>", 1)
+        completion = "<think>" + "".join(parts[:-1]) + think + "</think>" + final
+    else:
+        completion = "<think>" + "".join(parts)
+    return {"id": it["id"], "task": it["task"], "sample": s, "completion": completion, "n_tokens": used,
+            "truncated": stop_type == "limit", "stop_type": stop_type, "n_calls": len(calls), "calls": calls}
+
+
+items = []
+for path in args.inp:
+    with open(path, encoding="utf-8") as f:
+        rows = [json.loads(l) for l in f]
+    items += rows[: args.limit] if args.limit else rows
+done = set()
+if os.path.exists(args.out):
+    with open(args.out, encoding="utf-8") as f:
+        done = {(json.loads(l)["id"], json.loads(l)["sample"]) for l in f}
+os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+jobs = [(it, s) for it in items for s in range(args.n) if (it["id"], s) not in done]
+print(f"{len(jobs)} jobs ({len(done)} already done)", flush=True)
+lock, out_f, t0, cnt = threading.Lock(), open(args.out, "a", encoding="utf-8"), time.time(), [0]
+
+
+def run(job):
+    rec = solve(*job)
+    with lock:
+        out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        out_f.flush()
+        cnt[0] += 1
+        if cnt[0] % 20 == 0 or cnt[0] == len(jobs):
+            print(f"{cnt[0]}/{len(jobs)}  {(time.time() - t0) / 60:.1f} min", flush=True)
+
+
+with ThreadPoolExecutor(args.workers) as ex:
+    list(ex.map(run, jobs))

research/scripts/vis_autofix.py

@@ -0,0 +1,67 @@
+"""Deterministic fixes of visual answers before rendering (no model call), measured on recorded runs.
+
+  importmap  the page loads three.js as ES modules (imports from 'three', 'three/addons/...' or from the unpkg URL,
+             whose addons import the bare name 'three') but has no import map: insert the import map of the task
+             at the start of <head>.
+
+usage: vis_autofix.py --src results/vis_rep/vis.jsonl --out results/vis_rep_fix/vis.jsonl
+Records whose answer is changed get "autofix": [...]; the rest are copied unchanged (exactly paired).
+"""
+import argparse
+import json
+import os
+import re
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from render import extract  # noqa: E402
+
+IMPORT_MAP = ('<script type="importmap">{"imports": {"three": "https://unpkg.com/three@0.160.0/build/three.module.js", '
+              '"three/addons/": "https://unpkg.com/three@0.160.0/examples/jsm/"}}</script>')
+USES_THREE_MODULE = re.compile(r"""from\s+['"](three|three/addons/[^'"]+|https://unpkg\.com/three@[^'"]+)['"]""")
+
+
+def fix_importmap(src):
+    if 'type="importmap"' in src or "type='importmap'" in src or not USES_THREE_MODULE.search(src):
+        return src
+    m = re.search(r"<head[^>]*>", src, re.I)
+    if m:
+        return src[:m.end()] + "\n" + IMPORT_MAP + src[m.end():]
+    m = re.search(r"<script", src, re.I)
+    return src[:m.start()] + IMPORT_MAP + "\n" + src[m.start():] if m else src
+
+
+def autofix(text, kind):
+    """Returns (new answer text, list of applied fixes)."""
+    if kind != "html" or text is None:
+        return text, []
+    src = extract(text, kind)
+    if not src:
+        return text, []
+    new = fix_importmap(src)
+    if new == src:
+        return text, []
+    return text.replace(src, new), ["importmap"]
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+    a = ap.parse_args()
+    tasks = {json.loads(l)["id"]: json.loads(l) for l in open(a.tasks, encoding="utf-8")}
+    os.makedirs(os.path.dirname(a.out), exist_ok=True)
+    n = 0
+    with open(a.out, "w", encoding="utf-8") as f:
+        for l in open(a.src, encoding="utf-8"):
+            r = json.loads(l)
+            c = r["completion"]
+            if "</think>" in c:
+                head, fin = c.split("</think>", 1)
+                new, fixes = autofix(fin, tasks[r["id"]]["kind"])
+                if fixes:
+                    r = {**r, "completion": head + "</think>" + new, "autofix": fixes}
+                    n += 1
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    print(f"{n} answers changed -> {a.out}")

research/scripts/visual_eval.py

@@ -0,0 +1,76 @@
+"""Renders visual answers and scores them with objective checks (scripts/render.py).
+A result is "clean" when it parses, raises no JavaScript/console errors, loads every resource, is not blank and, for
+animation tasks, actually moves. Screenshots go to results/<tag>/png/ for side-by-side review.
+usage: visual_eval.py --tag vis_base                 (records results/<tag>/vis.jsonl from gen.py / posthoc)
+       visual_eval.py --tag vis_claude --dir data/visual/claude   (reference files <id>.svg / <id>.html)
+"""
+import argparse
+import collections
+import glob
+import json
+import os
+import re
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from render import Renderer, extract  # noqa: E402
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--tag", required=True)
+ap.add_argument("--dir")
+ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+a = ap.parse_args()
+tasks = {json.loads(l)["id"]: json.loads(l) for l in open(a.tasks, encoding="utf-8")}
+
+
+def clean(rep, task):
+    return (not rep.get("parse_error") and not rep["errors"] and not rep["failed"] and not rep["blank"]
+            and (rep["animated"] or not task["animated"]))
+
+
+out_dir = os.path.join("results", a.tag)
+os.makedirs(os.path.join(out_dir, "png"), exist_ok=True)
+sources = {}
+if a.dir:
+    for p in glob.glob(os.path.join(a.dir, "*.*")):
+        tid = "visual/" + os.path.splitext(os.path.basename(p))[0]
+        if tid in tasks:
+            sources[(tid, 0)] = open(p, encoding="utf-8").read()
+else:
+    for p in glob.glob(os.path.join(out_dir, "*.jsonl")):
+        for l in open(p, encoding="utf-8"):
+            r = json.loads(l)
+            if r["id"] in tasks:
+                fin = r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else None
+                sources[(r["id"], r.get("sample", 0))] = extract(fin, tasks[r["id"]]["kind"])
+
+rend = Renderer()
+rows = []
+for (tid, s), src in sorted(sources.items()):
+    task = tasks[tid]
+    png = os.path.join(out_dir, "png", f"{tid.split('/')[1]}_{s}.png")
+    if not src:
+        rep = {"errors": ["no code block"], "failed": [], "parse_error": "no code", "blank": True, "animated": False,
+               "coverage": 0.0, "overflow": False}
+    else:
+        rep = rend.render(task["kind"], src, png)
+    rep.update(id=tid, sample=s, cat=task["cat"], clean=clean(rep, task), chars=len(src or ""))
+    rows.append(rep)
+rend.close()
+with open(os.path.join(out_dir, "visual.json"), "w", encoding="utf-8") as f:
+    json.dump(rows, f, ensure_ascii=False, indent=1)
+by = collections.defaultdict(list)
+for r in rows:
+    by[r["cat"]].append(r)
+print(f"{a.tag}: {len(rows)} answers")
+for cat in ("svg", "html", "canvas", "three"):
+    v = by.get(cat, [])
+    if v:
+        print(f"  {cat:7s} n={len(v):3d} clean {100 * sum(r['clean'] for r in v) / len(v):5.1f}%  "
+              f"js-errors {sum(bool(r['errors']) for r in v):3d}  failed-loads {sum(bool(r['failed']) for r in v):2d}  "
+              f"blank {sum(r['blank'] for r in v):2d}  no-code {sum(r.get('parse_error') == 'no code' for r in v):2d}")
+print(f"  ALL     clean {100 * sum(r['clean'] for r in rows) / max(len(rows), 1):5.1f}%")
+for r in rows:
+    if not r["clean"]:
+        print("   ", r["id"], r["sample"], (r.get("parse_error") or "") , r["errors"][:2], r["failed"][:1],
+              "blank" if r["blank"] else "", "" if r["animated"] else "static")

research/scripts/visual_eyes.py

@@ -0,0 +1,209 @@
+"""Eyes for a text-only model: render -> look -> fix.
+
+The answer is rendered in headless Chrome (scripts/render.py). A vision-language model (LFM3-VL-2B: the vision
+sibling of LFM3 on the same language backbone, served by a second llama-server with --mmproj) describes the
+screenshot without being told the task. LFM3 then compares the description with the task and returns an improved
+file or NO CHANGES. An improvement is kept only if it renders with no more technical problems than before
+(same weights as visual_repair.py), so the module cannot break a working answer.
+
+--no-eyes is the control: the same revision turn without the description ("look at your file again"), to separate
+the value of seeing the result from the value of one more revision.
+Records whose answer is unchanged are copied, so comparisons with the source run are exactly paired.
+usage: visual_eyes.py --src results/vis_rep_fix/vis.jsonl --out results/vis_eyes/vis.jsonl [--no-eyes]
+"""
+import argparse
+import base64
+import json
+import os
+import re
+import sys
+import threading
+import time
+import urllib.request
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from render import Renderer, extract  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+ap = argparse.ArgumentParser()
+ap.add_argument("--src", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+ap.add_argument("--rounds", type=int, default=1)
+ap.add_argument("--think", type=int, default=6144)
+ap.add_argument("--reserve", type=int, default=4096)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+ap.add_argument("--port", type=int, default=8080)
+ap.add_argument("--vlm-port", type=int, default=8091)
+ap.add_argument("--workers", type=int, default=8)
+ap.add_argument("--model-tag", default="lfm3-orig-q5km", help="the generation server must serve this model file")
+ap.add_argument("--no-eyes", action="store_true", help="control: revision turn without the description")
+args = ap.parse_args()
+
+DESCRIBE = ("Describe this screenshot precisely for someone who cannot see it. List every visible element: shapes "
+            "and objects, their colours, sizes and positions (top, bottom, left, right, centre), all visible text "
+            "(quote it), and anything that looks broken, cut off, overlapping or empty. Describe only what is "
+            "actually visible; do not guess what it was meant to show.")
+EYES = ("I opened your {what} in Chrome (window 1024x768) and asked someone who cannot see your code to describe "
+        "the screen{frame}. Their description:\n\n\"\"\"\n{desc}\n\"\"\"\n\nCompare it with the task. If anything "
+        "requested is missing, wrong, misplaced, unreadable or looks unpolished, improve the file and give the "
+        "complete corrected file in a single ```{kind} code block. If the result already matches the task well, "
+        "answer only NO CHANGES.")
+SELF = ("Look at your {what} again and compare it with the task. If anything requested is missing, wrong, misplaced, "
+        "unreadable or looks unpolished, improve the file and give the complete corrected file in a single "
+        "```{kind} code block. If the result already matches the task well, answer only NO CHANGES.")
+
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
+            "samplers": ["top_k", "top_p", "temperature"]}
+tasks = {json.loads(l)["id"]: json.loads(l) for l in open(args.tasks, encoding="utf-8")}
+SHOTS = os.path.join(os.path.dirname(args.out), "shots")
+os.makedirs(SHOTS, exist_ok=True)
+
+
+def get_json(url, body=None, timeout=3600):
+    data = json.dumps(body).encode() if body is not None else None
+    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
+    return json.loads(urllib.request.urlopen(req, timeout=timeout).read())
+
+
+def wait_server():
+    """The night chain restarts the shared server with other models; generate only on the right one."""
+    while True:
+        try:
+            p = get_json(f"http://localhost:{args.port}/props", timeout=10)
+            if args.model_tag in p.get("model_path", ""):
+                return p["default_generation_settings"]["n_ctx"]
+        except Exception:
+            pass
+        time.sleep(15)
+
+
+def post(body):
+    for attempt in range(40):
+        n_ctx = wait_server()
+        try:
+            room = n_ctx - len(body["prompt"]) - 16  # the slot size changes when the chain restarts the server
+            return get_json(f"http://localhost:{args.port}/completion", {**body, "n_predict": min(body["n_predict"], room)})
+        except Exception:
+            if attempt == 39:
+                raise
+            time.sleep(10)
+
+
+def generate(msgs, seed):
+    ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_dict=False)
+    r = post({"prompt": ids, "n_predict": args.think + args.reserve, "seed": seed, "cache_prompt": False,
+              "return_tokens": True, **SAMPLING})
+    gen = [t for t in r["tokens"] if t != EOS]
+    if THINK_END in gen:
+        return r["content"].split("</think>", 1)[1], r["tokens_predicted"]
+    lim = min(args.think, len(gen))
+    cut = next((j for j in range(lim, 0, -1) if gen[j - 1] in ENDS), lim)
+    r2 = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.reserve, "seed": seed + 17,
+               "cache_prompt": False, **SAMPLING})
+    return r2["content"], cut + r2["tokens_predicted"]
+
+
+def look(png):
+    b64 = base64.b64encode(open(png, "rb").read()).decode()
+    body = {"messages": [{"role": "user", "content": [
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64," + b64}},
+        {"type": "text", "text": DESCRIBE}]}],
+        "temperature": 0.1, "top_k": 50, "top_p": 1.0, "min_p": 0.0, "repeat_penalty": 1.1, "max_tokens": 1500,
+        "seed": 1, "chat_template_kwargs": {"enable_thinking": False}}  # thinking: 6x slower, the description is enough
+    for attempt in range(5):
+        try:
+            msg = get_json(f"http://localhost:{args.vlm_port}/v1/chat/completions", body)["choices"][0]["message"]
+            text = re.sub(r"<think>.*?</think>", "", msg.get("content") or "", flags=re.S).strip()
+            return text or "(no description)"
+        except Exception:
+            if attempt == 4:
+                raise
+            time.sleep(10)
+
+
+def badness(rep, task):
+    """Technical problems with the weights of visual_repair.py (0 = renders cleanly)."""
+    if rep is None:
+        return 10
+    return (4 * bool(rep.get("parse_error")) + 3 * bool(rep["errors"]) + 3 * bool(rep["failed"]) + 3 * rep["blank"]
+            + (task["animated"] and not rep["animated"]) + bool(rep.get("overflow")))
+
+
+local = threading.local()
+
+
+def renderer():
+    if not hasattr(local, "r"):
+        local.r = Renderer()
+    return local.r
+
+
+def process(rec):
+    task, s = tasks[rec["id"]], rec.get("sample", 0)
+    name = f"{rec['id'].split('/')[-1]}_{s}"
+    fin = rec["completion"].split("</think>", 1)[1] if "</think>" in rec["completion"] else None
+    src = extract(fin, task["kind"])
+    if not src:
+        return {**rec, "eyes": "no source"}
+    png = os.path.join(SHOTS, f"{name}_r0.png")
+    bad = badness(renderer().render(task["kind"], src, png), task)
+    what = "SVG" if task["kind"] == "svg" else "page"
+    frame = " (one frame of the animation)" if task["animated"] else ""
+    msgs = [{"role": "user", "content": task["prompt"]}]
+    cur, used, log = fin, 0, []
+    seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + s + 9000
+    for rnd in range(args.rounds):
+        if args.no_eyes:
+            ask = SELF.format(what=what, kind=task["kind"])
+        else:
+            desc = look(png)
+            log.append({"round": rnd + 1, "description": desc})
+            ask = EYES.format(what=what, kind=task["kind"], frame=frame, desc=desc)
+        msgs = msgs + [{"role": "assistant", "content": cur}, {"role": "user", "content": ask}]
+        out, n = generate(msgs, seed + rnd)
+        used += n
+        new = extract(out, task["kind"])
+        if not new or new.strip() == extract(cur, task["kind"]).strip():
+            log.append({"round": rnd + 1, "result": "no changes"})
+            break
+        png2 = os.path.join(SHOTS, f"{name}_r{rnd + 1}.png")
+        b = badness(renderer().render(task["kind"], new, png2), task)
+        if b > bad:
+            log.append({"round": rnd + 1, "result": f"rejected: problems {bad} -> {b}"})
+            break
+        log.append({"round": rnd + 1, "result": f"kept: problems {bad} -> {b}"})
+        cur, bad, png = out, b, png2
+    changed = cur is not fin
+    print(f"{rec['id']} s{s}: {log[-1].get('result', '?') if log else '-'}", flush=True)
+    r = {**rec, "eyes_log": log, "eyes_changed": changed, "n_tokens": rec["n_tokens"] + used}
+    if changed:
+        r["completion"] = "<think>(eyes)</think>" + cur
+    return r
+
+
+def safe(rec):
+    """One failing record (browser crash, server error) must not lose the whole run: it is copied unchanged."""
+    try:
+        return process(rec)
+    except Exception:
+        import traceback
+        print(f"{rec['id']} s{rec.get('sample', 0)}: ERROR " + traceback.format_exc(), flush=True)
+        return {**rec, "eyes_log": [{"result": "error"}], "eyes_changed": False}
+
+
+recs = [json.loads(l) for l in open(args.src, encoding="utf-8")]
+t0 = time.time()
+with ThreadPoolExecutor(args.workers) as ex:
+    out = list(ex.map(safe, recs))
+with open(args.out, "w", encoding="utf-8") as f:
+    for r in out:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+print(f"{len(out)} records, {sum(r.get('eyes_changed', False) for r in out)} changed, "
+      f"{(time.time() - t0) / 60:.1f} min")

research/scripts/visual_loop.py

@@ -0,0 +1,96 @@
+"""Shared render-check-repair logic for visual answers, used by lfm3_deliberate.solve_visual.
+(The same rules as scripts/visual_repair.py, which produced the measured numbers.)
+
+problems(rep, animated, src, kind, hints) -> (weighted badness, feedback lines) from a render.py report;
+feedback(kind, lines) -> the repair turn text; fix_importmap from vis_autofix (deterministic, no model call).
+"""
+import re
+
+from vis_autofix import IMPORT_MAP, fix_importmap  # noqa: F401  (re-exported for callers)
+
+# browser error -> the fix a web developer would apply (general knowledge, never a solution of a particular task)
+HINTS = [
+    (r'Failed to resolve module specifier "three', "A module imports the bare name 'three' (three.js addons do this "
+     f"too). Put the import map in <head>, before any module script: {IMPORT_MAP}, then write "
+     "import * as THREE from 'three'; and import addons like import { OrbitControls } from "
+     "'three/addons/controls/OrbitControls.js';"),
+    (r"THREE is not defined", "three.js is not a global variable. Add the import map "
+     f"{IMPORT_MAP} and, inside <script type=\"module\">, write import * as THREE from 'three';"),
+    (r"THREE\.(\w+) is not a constructor", "THREE.{0} does not exist in three.js r160. Addons are imported "
+     "separately, e.g. import { OrbitControls } from 'three/addons/controls/OrbitControls.js'; then "
+     "new OrbitControls(camera, renderer.domElement). Otherwise check the class name."),
+    (r"Assignment to constant variable", "A variable declared with const is assigned again: declare variables "
+     "that change with let."),
+    (r"ReferenceError: (\w+) is not defined", "'{0}' is used where it is not declared: declare it with const/let "
+     "at the top level, before the functions that use it run."),
+    (r"Cannot read properties of (?:undefined|null) \(reading '(\w+)'\)", "Something is undefined or null when "
+     "its property '{0}' is read: an element id that does not exist, an object used before it is created, or an "
+     "array index out of range."),
+    (r"(\S+) is not a function", "{0} is not a function: a property is called like a method or a name is "
+     "misspelled (canvas.width/height are numbers; drawing methods are ctx.fillRect, ctx.arc, ctx.beginPath, "
+     "ctx.fill, ctx.stroke; angles in Math are radians, there is no Math.radians)."),
+]
+BLANK_THREE = ("Add lights for MeshStandardMaterial (AmbientLight and DirectionalLight), put the camera away from "
+               "the objects with camera.lookAt(0, 0, 0) and render every frame with "
+               "renderer.setAnimationLoop(() => renderer.render(scene, camera)).")
+BLANK_CANVAS = ("Set canvas.width = innerWidth and canvas.height = innerHeight, draw in every frame and use a fill "
+                "color that differs from the background.")
+STATIC_HINT = ("Compute positions from the elapsed time inside a requestAnimationFrame loop (or "
+               "renderer.setAnimationLoop) and redraw every frame.")
+SVG_HINT = "Close every tag, quote every attribute value, write & as &amp; and keep one root <svg> element."
+
+
+def hints_for(rep, animated, src):
+    out = []
+    for e in rep.get("errors", []):
+        for pat, text in HINTS:
+            m = re.search(pat, e)
+            if m:
+                h = text.replace("{0}", m.group(1)) if m.groups() else text
+                if h not in out:
+                    out.append(h)
+                break
+    if rep.get("parse_error"):
+        out.append(SVG_HINT)
+    if rep.get("blank") and not rep.get("parse_error") and not rep.get("errors"):
+        out.append(BLANK_THREE if "three" in (src or "") else BLANK_CANVAS if "<canvas" in (src or "") else
+                   "Make sure the content is visible: real sizes, colors that differ from the background, nothing "
+                   "hidden by CSS.")
+    if animated and not rep.get("animated") and not rep.get("blank"):
+        out.append(STATIC_HINT)
+    return out
+
+
+def problems(rep, animated, src, kind, hints=False):
+    """(weighted badness, feedback lines); 0 = renders cleanly."""
+    if not src:
+        return 10, [f"Your answer does not contain the complete file in a ```{kind} code block."]
+    bad, lines = 0, []
+    if rep.get("parse_error"):
+        bad += 4
+        lines.append(f"The SVG is not valid XML: {rep['parse_error']}")
+    if rep["errors"]:
+        bad += 3
+        lines += [f"JavaScript error: {e}" for e in rep["errors"][:4]]
+    if rep["failed"]:
+        bad += 3
+        lines += [f"This resource failed to load: {u}" for u in rep["failed"][:3]]
+    if rep["blank"]:
+        bad += 3
+        lines.append("Nothing is visible: the rendered page is blank (a single colour).")
+    if animated and not rep["animated"]:
+        bad += 1
+        lines.append("The animation does not move: two screenshots taken 0.7 seconds apart are identical.")
+    if rep.get("overflow"):
+        bad += 1
+        lines.append("The page is wider than the 1024 px window (there is a horizontal scrollbar).")
+    if hints and bad:
+        lines += ["Hint: " + h for h in hints_for(rep, animated, src)]
+    return bad, lines
+
+
+def feedback(kind, lines):
+    what = "SVG" if kind == "svg" else "page"
+    return (f"I opened your {what} in Chrome (window 1024x768) and found these problems:\n"
+            + "\n".join("- " + x for x in lines)
+            + f"\n\nFix them and give the complete corrected file in a single ```{kind} code block.")

research/scripts/visual_repair.py

@@ -0,0 +1,210 @@
+"""Render-check-repair module for visual answers (SVG / HTML / canvas / three.js).
+
+The answer is rendered in headless Chrome (scripts/render.py). If the browser reports problems (invalid SVG,
+JavaScript errors, resources that fail to load, a blank page, an animation that does not move, horizontal overflow),
+the model gets them as precise feedback in a new turn and returns a corrected file. A correction is kept only if it
+has strictly fewer problems (weighted) than the previous best. Up to --rounds repairs.
+Records without problems are copied unchanged, so the comparison with the source run is exactly paired.
+usage: visual_repair.py --src results/vis_base_dm/vis.jsonl --out results/vis_rep/vis.jsonl
+"""
+import argparse
+import json
+import os
+import re
+import sys
+import time
+import urllib.request
+import zlib
+
+from transformers import AutoTokenizer
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from render import Renderer, extract  # noqa: E402
+
+THINK_END, EOS = 124902, 124900
+ap = argparse.ArgumentParser()
+ap.add_argument("--src", required=True)
+ap.add_argument("--out", required=True)
+ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+ap.add_argument("--rounds", type=int, default=2)
+ap.add_argument("--think", type=int, default=6144)
+ap.add_argument("--reserve", type=int, default=4096)
+ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
+ap.add_argument("--port", type=int, default=8080)
+ap.add_argument("--hints", action="store_true", help="add a known fix to each recognised browser error")
+args = ap.parse_args()
+
+IMPORT_MAP = ('<script type="importmap">{"imports": {"three": "https://unpkg.com/three@0.160.0/build/three.module.js", '
+              '"three/addons/": "https://unpkg.com/three@0.160.0/examples/jsm/"}}</script>')
+# browser error -> the fix a web developer would apply (general knowledge, never a solution of a particular task)
+HINTS = [
+    (r'Failed to resolve module specifier "three', "A module imports the bare name 'three' (three.js addons do this "
+     f"too). Put the import map in <head>, before any module script: {IMPORT_MAP}, then write "
+     "import * as THREE from 'three'; and import addons like import { OrbitControls } from "
+     "'three/addons/controls/OrbitControls.js';"),
+    (r"THREE is not defined", "three.js is not a global variable. Add the import map "
+     f"{IMPORT_MAP} and, inside <script type=\"module\">, write import * as THREE from 'three';"),
+    (r"THREE\.(\w+) is not a constructor", "THREE.{0} does not exist in three.js r160. Addons are imported "
+     "separately, e.g. import { OrbitControls } from 'three/addons/controls/OrbitControls.js'; then "
+     "new OrbitControls(camera, renderer.domElement). Otherwise check the class name."),
+    (r"Assignment to constant variable", "A variable declared with const is assigned again: declare variables "
+     "that change with let."),
+    (r"ReferenceError: (\w+) is not defined", "'{0}' is used where it is not declared: declare it with const/let "
+     "at the top level, before the functions that use it run."),
+    (r"Cannot read properties of (?:undefined|null) \(reading '(\w+)'\)", "Something is undefined or null when "
+     "its property '{0}' is read: an element id that does not exist, an object used before it is created, or an "
+     "array index out of range."),
+    (r"(\S+) is not a function", "{0} is not a function: a property is called like a method or a name is "
+     "misspelled (canvas.width/height are numbers; drawing methods are ctx.fillRect, ctx.arc, ctx.beginPath, "
+     "ctx.fill, ctx.stroke; angles in Math are radians, there is no Math.radians)."),
+]
+BLANK_THREE = ("Add lights for MeshStandardMaterial (AmbientLight and DirectionalLight), put the camera away from "
+               "the objects with camera.lookAt(0, 0, 0) and render every frame with "
+               "renderer.setAnimationLoop(() => renderer.render(scene, camera)).")
+BLANK_CANVAS = ("Set canvas.width = innerWidth and canvas.height = innerHeight, draw in every frame and use a fill "
+                "color that differs from the background.")
+STATIC_HINT = ("Compute positions from the elapsed time inside a requestAnimationFrame loop (or "
+               "renderer.setAnimationLoop) and redraw every frame.")
+SVG_HINT = "Close every tag, quote every attribute value, write & as &amp; and keep one root <svg> element."
+
+
+def hints_for(rep, task, src):
+    out = []
+    for e in rep.get("errors", []):
+        for pat, text in HINTS:
+            m = re.search(pat, e)
+            if m:
+                h = text.replace("{0}", m.group(1)) if m.groups() else text
+                if h not in out:
+                    out.append(h)
+                break
+    if rep.get("parse_error"):
+        out.append(SVG_HINT)
+    if rep.get("blank") and not rep.get("parse_error") and not rep.get("errors"):
+        out.append(BLANK_THREE if "three" in (src or "") else BLANK_CANVAS if "<canvas" in (src or "") else
+                   "Make sure the content is visible: real sizes, colors that differ from the background, nothing "
+                   "hidden by CSS.")
+    if task["animated"] and not rep.get("animated") and not rep.get("blank"):
+        out.append(STATIC_HINT)
+    return out
+tok = AutoTokenizer.from_pretrained(args.tokenizer)
+ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
+SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
+            "samplers": ["top_k", "top_p", "temperature"]}
+tasks = {json.loads(l)["id"]: json.loads(l) for l in open(args.tasks, encoding="utf-8")}
+
+
+def post(body):
+    for attempt in range(5):
+        try:
+            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
+        except Exception:
+            if attempt == 4:
+                raise
+            time.sleep(5)
+
+
+def generate(msgs, seed):
+    ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_dict=False)
+    r = post({"prompt": ids, "n_predict": args.think + args.reserve, "seed": seed, "cache_prompt": False,
+              "return_tokens": True, **SAMPLING})
+    gen = [t for t in r["tokens"] if t != EOS]
+    if THINK_END in gen:
+        return r["content"].split("</think>", 1)[1], r["tokens_predicted"]
+    lim = min(args.think, len(gen))
+    cut = next((j for j in range(lim, 0, -1) if gen[j - 1] in ENDS), lim)
+    r2 = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.reserve, "seed": seed + 17,
+               "cache_prompt": False, **SAMPLING})
+    return r2["content"], cut + r2["tokens_predicted"]
+
+
+def problems(rep, task, src):
+    """(weighted badness, feedback lines)."""
+    if not src:
+        return 10, [f"Your answer does not contain the complete file in a ```{task['kind']} code block."]
+    bad, lines = 0, []
+    if rep.get("parse_error"):
+        bad += 4
+        lines.append(f"The SVG is not valid XML: {rep['parse_error']}")
+    if rep["errors"]:
+        bad += 3
+        lines += [f"JavaScript error: {e}" for e in rep["errors"][:4]]
+    if rep["failed"]:
+        bad += 3
+        lines += [f"This resource failed to load: {u}" for u in rep["failed"][:3]]
+    if rep["blank"]:
+        bad += 3
+        lines.append("Nothing is visible: the rendered page is blank (a single colour).")
+    if task["animated"] and not rep["animated"]:
+        bad += 1
+        lines.append("The animation does not move: two screenshots taken 0.7 seconds apart are identical.")
+    if rep.get("overflow"):
+        bad += 1
+        lines.append("The page is wider than the 1024 px window (there is a horizontal scrollbar).")
+    if args.hints and bad:
+        lines += ["Hint: " + h for h in hints_for(rep, task, src)]
+    return bad, lines
+
+
+FEEDBACK = ("I opened your {what} in Chrome (window 1024x768) and found these problems:\n{lines}\n\nFix them and give "
+            "the complete corrected file in a single ```{kind} code block.")
+
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+local = threading.local()
+
+
+def renderer():
+    if not hasattr(local, "r"):
+        local.r = Renderer()
+    return local.r
+
+
+def process(rec):
+    task = tasks[rec["id"]]
+    fin = rec["completion"].split("</think>", 1)[1] if "</think>" in rec["completion"] else None
+    src = extract(fin, task["kind"])
+    png = os.path.join(os.path.dirname(args.out), f"_probe_{threading.get_ident()}.png")
+    rep = renderer().render(task["kind"], src, png) if src else None
+    bad, lines = problems(rep, task, src)
+    if bad == 0:
+        return {**rec, "vis_bad": 0}
+    msgs = [{"role": "user", "content": task["prompt"]}]
+    best = (bad, fin, 0)
+    cur, used = fin or "", 0
+    seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec.get("sample", 0) + 7000
+    for rnd in range(args.rounds):
+        what = "SVG" if task["kind"] == "svg" else "page"
+        msgs = msgs + [{"role": "assistant", "content": cur or "(no answer)"},
+                       {"role": "user", "content": FEEDBACK.format(what=what, kind=task["kind"],
+                                                                 lines="\n".join("- " + x for x in lines))}]
+        cur, n = generate(msgs, seed + rnd)
+        used += n
+        src = extract(cur, task["kind"])
+        rep = renderer().render(task["kind"], src, png) if src else None
+        b, lines = problems(rep, task, src)
+        if b < best[0]:
+            best = (b, cur, rnd + 1)
+        if b == 0:
+            break
+    r = {**rec, "vis_bad_first": bad, "vis_bad": best[0], "vis_kept": best[2], "n_tokens": rec["n_tokens"] + used}
+    if best[2]:
+        r["completion"] = "<think>(render-repaired)</think>" + best[1]
+    print(f"{rec['id']} s{rec.get('sample', 0)}: badness {bad} -> {best[0]} (round {best[2]})", flush=True)
+    return r
+
+
+recs = [json.loads(l) for l in open(args.src, encoding="utf-8")]
+t0 = time.time()
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+with ThreadPoolExecutor(8) as ex:
+    out = list(ex.map(process, recs))
+os.makedirs(os.path.dirname(args.out), exist_ok=True)
+with open(args.out, "w", encoding="utf-8") as f:
+    for r in out:
+        f.write(json.dumps(r, ensure_ascii=False) + "\n")
+print(f"{len(out)} records, {sum('vis_bad_first' in r for r in out)} had problems, "
+      f"{sum(r.get('vis_bad', 0) == 0 and 'vis_bad_first' in r for r in out)} fixed fully, {(time.time() - t0) / 60:.1f} min")

research/scripts/vote.py

@@ -0,0 +1,40 @@
+"""Majority vote over the K recorded answers of each problem (math tasks), against the mean single-sample accuracy.
+Also reports "vote with the module's answers" when a module run is given.
+usage: vote.py results/orig_dm/aime.jsonl data/eval/aime.jsonl"""
+import collections
+import json
+import sys
+
+sys.path.insert(0, __import__("os").path.dirname(__import__("os").path.abspath(__file__)))
+from grade import grade_math, last_boxed  # noqa: E402
+
+recs = [json.loads(l) for l in open(sys.argv[1], encoding="utf-8")]
+items = {}
+for p in sys.argv[2:]:
+    for l in open(p, encoding="utf-8"):
+        r = json.loads(l)
+        items[r["id"]] = r
+by = collections.defaultdict(list)
+for r in recs:
+    if r["id"] in items:
+        by[r["id"]].append(r)
+single = maj = oracle = 0.0
+k_seen = collections.Counter()
+for pid, g in by.items():
+    gold = items[pid]["answer"]
+    answers, oks = [], []
+    for r in g:
+        fin = r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else ""
+        b = last_boxed(fin)
+        answers.append(b.strip().replace(" ", "") if b else None)
+        oks.append(bool(grade_math(fin, gold)))
+    k_seen[len(g)] += 1
+    single += sum(oks) / len(oks)
+    oracle += any(oks)
+    votes = collections.Counter(a for a in answers if a)
+    if votes:
+        top = votes.most_common(1)[0][0]
+        maj += grade_math("\\boxed{" + top + "}", gold)
+n = len(by)
+print(f"{sys.argv[1]}: {n} problems, samples/problem {dict(k_seen)}; single {100 * single / n:.1f}  "
+      f"majority {100 * maj / n:.1f}  any-correct {100 * oracle / n:.1f}")

research/scripts/vote_design.py

@@ -0,0 +1,60 @@
+"""Designs the vote rule on DEV records (module runs with a 'module' field: none = the model finished by itself,
+fallback = answer read out at the budget). Each answer gets weight w_nat (natural) or w_fb (read-out), optionally
+times (median_len / len)^alpha (shorter reasoning = more reliable). Prints accuracy per rule.
+usage: vote_design.py results/m180_fb4/mini.jsonl data/coc/math180.jsonl [--apply w_fb alpha]"""
+import collections
+import json
+import os
+import statistics
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from grade import grade_math, grade_mmlupro, last_boxed  # noqa: E402
+import re  # noqa: E402
+
+args = [a for a in sys.argv[1:] if not a.startswith("--")]
+apply = None
+if "--apply" in sys.argv:
+    i = sys.argv.index("--apply")
+    apply = (float(sys.argv[i + 1]), float(sys.argv[i + 2]))
+    args = sys.argv[1:i]
+recs = [json.loads(l) for l in open(args[0], encoding="utf-8")]
+items = {}
+for p in args[1:]:
+    for l in open(p, encoding="utf-8"):
+        r = json.loads(l)
+        items[r["id"]] = r
+by = collections.defaultdict(list)
+for r in recs:
+    if r["id"] in items:
+        fin = r["completion"].split("</think>", 1)[1] if "</think>" in r["completion"] else ""
+        if items[r["id"]]["task"] == "mmlupro":
+            m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", fin)
+            ans, ok = (m[-1] if m else None), bool(grade_mmlupro(fin, items[r["id"]]["answer"]))
+        else:
+            b = last_boxed(fin)
+            ans, ok = (b.strip().replace(" ", "") if b else None), bool(grade_math(fin, items[r["id"]]["answer"]))
+        by[r["id"]].append({"ans": ans, "fb": r.get("module") == "fallback", "len": r["n_tokens"], "ok": ok})
+med = statistics.median(x["len"] for g in by.values() for x in g)
+
+
+def acc(w_fb, alpha):
+    tot = 0
+    for pid, g in by.items():
+        sc = collections.defaultdict(float)
+        for x in g:
+            if x["ans"]:
+                sc[x["ans"]] += (w_fb if x["fb"] else 1.0) * (med / max(x["len"], 200)) ** alpha
+        if sc:
+            top = max(sc, key=sc.get)
+            tot += next(x["ok"] for x in g if x["ans"] == top)
+    return 100 * tot / len(by)
+
+
+print(f"{args[0]}: {len(by)} problems, single {100 * sum(x['ok'] for g in by.values() for x in g) / sum(len(g) for g in by.values()):.1f}, "
+      f"any {100 * sum(any(x['ok'] for x in g) for g in by.values()) / len(by):.1f}")
+if apply:
+    print(f"rule w_fb={apply[0]} alpha={apply[1]}: {acc(*apply):.1f}")
+else:
+    for w in (1.0, 0.75, 0.5, 0.25, 0.1):
+        print(f"w_fb={w:<5}" + "".join(f"  a={a}: {acc(w, a):5.1f}" for a in (0.0, 0.5, 1.0, 2.0)))

research/scripts/vqa_validate.py

@@ -0,0 +1,39 @@
+"""Does VQAScore (LFM3-VL) agree with my 1-5 quality ratings (results/vis_ratings.json)?
+Model ratings are for sample 0 of vis_rep, reference ratings for my own files (vis_claude).
+usage: vqa_validate.py
+"""
+import json
+
+from scipy.stats import spearmanr
+
+R = json.load(open("results/vis_ratings.json", encoding="utf-8"))
+model = json.load(open("results/vis_rep/vqa.json"))
+ref = json.load(open("results/vis_claude/vqa.json"))
+xs, ys, cats = [], [], []
+for tid, r in R["model"].items():
+    if f"{tid}_0" in model:
+        xs.append(model[f"{tid}_0"]), ys.append(r), cats.append(("model", tid))
+for tid, r in R["claude"].items():
+    if f"{tid}_0" in ref:
+        xs.append(ref[f"{tid}_0"]), ys.append(r), cats.append(("ref", tid))
+rho, p = spearmanr(xs, ys)
+print(f"all {len(xs)} rated screenshots: Spearman {rho:+.2f} (p={p:.1g})")
+m = [(x, y) for x, y, c in zip(xs, ys, cats) if c[0] == "model"]
+rho, p = spearmanr([a for a, _ in m], [b for _, b in m])
+print(f"model answers only ({len(m)}): Spearman {rho:+.2f} (p={p:.1g})")
+for k in "vhct":
+    mm = [(x, y) for x, y, c in zip(xs, ys, cats) if c[0] == "model" and c[1][0] == k]
+    if len(mm) > 3:
+        rho, _ = spearmanr([a for a, _ in mm], [b for _, b in mm])
+        print(f"  {k}: n={len(mm)} Spearman {rho:+.2f}")
+by = {}
+for x, y in zip(xs, ys):
+    by.setdefault(y, []).append(x)
+print("mean VQAScore by my rating: " + ", ".join(f"{k}: {sum(v) / len(v):.2f} (n={len(v)})" for k, v in sorted(by.items())))
+# pairwise agreement: for two answers of the same task with different ratings, does VQAScore order them the same way?
+agree = tot = 0
+for tid, r in R["claude"].items():
+    if tid in R["model"] and r != R["model"][tid] and f"{tid}_0" in model and f"{tid}_0" in ref:
+        tot += 1
+        agree += (ref[f"{tid}_0"] > model[f"{tid}_0"]) == (r > R["model"][tid])
+print(f"same task, model vs my file: VQAScore orders {agree}/{tot} pairs like my ratings")

research/scripts/vqascore.py

@@ -0,0 +1,69 @@
+"""VQAScore with the vision sibling LFM3-VL-2B: P("Yes") to 'Does this image show: <task>?' for a screenshot
+(Lin et al. 2024 showed this probability tracks human judgements of text-to-image fit better than CLIPScore).
+Used to validate against my 1-5 ratings and to choose between candidate answers.
+usage: vqascore.py --png results/vis_rep/png --out results/vis_rep/vqa.json [--names v01_0 ...]
+"""
+import argparse
+import base64
+import glob
+import json
+import math
+import os
+import time
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor
+
+QUESTION = 'Does this image show the following: "{desc}"? Answer yes or no.'
+
+
+def describe_task(prompt):
+    return prompt.split("Answer with")[0].strip()
+
+
+def score(png, desc, port=8091):
+    b64 = base64.b64encode(open(png, "rb").read()).decode()
+    body = {"messages": [{"role": "user", "content": [
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64," + b64}},
+        {"type": "text", "text": QUESTION.format(desc=desc)}]}],
+        "temperature": 0.0, "max_tokens": 1, "logprobs": True, "top_logprobs": 20,
+        "chat_template_kwargs": {"enable_thinking": False}}
+    for attempt in range(5):
+        try:
+            req = urllib.request.Request(f"http://localhost:{port}/v1/chat/completions", data=json.dumps(body).encode(),
+                                         headers={"Content-Type": "application/json"})
+            r = json.loads(urllib.request.urlopen(req, timeout=1800).read())
+            break
+        except Exception:
+            if attempt == 4:
+                raise
+            time.sleep(10)
+    top = r["choices"][0]["logprobs"]["content"][0]["top_logprobs"]
+    yes = sum(math.exp(t["logprob"]) for t in top if t["token"].strip().lower() == "yes")
+    no = sum(math.exp(t["logprob"]) for t in top if t["token"].strip().lower() == "no")
+    return yes / (yes + no) if yes + no > 0 else 0.5
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--png", required=True, help="folder with <id>_<sample>.png screenshots")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--names", nargs="*")
+    ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+    ap.add_argument("--port", type=int, default=8091)
+    ap.add_argument("--workers", type=int, default=2)
+    a = ap.parse_args()
+    tasks = {json.loads(l)["id"].split("/")[-1]: json.loads(l) for l in open(a.tasks, encoding="utf-8")}
+    names = a.names or sorted(os.path.splitext(os.path.basename(p))[0] for p in glob.glob(os.path.join(a.png, "*.png")))
+    names = [n for n in names if n.rsplit("_", 1)[0] in tasks]
+    old = json.load(open(a.out)) if os.path.exists(a.out) else {}
+    todo = [n for n in names if n not in old]
+
+    def one(n):
+        return n, score(os.path.join(a.png, n + ".png"), describe_task(tasks[n.rsplit("_", 1)[0]]["prompt"]), a.port)
+
+    t0 = time.time()
+    with ThreadPoolExecutor(a.workers) as ex:
+        for n, s in ex.map(one, todo):
+            old[n] = round(s, 4)
+    json.dump(old, open(a.out, "w"), indent=0)
+    print(f"{len(todo)} scored ({(time.time() - t0) / 60:.1f} min) -> {a.out}")

research/scripts/wiki_index.py

@@ -0,0 +1,66 @@
+"""Builds a full-text (BM25, SQLite FTS5) index over English Wikipedia passages for the knowledge module.
+
+Each article is cut into passages of about --words words (paragraph-aligned); only the first --max-passages
+passages of an article are kept (the lead and early sections hold most encyclopedic facts). Parquet files are data
+(wikimedia/wikipedia 20231101.en) and are read with pyarrow only.
+usage: python -I wiki_index.py <dir with parquet files> <out .sqlite> [--words 120] [--max-passages 4]
+"""
+import argparse
+import glob
+import os
+import sqlite3
+import time
+
+import pyarrow.parquet as pq
+
+ap = argparse.ArgumentParser()
+ap.add_argument("src")
+ap.add_argument("out")
+ap.add_argument("--words", type=int, default=120)
+ap.add_argument("--max-passages", type=int, default=4)
+a = ap.parse_args()
+
+db = sqlite3.connect(a.out)
+db.execute("PRAGMA journal_mode=OFF")
+db.execute("PRAGMA synchronous=OFF")
+db.execute("CREATE VIRTUAL TABLE IF NOT EXISTS p USING fts5(title, body, tokenize='porter unicode61')")
+done = {r[0] for r in db.execute("CREATE TABLE IF NOT EXISTS files(name TEXT PRIMARY KEY)") or []}
+done = {r[0] for r in db.execute("SELECT name FROM files")}
+
+
+def passages(text):
+    out, cur = [], []
+    for para in text.split("\n"):
+        para = para.strip()
+        if not para or para.startswith(("References", "External links", "See also")):
+            continue
+        cur += para.split()
+        while len(cur) >= a.words:
+            out.append(" ".join(cur[: a.words]))
+            cur = cur[a.words:]
+            if len(out) >= a.max_passages:
+                return out
+    if cur and len(out) < a.max_passages and len(cur) > 20:
+        out.append(" ".join(cur))
+    return out
+
+
+t0, n = time.time(), 0
+for f in sorted(glob.glob(os.path.join(a.src, "*.parquet"))):
+    name = os.path.basename(f)
+    if name in done:
+        continue
+    pf = pq.ParquetFile(f)
+    for batch in pf.iter_batches(batch_size=2000, columns=["title", "text"]):
+        rows = []
+        for title, text in zip(batch.column("title").to_pylist(), batch.column("text").to_pylist()):
+            for ps in passages(text or ""):
+                rows.append((title, ps))
+        db.executemany("INSERT INTO p(title, body) VALUES (?, ?)", rows)
+        n += len(rows)
+    db.execute("INSERT INTO files VALUES (?)", (name,))
+    db.commit()
+    print(f"{name}: {n} passages total, {(time.time() - t0) / 60:.1f} min", flush=True)
+db.execute("INSERT INTO p(p) VALUES('optimize')")
+db.commit()
+print("done", n)

runtime/grade.py

@@ -0,0 +1,302 @@
+"""Grades generations produced by gen.py against the frozen eval suite.
+
+Only the final answer (text after </think>) is graded; a response whose thinking never closed counts as wrong.
+Writes results/<tag>/graded.jsonl (one row per sample) and prints per-task accuracy.
+
+usage: grade.py <tag>
+"""
+import glob
+import json
+import os
+import re
+import subprocess
+
+import sandbox  # noqa: E402  (same folder)
+import sys
+import tempfile
+import threading
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
+TMP = os.path.join(tempfile.gettempdir(), "lfm3_codeexec")
+os.makedirs(TMP, exist_ok=True)
+
+
+def final_answer(completion):
+    if "</think>" not in completion:
+        return None
+    return completion.split("</think>", 1)[1].strip()
+
+
+def last_boxed(text):
+    idx = text.rfind("\\boxed")
+    if idx < 0:
+        idx = text.rfind("\\fbox")
+        if idx < 0:
+            return None
+    i = text.find("{", idx)
+    if i < 0:
+        return None
+    depth = 0
+    for j in range(i, len(text)):
+        if text[j] == "{":
+            depth += 1
+        elif text[j] == "}":
+            depth -= 1
+            if depth == 0:
+                return text[i + 1: j]
+    return None
+
+
+def grade_math(ans_text, gold):
+    from math_verify import parse, verify
+    if ans_text is None:
+        return False
+    pred = last_boxed(ans_text)
+    if pred is None:
+        return False
+    try:
+        g = parse("$" + gold + "$", parsing_timeout=None)
+        p = parse("$" + pred + "$", parsing_timeout=None)
+        if verify(g, p, timeout_seconds=None):
+            return True
+    except Exception:
+        pass
+    norm = lambda s: re.sub(r"\\text\{(.*?)\}|\\!|\s|\$|\\left|\\right|\.$", r"\1", s).replace("dfrac", "frac")
+    return norm(pred) == norm(gold)
+
+
+def grade_gsm8k(ans_text, gold):
+    if ans_text is None:
+        return False
+    pred = last_boxed(ans_text)
+    if pred is None:
+        nums = re.findall(r"-?\d[\d,]*\.?\d*", ans_text)
+        if not nums:
+            return False
+        pred = nums[-1]
+    pred = re.sub(r"\\text\{.*?\}|\\\$|\$|,|\\%|%|\\!|\s", "", pred)
+    m = re.search(r"-?\d+(\.\d+)?", pred.replace("{", "").replace("}", ""))
+    if not m:
+        return False
+    try:
+        return abs(float(m.group(0)) - float(gold)) < 1e-6
+    except ValueError:
+        return False
+
+
+def extract_code(ans_text):
+    if ans_text is None:
+        return None
+    blocks = re.findall(r"```(?:python|py|Python)?\s*\n(.*?)```", ans_text, re.S)
+    if not blocks:
+        return None
+    return max(blocks, key=len) if len(blocks) > 1 and "def " not in blocks[-1] else blocks[-1]
+
+
+def run_python(src, timeout):
+    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(src)
+    try:
+        r = sandbox.run([sys.executable, "-I", path], capture_output=True, timeout=timeout, cwd=TMP,
+                           creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
+        return r.returncode == 0
+    except subprocess.TimeoutExpired:
+        return False
+    finally:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+def grade_humaneval(ans_text, item):
+    code = extract_code(ans_text)
+    if code is None:
+        return False
+    if f"def {item['entry_point']}" not in code:
+        code = item["code_prompt"] + code
+    src = code + "\n\n" + item["test"] + f"\n\ncheck({item['entry_point']})\n"
+    return run_python(src, 60)
+
+
+def grade_mbpp(ans_text, item):
+    code = extract_code(ans_text)
+    if code is None:
+        return False
+    src = "\n".join(item.get("test_imports") or []) + "\n" + code + "\n\n" + item["test"] + "\n"
+    return run_python(src, 60)
+
+
+def run_io(code, inp, out, timeout=6):
+    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(code)
+    try:
+        r = sandbox.run([sys.executable, "-I", path], input=inp.encode(), capture_output=True, timeout=timeout,
+                           cwd=TMP, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
+        if r.returncode != 0:
+            return False
+        return r.stdout.decode(errors="replace").split() == out.split()
+    except subprocess.TimeoutExpired:
+        return False
+    finally:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+def grade_code_io(ans_text, item):
+    code = extract_code(ans_text)
+    if code is None:
+        return False
+    return all(run_io(code, i, o) for i, o in zip(item["inputs"], item["outputs"]))
+
+
+def grade_lcb(ans_text, item):
+    from lcb_tests import run_lcb
+    code = extract_code(ans_text)
+    if code is None:
+        return False
+    k, n, _ = run_lcb(code, item, item["hidden"])
+    return k == n
+
+
+_IFEVAL_LOCK = threading.Lock()
+
+
+def grade_ifeval(ans_text, item):
+    # lm_eval fills empty instruction arguments with random values and langdetect is random too: seed both under a
+    # lock so that a response always gets the same grade
+    import random
+    from langdetect import DetectorFactory
+    from lm_eval.tasks.ifeval.utils import process_results
+    doc = {"key": 0, "prompt": item["prompt"], "instruction_id_list": item["instruction_id_list"],
+           "kwargs": item["kwargs"]}
+    with _IFEVAL_LOCK:
+        random.seed(0)
+        DetectorFactory.seed = 0
+        res = process_results(doc, [ans_text or ""])
+    return {"strict": res["prompt_level_strict_acc"], "loose": res["prompt_level_loose_acc"],
+            "inst_strict": res["inst_level_strict_acc"]}
+
+
+def grade_mmlupro(ans_text, gold):
+    if ans_text is None:
+        return False
+    m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", ans_text)
+    if not m:
+        m = re.findall(r"\\boxed\{\(?([A-J])\)?\}", ans_text)
+    return bool(m) and m[-1] == gold
+
+
+def grade_bbh(ans_text, gold, prompt=None):
+    """Exact match after extraction. With the prompt (flexible mode) it also accepts answers that state the same
+    thing in an unambiguous other form: the text of the chosen option instead of its letter, a word list separated by
+    commas, the whole Dyck sequence (input + completion) instead of the completion, plausible/not plausible for
+    yes/no. Each rule maps the answer to exactly one gold-comparable form, so it cannot accept a different answer."""
+    if ans_text is None:
+        return False
+    m = re.findall(r"answer is:?\s*(.+)", ans_text, re.I)
+    pred = m[-1] if m else (last_boxed(ans_text) or "")
+    pred = re.sub(r"\\boxed\{(.*)\}", r"\1", pred).strip().strip("*").strip().rstrip(".").strip()
+    gold = gold.strip()
+    norm = lambda s: re.sub(r"\s+", " ", s.lower().replace("`", "").replace("*", "")).strip().rstrip(".").strip()
+    if re.fullmatch(r"\([A-R]\)", gold):
+        mm = re.search(r"\(?([A-R])\)", pred) or re.fullmatch(r"([A-R])", pred)
+        if mm or prompt is None:
+            return bool(mm) and mm.group(1) == gold[1]
+        opts = {norm(t): k for k, t in re.findall(r"^\(([A-R])\)\s*(.+?)\s*$", prompt, re.M)}
+        return opts.get(norm(pred)) == gold[1]
+    if norm(pred) == norm(gold):
+        return True
+    if prompt is None:
+        return False
+    if "Sort the following words alphabetically" in prompt:
+        return re.findall(r"[^\s,]+", norm(pred)) == norm(gold).split()
+    if "parentheses are closed properly" in prompt:
+        inp = re.search(r"Input:\s*(.+)", prompt).group(1)
+        sp = lambda s: re.sub(r"\s+", "", norm(s))
+        return sp(pred) == sp(inp) + sp(gold)
+    yn = {"plausible": "yes", "not plausible": "no", "implausible": "no"}
+    return norm(gold) in ("yes", "no") and yn.get(norm(pred)) == norm(gold)
+
+
+def load_suite(pattern=os.path.join(ROOT, "data", "eval", "*.jsonl")):
+    items = {}
+    for path in glob.glob(pattern):
+        with open(path, encoding="utf-8") as f:
+            for l in f:
+                r = json.loads(l)
+                items[r["id"]] = r
+    return items
+
+
+def grade_one(rec, item):
+    a = final_answer(rec["completion"])
+    t = item["task"]
+    extra = {}
+    if t in ("math500", "aime"):
+        ok = grade_math(a, item["answer"])
+    elif t == "gsm8k":
+        ok = grade_gsm8k(a, item["answer"])
+    elif t == "humaneval":
+        ok = grade_humaneval(a, item)
+    elif t in ("mbpp", "code_fn"):
+        ok = grade_mbpp(a, item)
+    elif t == "code_io":
+        ok = grade_code_io(a, item)
+    elif t == "lcb":
+        ok = grade_lcb(a, item)
+    elif t == "math":
+        ok = grade_math(a, item["answer"])
+    elif t == "ifeval":
+        extra = grade_ifeval(a, item)
+        ok = extra["strict"]
+    elif t == "mmlupro":
+        ok = grade_mmlupro(a, item["answer"])
+    elif t == "bbh":
+        ok = grade_bbh(a, item["answer"], item["prompt"])
+        extra = {"strict": grade_bbh(a, item["answer"])}
+    else:
+        raise ValueError(t)
+    return {"id": rec["id"], "task": t, "sample": rec["sample"], "correct": bool(ok), "n_tokens": rec["n_tokens"],
+            "truncated": rec["truncated"], "closed_think": a is not None,
+            "subtask": item.get("subtask") or item.get("category") or item.get("level"), **extra}
+
+
+def main(tag, pattern=None):
+    items = load_suite(pattern) if pattern else load_suite()
+    recs = []
+    for path in glob.glob(os.path.join(ROOT, "results", tag, "*.jsonl")):
+        if path.endswith("graded.jsonl"):
+            continue
+        with open(path, encoding="utf-8") as f:
+            recs += [json.loads(l) for l in f]
+    with ThreadPoolExecutor(12) as ex:
+        graded = list(ex.map(lambda r: grade_one(r, items[r["id"]]), recs))
+    with open(os.path.join(ROOT, "results", tag, "graded.jsonl"), "w", encoding="utf-8") as f:
+        for g in graded:
+            f.write(json.dumps(g) + "\n")
+    agg = defaultdict(list)
+    for g in graded:
+        agg[g["task"]].append(g)
+    print(f"{'task':10s} {'n':>5s} {'acc':>7s} {'trunc':>6s} {'no</think>':>10s} {'avg_tok':>8s}")
+    for t in sorted(agg):
+        rows = agg[t]
+        n = len(rows)
+        print(f"{t:10s} {n:5d} {100 * sum(r['correct'] for r in rows) / n:7.2f} "
+              f"{100 * sum(r['truncated'] for r in rows) / n:6.1f} {100 * sum(not r['closed_think'] for r in rows) / n:10.1f} "
+              f"{sum(r['n_tokens'] for r in rows) / n:8.0f}")
+        if t == "ifeval":
+            print(f"{'  loose':10s} {'':5s} {100 * sum(r['loose'] for r in rows) / n:7.2f}   inst_strict "
+                  f"{100 * sum(sum(r['inst_strict']) for r in rows) / sum(len(r['inst_strict']) for r in rows):.2f}")
+
+
+if __name__ == "__main__":
+    main(*sys.argv[1:])

runtime/lcb_tests.py

@@ -0,0 +1,95 @@
+"""Test runner for LiveCodeBench items (stdin programs and LeetCode-style functional solutions).
+run_lcb(code, item, tests) -> (passed, total, first failure description); stops at the first failing test."""
+import json
+import os
+import subprocess
+
+import sandbox  # noqa: E402  (same folder)
+import sys
+import tempfile
+
+TMP = os.path.join(tempfile.gettempdir(), "lfm3_lcb")
+os.makedirs(TMP, exist_ok=True)
+
+HARNESS = r'''
+import sys, json
+from typing import *
+import collections, heapq, bisect, itertools, functools, math, string, re
+from collections import *
+from heapq import *
+from bisect import *
+from itertools import *
+from functools import *
+from math import *
+sys.setrecursionlimit(1 << 20)
+__CODE__
+_tests = json.loads(sys.stdin.read())
+for _k, (_inp, _exp) in enumerate(_tests):
+    _args = [json.loads(_l) for _l in _inp.strip().split("\n")]
+    try:
+        _res = getattr(Solution(), __FUNC__)(*_args)
+    except Exception as _e:
+        print(json.dumps({"k": _k, "err": repr(_e)[:300]}))
+        sys.exit(0)
+    _exp_v = json.loads(_exp)
+    _ok = _res == _exp_v or (isinstance(_res, tuple) and list(_res) == _exp_v) or (
+        isinstance(_res, float) and isinstance(_exp_v, (int, float)) and abs(_res - _exp_v) < 1e-6)
+    if not _ok:
+        print(json.dumps({"k": _k, "got": repr(_res)[:300]}))
+        sys.exit(0)
+print(json.dumps({"k": len(_tests)}))
+'''
+
+
+def _same(got, exp):
+    a, b = got.split(), exp.split()
+    if len(a) != len(b):
+        return False
+    for x, y in zip(a, b):
+        if x != y:
+            try:
+                if abs(float(x) - float(y)) > 1e-6 * max(1.0, abs(float(y))):
+                    return False
+            except ValueError:
+                return False
+    return True
+
+
+def _run(src, stdin, timeout):
+    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(src)
+    try:
+        r = sandbox.run([sys.executable, "-I", path], input=stdin.encode(), capture_output=True, timeout=timeout,
+                           cwd=TMP, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
+        return r.returncode, r.stdout.decode(errors="replace"), r.stderr.decode(errors="replace")
+    except subprocess.TimeoutExpired:
+        return -1, "", "TimeoutError: time limit exceeded"
+    finally:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+def run_lcb(code, item, tests, per_test_timeout=6):
+    if item["testtype"] == "functional":
+        src = HARNESS.replace("__CODE__", code).replace("__FUNC__", repr(item["func_name"]))
+        rc, so, se = _run(src, json.dumps(tests), per_test_timeout * max(1, min(len(tests), 10)))
+        try:
+            res = json.loads(so.strip().splitlines()[-1])
+        except Exception:
+            err = "\n".join(se.strip().splitlines()[-5:]) or "TimeoutError: time limit exceeded"
+            return 0, len(tests), f"Running the solution failed:\n{err}"
+        k = res["k"]
+        if k == len(tests):
+            return k, len(tests), None
+        inp, exp = tests[k]
+        got = res.get("got") or ("raised " + res.get("err", ""))
+        return k, len(tests), f"Input:\n{inp[:800]}\nExpected:\n{exp[:400]}\nYour solution returned:\n{got}"
+    for k, (inp, exp) in enumerate(tests):
+        rc, so, se = _run(code, inp, per_test_timeout)
+        if not (rc == 0 and _same(so, exp)):
+            got = so.strip()[:400] if rc == 0 else (so.strip()[:200] + "\n" + "\n".join(se.strip().splitlines()[-5:])).strip()
+            return k, len(tests), f"Input:\n{inp[:800]}\nExpected output:\n{exp[:400]}\nYour program's output:\n{got or '(nothing)'}"
+    return len(tests), len(tests), None

runtime/public_tests.py

@@ -0,0 +1,147 @@
+"""Public tests a programmer can see in the task statement itself (never the hidden grading tests):
+  code_io   -> Input/Output example pairs of the statement (Codeforces / AtCoder / "Sample Input" styles)
+  humaneval -> doctest lines (>>> call / expected value) of the docstring
+  mbpp      -> the assert lines quoted in the prompt
+run_public(task, code, prompt) -> (n_pass, n_total, first_failure_text)"""
+import os
+import re
+import subprocess
+
+import sandbox  # noqa: E402  (same folder)
+import sys
+import tempfile
+
+TMP = os.path.join(tempfile.gettempdir(), "lfm3_pub")
+os.makedirs(TMP, exist_ok=True)
+
+_HDR_IN = re.compile(r"^\s*(?:-+\s*)?(?:sample|example)?\s*input\s*(?:\d+|#\d+)?\s*:?\s*(?:-+)?\s*$", re.I)
+_HDR_OUT = re.compile(r"^\s*(?:-+\s*)?(?:sample|example)?\s*output\s*(?:\d+|#\d+)?\s*:?\s*(?:-+)?\s*$", re.I)
+_STOP = re.compile(r"^\s*(?:-{3,}.*|note:?|notes:?|explanation:?.*|constraints:?|write a python 3 program.*)\s*$", re.I)
+
+
+_EX_HDR = re.compile(r"^\s*-*\s*(?:examples?|samples?|sample tests?)\s*\d*\s*:?\s*-*\s*$", re.I)
+
+
+def _prose(text):
+    return len(re.findall(r"[A-Za-z]{3,}", text)) >= 12
+
+
+def io_examples(prompt):
+    """Example pairs after the statement's Example(s) header (format sections named Input/Output are skipped)."""
+    lines = prompt.replace("\r", "").split("\n")
+    pairs = []
+    i = next((k + 1 for k, l in enumerate(lines) if _EX_HDR.match(l)), 0)
+    while i < len(lines):
+        if _HDR_IN.match(lines[i]) and not re.match(r"^\s*-+\s*input\s*-+\s*$", lines[i], re.I):
+            j = i + 1
+            inp = []
+            while j < len(lines) and not _HDR_OUT.match(lines[j]):
+                if _HDR_IN.match(lines[j]) or _STOP.match(lines[j]):
+                    break
+                inp.append(lines[j])
+                j += 1
+            if j < len(lines) and _HDR_OUT.match(lines[j]):
+                k = j + 1
+                out = []
+                while k < len(lines) and not _HDR_IN.match(lines[k]) and not _STOP.match(lines[k]):
+                    out.append(lines[k])
+                    k += 1
+                for b in range(1, len(out)):  # an explanation paragraph after the expected output is not output
+                    if not out[b].strip() and len(re.findall(r"[A-Za-z]{3,}", " ".join(out[b + 1:b + 4]))) >= 4:
+                        out = out[:b]
+                        break
+                inp_s, out_s = "\n".join(inp).strip("\n"), "\n".join(out).strip()
+                if inp_s.strip() and out_s and not _prose(inp_s):
+                    pairs.append((inp_s + "\n", out_s))
+                i = k
+                continue
+        i += 1
+    return pairs
+
+
+def doctests(prompt):
+    tests, lines = [], prompt.split("\n")
+    for i, l in enumerate(lines):
+        s = l.strip()
+        if s.startswith(">>>") and i + 1 < len(lines):
+            exp = lines[i + 1].strip()
+            if exp and not exp.startswith(">>>") and not exp.startswith('"""'):
+                tests.append((s[3:].strip(), exp))
+    return tests
+
+
+def asserts(prompt):
+    return [l.strip() for l in prompt.split("\n") if l.strip().startswith("assert ")]
+
+
+def _run(src, stdin="", timeout=6):
+    fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(src)
+    try:
+        r = sandbox.run([sys.executable, "-I", path], input=stdin.encode(), capture_output=True, timeout=timeout,
+                           cwd=TMP, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
+        return r.returncode, r.stdout.decode(errors="replace"), r.stderr.decode(errors="replace")
+    except subprocess.TimeoutExpired:
+        return -1, "", "TimeoutError: time limit exceeded (6 s)"
+    finally:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+def _same(got, exp):
+    a, b = got.split(), exp.split()
+    if len(a) != len(b):
+        return False
+    for x, y in zip(a, b):
+        if x != y:
+            try:
+                if abs(float(x) - float(y)) > 1e-6 * max(1.0, abs(float(y))):
+                    return False
+            except ValueError:
+                return False
+    return True
+
+
+def _tail(err, n=6):
+    return "\n".join(err.strip().splitlines()[-n:])
+
+
+def run_public(task, code, prompt, item=None):
+    """Returns (passed, total, failure description or None)."""
+    if task == "lcb":
+        from lcb_tests import run_lcb
+        return run_lcb(code, item, item["public"])
+    if task == "code_io":
+        tests = io_examples(prompt)
+        fails = None
+        ok = 0
+        for k, (inp, out) in enumerate(tests):
+            rc, so, se = _run(code, inp)
+            if rc == 0 and _same(so, out):
+                ok += 1
+            elif fails is None:
+                got = so.strip()[:600] if rc == 0 else (so.strip()[:300] + "\n" + _tail(se)).strip()
+                fails = f"Example {k + 1}\nInput:\n{inp.strip()[:800]}\nExpected output:\n{out[:600]}\n" \
+                        f"Your program's output:\n{got or '(nothing)'}"
+        return ok, len(tests), fails
+    if task == "humaneval":
+        tests = doctests(prompt)
+        checks = [f"_r = {call}\n_e = {exp}\nprint('PASS' if _r == _e else 'FAIL ' + repr(_r))" for call, exp in tests]
+    elif task == "mbpp":
+        tests = asserts(prompt)
+        checks = [f"try:\n    {a}\n    print('PASS')\nexcept AssertionError:\n    print('FAIL')" for a in tests]
+    else:
+        return 0, 0, None
+    ok, fails = 0, None
+    for k, (t, chk) in enumerate(zip(tests, checks)):
+        rc, so, se = _run(code + "\n\n" + chk)
+        if rc == 0 and so.strip().endswith("PASS"):
+            ok += 1
+        elif fails is None:
+            shown = t if isinstance(t, str) else f">>> {t[0]}\nexpected: {t[1]}"
+            res = so.strip().splitlines()[-1] if rc == 0 and so.strip() else _tail(se)
+            fails = f"Test:\n{shown}\nResult: {res}"
+    return ok, len(tests), fails

runtime/render.py

@@ -0,0 +1,122 @@
+"""Headless renderer for visual answers (SVG / HTML / canvas / three.js) with objective checks.
+
+render(kind, source, png_path) opens the page in headless Chrome (the installed browser, via playwright) at 1024x768,
+waits for it to settle and returns a report:
+  errors      JavaScript exceptions and console errors (messages)
+  failed      network requests that failed (e.g. a wrong library URL)
+  blank       the screenshot is (almost) a single colour
+  overflow    the page is wider than the viewport (horizontal scroll)
+  animated    two screenshots 0.7 s apart differ (for animation tasks)
+  coverage    share of the viewport that is not background colour
+  parse_error for SVG: the XML does not parse
+The screenshot is saved to png_path. Used by the visual benchmark and by the render-check-repair module.
+"""
+import io
+import os
+import re
+import tempfile
+import xml.dom.minidom
+
+import numpy as np
+from PIL import Image
+from playwright.sync_api import sync_playwright
+
+W, H = 1024, 768
+TMP = os.path.join(tempfile.gettempdir(), "lfm3_render")
+os.makedirs(TMP, exist_ok=True)
+
+SVG_PAGE = """<!doctype html><html><head><meta charset="utf-8"><style>html,body{margin:0;background:#fff}
+#c{width:1024px;height:768px;display:flex;align-items:center;justify-content:center}
+#c>svg{max-width:1024px;max-height:768px;width:100%;height:100%}</style></head><body><div id="c">__SVG__</div></body></html>"""
+
+
+def _stats(png):
+    a = np.asarray(Image.open(io.BytesIO(png)).convert("RGB")).astype(np.int16)
+    flat = a.reshape(-1, 3)
+    vals, counts = np.unique(flat[::7], axis=0, return_counts=True)
+    bg = vals[counts.argmax()]
+    diff = np.abs(a - bg).sum(axis=2) > 24
+    return a, float(diff.mean())
+
+
+def extract(text, kind):
+    if text is None:
+        return None
+    blocks = re.findall(r"```(\w*)\s*\n(.*?)```", text, re.S)
+    for lang, body in reversed(blocks):
+        if kind == "svg" and ("<svg" in body):
+            return body[body.index("<svg"):].strip()
+        if kind == "html" and ("<html" in body.lower() or "<body" in body.lower() or "<canvas" in body or "<script" in body):
+            return body.strip()
+    if kind == "svg":
+        m = re.search(r"<svg.*</svg>", text, re.S)
+        return m.group(0) if m else None
+    m = re.search(r"<!doctype html.*</html>|<html.*</html>", text, re.S | re.I)
+    return m.group(0) if m else None
+
+
+class Renderer:
+    def __init__(self, channel="chrome"):
+        self.pw = sync_playwright().start()
+        self.browser = self.pw.chromium.launch(channel=channel, headless=True,
+                                               args=["--use-angle=swiftshader", "--enable-unsafe-swiftshader"])
+
+    def close(self):
+        self.browser.close()
+        self.pw.stop()
+
+    def render(self, kind, source, png_path, wait_ms=1500):
+        rep = {"errors": [], "failed": [], "parse_error": None}
+        if kind == "svg":
+            try:
+                xml.dom.minidom.parseString(source.encode("utf-8"))
+            except Exception as e:
+                rep["parse_error"] = str(e)[:200]
+            html = SVG_PAGE.replace("__SVG__", source)
+        else:
+            html = source
+        path = os.path.join(TMP, f"page_{os.getpid()}_{abs(hash(png_path)) % 10**8}.html")
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(html)
+        page = self.browser.new_page(viewport={"width": W, "height": H})
+        page.on("pageerror", lambda e: rep["errors"].append(f"Uncaught {e.name}: {e.message}"[:300]))
+        page.on("console", lambda m: rep["errors"].append(("console.error: " + m.text)[:300]) if m.type == "error" else None)
+        page.on("requestfailed", lambda r: rep["failed"].append(r.url[:200]))
+        try:
+            page.goto("file:///" + path.replace("\\", "/"), wait_until="load", timeout=20000)
+            page.wait_for_timeout(wait_ms)
+            png1 = page.screenshot()
+            page.wait_for_timeout(700)
+            png2 = page.screenshot()
+            rep["overflow"] = bool(page.evaluate("document.documentElement.scrollWidth > window.innerWidth + 2"))
+            rep["page_height"] = int(page.evaluate("document.documentElement.scrollHeight"))
+        except Exception as e:
+            rep["errors"].append(f"page failed to load: {str(e)[:200]}")
+            png1 = png2 = page.screenshot() if not page.is_closed() else None
+            rep["overflow"], rep["page_height"] = False, 0
+        finally:
+            page.close()
+            try:
+                os.remove(path)
+            except OSError:
+                pass
+        if png1 is None:
+            rep.update(blank=True, coverage=0.0, animated=False)
+            return rep
+        with open(png_path, "wb") as f:
+            f.write(png2)
+        a1, cov = _stats(png1)
+        a2, _ = _stats(png2)
+        rep["coverage"] = round(cov, 4)
+        rep["blank"] = cov < 0.003
+        rep["animated"] = bool((np.abs(a1 - a2).sum(axis=2) > 24).mean() > 0.0005)
+        return rep
+
+
+if __name__ == "__main__":
+    import json
+    import sys
+    r = Renderer()
+    kind, src, out = sys.argv[1], open(sys.argv[2], encoding="utf-8").read(), sys.argv[3]
+    print(json.dumps(r.render(kind, src, out), ensure_ascii=False))
+    r.close()

runtime/sandbox.py

@@ -0,0 +1,42 @@
+"""subprocess.run replacement for running untrusted model programs on Windows.
+
+Two Windows traps make subprocess.run(..., timeout=) hang forever on a stuck program:
+  * in a venv, python.exe is a launcher that starts the real interpreter as a child; the timeout kills only the
+    launcher and the child keeps the pipes open;
+  * Popen.communicate writes `input` synchronously before it starts waiting, so a program that stops reading a large
+    input blocks the write and the timeout is never checked.
+run() feeds stdin from a temporary file and kills the whole process tree on timeout, then raises
+subprocess.TimeoutExpired like subprocess.run."""
+import os
+import subprocess
+import tempfile
+
+
+def run(cmd, input=None, timeout=None, capture_output=True, text=False, encoding=None, errors=None, **kw):
+    kw.setdefault("creationflags", getattr(subprocess, "CREATE_NO_WINDOW", 0))
+    fin, path = subprocess.DEVNULL, None
+    if input is not None:
+        fd, path = tempfile.mkstemp(suffix=".in")
+        with os.fdopen(fd, "wb") as f:
+            f.write(input if isinstance(input, bytes) else input.encode(encoding or "utf-8"))
+        fin = open(path, "rb")
+    try:
+        p = subprocess.Popen(cmd, stdin=fin, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=text,
+                             encoding=encoding, errors=errors, **kw)
+        try:
+            out, err = p.communicate(timeout=timeout)
+        except subprocess.TimeoutExpired:
+            subprocess.run(["taskkill", "/F", "/T", "/PID", str(p.pid)], capture_output=True)
+            try:
+                p.communicate(timeout=10)
+            except subprocess.TimeoutExpired:
+                pass
+            raise subprocess.TimeoutExpired(cmd, timeout)
+        return subprocess.CompletedProcess(cmd, p.returncode, out, err)
+    finally:
+        if path:
+            fin.close()
+            try:
+                os.remove(path)
+            except OSError:
+                pass

runtime/vis_autofix.py

@@ -0,0 +1,67 @@
+"""Deterministic fixes of visual answers before rendering (no model call), measured on recorded runs.
+
+  importmap  the page loads three.js as ES modules (imports from 'three', 'three/addons/...' or from the unpkg URL,
+             whose addons import the bare name 'three') but has no import map: insert the import map of the task
+             at the start of <head>.
+
+usage: vis_autofix.py --src results/vis_rep/vis.jsonl --out results/vis_rep_fix/vis.jsonl
+Records whose answer is changed get "autofix": [...]; the rest are copied unchanged (exactly paired).
+"""
+import argparse
+import json
+import os
+import re
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from render import extract  # noqa: E402
+
+IMPORT_MAP = ('<script type="importmap">{"imports": {"three": "https://unpkg.com/three@0.160.0/build/three.module.js", '
+              '"three/addons/": "https://unpkg.com/three@0.160.0/examples/jsm/"}}</script>')
+USES_THREE_MODULE = re.compile(r"""from\s+['"](three|three/addons/[^'"]+|https://unpkg\.com/three@[^'"]+)['"]""")
+
+
+def fix_importmap(src):
+    if 'type="importmap"' in src or "type='importmap'" in src or not USES_THREE_MODULE.search(src):
+        return src
+    m = re.search(r"<head[^>]*>", src, re.I)
+    if m:
+        return src[:m.end()] + "\n" + IMPORT_MAP + src[m.end():]
+    m = re.search(r"<script", src, re.I)
+    return src[:m.start()] + IMPORT_MAP + "\n" + src[m.start():] if m else src
+
+
+def autofix(text, kind):
+    """Returns (new answer text, list of applied fixes)."""
+    if kind != "html" or text is None:
+        return text, []
+    src = extract(text, kind)
+    if not src:
+        return text, []
+    new = fix_importmap(src)
+    if new == src:
+        return text, []
+    return text.replace(src, new), ["importmap"]
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
+    a = ap.parse_args()
+    tasks = {json.loads(l)["id"]: json.loads(l) for l in open(a.tasks, encoding="utf-8")}
+    os.makedirs(os.path.dirname(a.out), exist_ok=True)
+    n = 0
+    with open(a.out, "w", encoding="utf-8") as f:
+        for l in open(a.src, encoding="utf-8"):
+            r = json.loads(l)
+            c = r["completion"]
+            if "</think>" in c:
+                head, fin = c.split("</think>", 1)
+                new, fixes = autofix(fin, tasks[r["id"]]["kind"])
+                if fixes:
+                    r = {**r, "completion": head + "</think>" + new, "autofix": fixes}
+                    n += 1
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    print(f"{n} answers changed -> {a.out}")

runtime/visual_loop.py

@@ -0,0 +1,96 @@
+"""Shared render-check-repair logic for visual answers, used by lfm3_deliberate.solve_visual.
+(The same rules as scripts/visual_repair.py, which produced the measured numbers.)
+
+problems(rep, animated, src, kind, hints) -> (weighted badness, feedback lines) from a render.py report;
+feedback(kind, lines) -> the repair turn text; fix_importmap from vis_autofix (deterministic, no model call).
+"""
+import re
+
+from vis_autofix import IMPORT_MAP, fix_importmap  # noqa: F401  (re-exported for callers)
+
+# browser error -> the fix a web developer would apply (general knowledge, never a solution of a particular task)
+HINTS = [
+    (r'Failed to resolve module specifier "three', "A module imports the bare name 'three' (three.js addons do this "
+     f"too). Put the import map in <head>, before any module script: {IMPORT_MAP}, then write "
+     "import * as THREE from 'three'; and import addons like import { OrbitControls } from "
+     "'three/addons/controls/OrbitControls.js';"),
+    (r"THREE is not defined", "three.js is not a global variable. Add the import map "
+     f"{IMPORT_MAP} and, inside <script type=\"module\">, write import * as THREE from 'three';"),
+    (r"THREE\.(\w+) is not a constructor", "THREE.{0} does not exist in three.js r160. Addons are imported "
+     "separately, e.g. import { OrbitControls } from 'three/addons/controls/OrbitControls.js'; then "
+     "new OrbitControls(camera, renderer.domElement). Otherwise check the class name."),
+    (r"Assignment to constant variable", "A variable declared with const is assigned again: declare variables "
+     "that change with let."),
+    (r"ReferenceError: (\w+) is not defined", "'{0}' is used where it is not declared: declare it with const/let "
+     "at the top level, before the functions that use it run."),
+    (r"Cannot read properties of (?:undefined|null) \(reading '(\w+)'\)", "Something is undefined or null when "
+     "its property '{0}' is read: an element id that does not exist, an object used before it is created, or an "
+     "array index out of range."),
+    (r"(\S+) is not a function", "{0} is not a function: a property is called like a method or a name is "
+     "misspelled (canvas.width/height are numbers; drawing methods are ctx.fillRect, ctx.arc, ctx.beginPath, "
+     "ctx.fill, ctx.stroke; angles in Math are radians, there is no Math.radians)."),
+]
+BLANK_THREE = ("Add lights for MeshStandardMaterial (AmbientLight and DirectionalLight), put the camera away from "
+               "the objects with camera.lookAt(0, 0, 0) and render every frame with "
+               "renderer.setAnimationLoop(() => renderer.render(scene, camera)).")
+BLANK_CANVAS = ("Set canvas.width = innerWidth and canvas.height = innerHeight, draw in every frame and use a fill "
+                "color that differs from the background.")
+STATIC_HINT = ("Compute positions from the elapsed time inside a requestAnimationFrame loop (or "
+               "renderer.setAnimationLoop) and redraw every frame.")
+SVG_HINT = "Close every tag, quote every attribute value, write & as &amp; and keep one root <svg> element."
+
+
+def hints_for(rep, animated, src):
+    out = []
+    for e in rep.get("errors", []):
+        for pat, text in HINTS:
+            m = re.search(pat, e)
+            if m:
+                h = text.replace("{0}", m.group(1)) if m.groups() else text
+                if h not in out:
+                    out.append(h)
+                break
+    if rep.get("parse_error"):
+        out.append(SVG_HINT)
+    if rep.get("blank") and not rep.get("parse_error") and not rep.get("errors"):
+        out.append(BLANK_THREE if "three" in (src or "") else BLANK_CANVAS if "<canvas" in (src or "") else
+                   "Make sure the content is visible: real sizes, colors that differ from the background, nothing "
+                   "hidden by CSS.")
+    if animated and not rep.get("animated") and not rep.get("blank"):
+        out.append(STATIC_HINT)
+    return out
+
+
+def problems(rep, animated, src, kind, hints=False):
+    """(weighted badness, feedback lines); 0 = renders cleanly."""
+    if not src:
+        return 10, [f"Your answer does not contain the complete file in a ```{kind} code block."]
+    bad, lines = 0, []
+    if rep.get("parse_error"):
+        bad += 4
+        lines.append(f"The SVG is not valid XML: {rep['parse_error']}")
+    if rep["errors"]:
+        bad += 3
+        lines += [f"JavaScript error: {e}" for e in rep["errors"][:4]]
+    if rep["failed"]:
+        bad += 3
+        lines += [f"This resource failed to load: {u}" for u in rep["failed"][:3]]
+    if rep["blank"]:
+        bad += 3
+        lines.append("Nothing is visible: the rendered page is blank (a single colour).")
+    if animated and not rep["animated"]:
+        bad += 1
+        lines.append("The animation does not move: two screenshots taken 0.7 seconds apart are identical.")
+    if rep.get("overflow"):
+        bad += 1
+        lines.append("The page is wider than the 1024 px window (there is a horizontal scrollbar).")
+    if hints and bad:
+        lines += ["Hint: " + h for h in hints_for(rep, animated, src)]
+    return bad, lines
+
+
+def feedback(kind, lines):
+    what = "SVG" if kind == "svg" else "page"
+    return (f"I opened your {what} in Chrome (window 1024x768) and found these problems:\n"
+            + "\n".join("- " + x for x in lines)
+            + f"\n\nFix them and give the complete corrected file in a single ```{kind} code block.")

tokenizer/chat_template.jinja

@@ -0,0 +1,280 @@
+{{- bos_token -}}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{%- set enable_thinking = enable_thinking | default(true) -%}
+
+{%- macro render_value(v) -%}
+    {%- if v is string -%}
+        {{- v -}}
+    {%- elif v is boolean -%}
+        {{- "true" if v else "false" -}}
+    {%- elif v is none -%}
+        {{- "null" -}}
+    {%- elif v is mapping or (v is sequence and v is not string) -%}
+        {{- v | tojson -}}
+    {%- else -%}
+        {{- v -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro parse_content(content) -%}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping and content.get("type") == "image" -%}
+        {{- "<|vision_start|><|image|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") == "video" -%}
+        {{- "<|vision_start|><|video|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") == "text" -%}
+        {{- (content.get("text") or "") | string -}}
+    {%- elif content is mapping -%}
+        {{- content | tojson -}}
+    {%- elif content is iterable -%}
+        {%- set _ns = namespace(result="") -%}
+        {%- for item in content -%}
+            {%- if item is string -%}
+                {%- set _ns.result = _ns.result + item -%}
+            {%- elif item is mapping and item.get("type") == "image" -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|image|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") == "video" -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|video|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") == "text" -%}
+                {%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
+            {%- else -%}
+                {%- set _ns.result = _ns.result + (item | tojson) -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {{- _ns.result -}}
+    {%- elif content is not none -%}
+        {{- content | tojson -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_calls(tool_calls) -%}
+    {%- for tool_call in tool_calls -%}
+        {%- if not loop.first -%}{{- "\n\n" -}}{%- endif -%}
+        {%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
+        {%- set func_args = func.get("arguments") -%}
+        {%- if func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
+            {{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
+        {%- endif -%}
+        {{- "<tool_call>\n<function=" + func["name"] + ">\n" -}}
+        {%- if func_args is mapping -%}
+            {%- for arg_name, arg_value in func_args.items() -%}
+                {{- "<parameter=" + arg_name + ">\n" + render_value(arg_value) + "\n</parameter>\n" -}}
+            {%- endfor -%}
+        {%- endif -%}
+        {{- "</function>\n</tool_call>" -}}
+    {%- endfor -%}
+{%- endmacro -%}
+
+{%- macro tool_to_json(tool) -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- set ns_t = namespace(first=true) -%}
+    {{- "{" -}}
+    {%- for k, v in fn.items() -%}
+        {%- if k != "defer_loading" -%}
+            {%- if not ns_t.first -%}{{- ", " -}}{%- endif -%}
+            {%- set ns_t.first = false -%}
+            {{- k | tojson -}}{{- ": " -}}{{- v | tojson -}}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- "}" -}}
+{%- endmacro -%}
+
+{%- macro is_tool_reference(item) -%}
+    {{- "1" if item is mapping and item.get("type") == "tool_reference" else "" -}}
+{%- endmacro -%}
+
+{%- macro referenced_tool_to_json(name) -%}
+    {%- set _ns = namespace(found=none) -%}
+    {%- for tool in tools or [] -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- if fn.name == name -%}
+            {%- set _ns.found = fn -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if _ns.found is none -%}
+        {{- raise_exception("Tool reference '" + name | string + "' not found in available tools. Every tool referenced by a tool_reference must be in `tools` on every request.") -}}
+    {%- endif -%}
+    {{- tool_to_json(_ns.found) -}}
+{%- endmacro -%}
+
+{%- macro render_tool_response(content) -%}
+    {%- if content is string and "<tool_response>" in content -%}
+        {{- content -}}
+    {%- else -%}
+    {{- "<tool_response>\n" -}}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping -%}
+        {{- parse_content(content) -}}
+    {%- elif content is iterable -%}
+        {%- for item in content -%}
+            {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+            {%- if is_tool_reference(item) -%}
+                {%- if loop.first or not is_tool_reference(loop.previtem) -%}
+                    {{- "<additional_tools>\n" -}}
+                {%- endif -%}
+                {{- referenced_tool_to_json(item.name) -}}
+                {%- if loop.last or not is_tool_reference(loop.nextitem) -%}
+                    {{- "\n</additional_tools>" -}}
+                {%- endif -%}
+            {%- elif item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
+                {{- "<additional_tools>\n" -}}
+                {%- set _tl = item.result.tools -%}
+                {%- if _tl is string -%}
+                    {{- _tl -}}
+                {%- elif _tl is mapping -%}
+                    {{- tool_to_json(_tl) -}}
+                {%- else -%}
+                    {%- for t in _tl -%}
+                        {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+                        {{- tool_to_json(t) -}}
+                    {%- endfor -%}
+                {%- endif -%}
+                {{- "\n</additional_tools>" -}}
+            {%- else -%}
+                {{- parse_content(item) -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- elif content is not none -%}
+        {{- content | tojson -}}
+    {%- endif -%}
+    {{- "\n</tool_response>" -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
+{%- if messages and messages[0]["role"] == "system" -%}
+    {%- if messages[0].get("content") -%}
+        {%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
+    {%- endif -%}
+    {%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- if tools -%}
+    {%- set tools_ns = namespace(lines=[]) -%}
+    {%- for tool in tools -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- set _dyn = (fn.defer_loading is defined and fn.defer_loading) or (tool.defer_loading is defined and tool.defer_loading) -%}
+        {%- if not _dyn -%}
+            {%- set tools_ns.lines = tools_ns.lines + [tool_to_json(fn)] -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if tools_ns.lines -%}
+        {%- set ns.system_prompt = ns.system_prompt + ("\n\n" if ns.system_prompt else "") + "<tools>\n" + (tools_ns.lines | join("\n")) + "\n</tools>" -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if ns.system_prompt -%}
+    {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if message["role"] == "user" -%}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- set carry = namespace(pending=[]) -%}
+{%- for message in messages -%}
+    {%- set _next = messages[loop.index0 + 1] if not loop.last else none -%}
+    {%- set _merge = message.role == "tool" and not message.get("additional_tools")
+        and _next is not none and _next.role == "tool" and not _next.get("additional_tools")
+        and message.get("tool_call_id") and _next.get("tool_call_id") == message.get("tool_call_id") -%}
+    {%- set _items = [] -%}
+    {%- if message.role == "tool" and not message.get("additional_tools") -%}
+        {%- set _c = message.get("content") -%}
+        {%- if _c is string -%}
+            {%- set _items = [{"type": "text", "text": _c}] if _c else [] -%}
+        {%- elif _c is mapping -%}
+            {%- set _items = [_c] -%}
+        {%- elif _c is iterable -%}
+            {%- set _items = _c | list -%}
+        {%- elif _c is not none -%}
+            {%- set _items = [_c] -%}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if _merge -%}
+        {%- set carry.pending = carry.pending + _items -%}
+    {%- else -%}
+    {{- "<|im_start|>" + message.role + "\n" -}}
+    {%- if message.role == "assistant" -%}
+        {%- generation -%}
+        {%- set thinking = message.thinking if message.thinking is defined else message.reasoning if message.reasoning is defined else message.reasoning_content -%}
+        {%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
+        {%- if thinking is defined and keep_thinking -%}
+            {{- "<think>" + (thinking or "") + "</think>" -}}
+        {%- elif keep_thinking and "</think>" not in (message.get("content") if message.get("content") is string else "") -%}
+            {{- "<think></think>" -}}
+        {%- endif -%}
+        {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
+        {%- set _has_cfm = false -%}
+        {%- set content = "" -%}
+        {%- if message.get("content") -%}
+            {%- set content = parse_content(message.content) -%}
+        {%- endif -%}
+        {%- if not keep_thinking -%}
+            {%- if "</think>" in content -%}
+                {%- set content = content.split("</think>")[-1] | trim -%}
+            {%- endif -%}
+            {%- set content = "<think></think>" + content -%}
+        {%- endif -%}
+        {%- if content.endswith(_cfm_tag) -%}
+            {%- set _has_cfm = true -%}
+            {%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
+            {%- set content = content[:_trunc_len] -%}
+        {%- endif -%}
+        {{- content -}}
+        {%- if message.tool_calls -%}
+            {{- render_tool_calls(message.tool_calls) -}}
+        {%- endif -%}
+        {%- if _has_cfm -%}
+            {{- _cfm_tag -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+        {%- endgeneration -%}
+    {%- elif message.role == "tool" -%}
+        {%- if message.get("additional_tools") -%}
+            {%- set _at = message["additional_tools"] -%}
+            {%- if _at is mapping or (_at is sequence and _at is not string) -%}
+                {%- set _src = _at -%}
+            {%- else -%}
+                {%- set _src = message.get("content") -%}
+            {%- endif -%}
+            {%- if _src is string and "<tool_response>" in _src -%}
+                {{- _src -}}
+            {%- elif _src is string and "<additional_tools>" in _src -%}
+                {{- "<tool_response>\n" + _src + "\n</tool_response>" -}}
+            {%- elif _src -%}
+                {{- "<tool_response>\n<additional_tools>\n" -}}
+                {%- if _src is string -%}
+                    {{- _src -}}
+                {%- elif _src is mapping -%}
+                    {{- tool_to_json(_src) -}}
+                {%- else -%}
+                    {%- for t in _src -%}
+                        {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+                        {{- tool_to_json(t) -}}
+                    {%- endfor -%}
+                {%- endif -%}
+                {{- "\n</additional_tools>\n</tool_response>" -}}
+            {%- endif -%}
+        {%- elif carry.pending -%}
+            {{- render_tool_response(carry.pending + _items) -}}
+            {%- set carry.pending = [] -%}
+        {%- else -%}
+            {{- render_tool_response("" if message.get("content") is none else message.get("content")) -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- else %}
+        {%- if message.get("content") -%}
+            {{- parse_content(message["content"]) -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- endif %}
+    {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- "<|im_start|>assistant\n" -}}
+    {%- if enable_thinking -%}
+        {{- "<think>" -}}
+    {%- else -%}
+        {{- "<think></think>" -}}
+    {%- endif -%}
+{%- endif -%}

tokenizer/config.json

@@ -0,0 +1,68 @@
+{
+  "architectures": [
+    "Lfm3ForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_lfm3.Lfm3Config",
+    "AutoModel": "modeling_lfm3.Lfm3Model",
+    "AutoModelForCausalLM": "modeling_lfm3.Lfm3ForCausalLM"
+  },
+  "bos_token_id": 124894,
+  "conv_L_cache": 7,
+  "conv_bias": false,
+  "dtype": "bfloat16",
+  "eos_token_id": 124900,
+  "head_dim": 128,
+  "hidden_size": 2048,
+  "intermediate_size": 7680,
+  "layer_types": [
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "sliding_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "sliding_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "sliding_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "sliding_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "sliding_attention",
+    "conv",
+    "full_attention",
+    "conv"
+  ],
+  "max_position_embeddings": 262144,
+  "model_type": "lfm3",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 4,
+  "pad_token_id": 124893,
+  "rope_parameters": {
+    "full_attention": {
+      "partial_rotary_factor": 0.5,
+      "rope_theta": 1000000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000,
+      "rope_type": "default"
+    }
+  },
+  "sliding_window": 256,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.1",
+  "use_cache": true,
+  "use_sandwich_norm": true,
+  "vocab_size": 128000
+}