diff --git a/research/scripts/rate.py b/research/scripts/rate.py new file mode 100644 index 0000000000000000000000000000000000000000..bcff52645e6a42b107a95fd3b5518b07e5649cfe --- /dev/null +++ b/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)) diff --git a/research/scripts/rbud_test.py b/research/scripts/rbud_test.py new file mode 100644 index 0000000000000000000000000000000000000000..bdd173bf299b50bd0469ae59b9b6a8a50e89dfdb --- /dev/null +++ b/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") diff --git a/research/scripts/ref_check.py b/research/scripts/ref_check.py new file mode 100644 index 0000000000000000000000000000000000000000..0892efd01546e3eff9b87d49d9627c298b021dae --- /dev/null +++ b/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 print the statement of sample problem n (0-based) + ref_check.py check run data/claude_ref/.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") diff --git a/research/scripts/render.py b/research/scripts/render.py new file mode 100644 index 0000000000000000000000000000000000000000..1b7c77fa59967854d679f95db34e53b527ff7e20 --- /dev/null +++ b/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 = """
__SVG__
""" + + +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 ("", text, re.S) + return m.group(0) if m else None + m = re.search(r"|", 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() diff --git a/research/scripts/repair.py b/research/scripts/repair.py new file mode 100644 index 0000000000000000000000000000000000000000..59eb3107b4864a4bf6cc51450ef93b51ee51a1f7 --- /dev/null +++ b/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("", 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 "" not in rec["completion"]: + return rec + t = task_of(item) + head, final = rec["completion"].split("", 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 + "" + 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": "(repaired)" + 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") diff --git a/research/scripts/rl_iter.sh b/research/scripts/rl_iter.sh new file mode 100644 index 0000000000000000000000000000000000000000..cc38a848d16b89db8f20971084474766194f351c --- /dev/null +++ b/research/scripts/rl_iter.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# usage: rl_iter.sh [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 diff --git a/research/scripts/rl_iter_bc.sh b/research/scripts/rl_iter_bc.sh new file mode 100644 index 0000000000000000000000000000000000000000..51f40589f135e7a8480d5764679b3464a90e3b2b --- /dev/null +++ b/research/scripts/rl_iter_bc.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# RL iteration with adaptive sampling + commitment bootstrap. +# usage: rl_iter_bc.sh [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 - <&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 diff --git a/research/scripts/rl_train.py b/research/scripts/rl_train.py new file mode 100644 index 0000000000000000000000000000000000000000..35bb0526578e45510475e99d6f97c4910b805662 --- /dev/null +++ b/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) diff --git a/research/scripts/run_dev.sh b/research/scripts/run_dev.sh new file mode 100644 index 0000000000000000000000000000000000000000..38c8482df9ffc00b26663c985e0242b35bb7e7bf --- /dev/null +++ b/research/scripts/run_dev.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# usage: run_dev.sh [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 diff --git a/research/scripts/run_eval.sh b/research/scripts/run_eval.sh new file mode 100644 index 0000000000000000000000000000000000000000..594376a0290db709788e2e165c6d64bc4ad45aaa --- /dev/null +++ b/research/scripts/run_eval.sh @@ -0,0 +1,9 @@ +#!/bin/bash +# usage: run_eval.sh (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 diff --git a/research/scripts/run_mini.sh b/research/scripts/run_mini.sh new file mode 100644 index 0000000000000000000000000000000000000000..a6c916a2c7d91ae1ef726e382e70389286d190a5 --- /dev/null +++ b/research/scripts/run_mini.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# usage: run_mini.sh [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 diff --git a/research/scripts/run_mini2.sh b/research/scripts/run_mini2.sh new file mode 100644 index 0000000000000000000000000000000000000000..ae8911cb769dce2bf26da50157b1dc58235b7029 --- /dev/null +++ b/research/scripts/run_mini2.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# usage: [ENV=..] run_mini2.sh -- 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 diff --git a/research/scripts/run_quick.sh b/research/scripts/run_quick.sh new file mode 100644 index 0000000000000000000000000000000000000000..1a95e7cf751e80b2fa7f8cb23a649ffe5552a013 --- /dev/null +++ b/research/scripts/run_quick.sh @@ -0,0 +1,9 @@ +#!/bin/bash +# usage: run_quick.sh [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 diff --git a/research/scripts/sample_prompts.py b/research/scripts/sample_prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..aa9152832c14c0a26b072444cadc1ec1143f7dbc --- /dev/null +++ b/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)) diff --git a/research/scripts/sandbox.py b/research/scripts/sandbox.py new file mode 100644 index 0000000000000000000000000000000000000000..522e6360793c1e170504b3d0dc5b1bb22af77a07 --- /dev/null +++ b/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 diff --git a/research/scripts/serve.sh b/research/scripts/serve.sh new file mode 100644 index 0000000000000000000000000000000000000000..eb2d194ef89c55255b54511bee60fbe89b719e7f --- /dev/null +++ b/research/scripts/serve.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# usage: serve.sh [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 diff --git a/research/scripts/sft_train.py b/research/scripts/sft_train.py new file mode 100644 index 0000000000000000000000000000000000000000..29b377871e353530298f5dffdcc8b9c3b3c0d103 --- /dev/null +++ b/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) diff --git a/research/scripts/shadow_split.py b/research/scripts/shadow_split.py new file mode 100644 index 0000000000000000000000000000000000000000..0e5e4aa7626b2486d2b9c15bc509ff2d34e734d3 --- /dev/null +++ b/research/scripts/shadow_split.py @@ -0,0 +1,16 @@ +"""Split a shadow CoC run into two gradeable result dirs on the SAME trajectories: +results/_full/ (generation run to the end) and results/_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") diff --git a/research/scripts/sim_coc.py b/research/scripts/sim_coc.py new file mode 100644 index 0000000000000000000000000000000000000000..6ad1ac73c48ca66154bb374d25b803aeb5edf592 --- /dev/null +++ b/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": "x" + 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}") diff --git a/research/scripts/sim_multi.py b/research/scripts/sim_multi.py new file mode 100644 index 0000000000000000000000000000000000000000..89d70bfed266b931655a24e92651d6fe3ace2c0b --- /dev/null +++ b/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("", 1)[1] if "" 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": "x" + 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}") diff --git a/research/scripts/sim_single.py b/research/scripts/sim_single.py new file mode 100644 index 0000000000000000000000000000000000000000..314a1e8d2d8f989358759a744049b4e7283ab7c5 --- /dev/null +++ b/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": "x" + 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}") diff --git a/research/scripts/soft_post.py b/research/scripts/soft_post.py new file mode 100644 index 0000000000000000000000000000000000000000..956a544c119bd4914b33f78d3007c555de20245f --- /dev/null +++ b/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": "" + tok.decode(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}") diff --git a/research/scripts/soft_stats.py b/research/scripts/soft_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..005a5a2208e00a386cd0236d1497ffd6b321cdf1 --- /dev/null +++ b/research/scripts/soft_stats.py @@ -0,0 +1,18 @@ +"""Summary of the latent-thinking runs (results/st2_): 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}") diff --git a/research/scripts/soft_think.py b/research/scripts/soft_think.py new file mode 100644 index 0000000000000000000000000000000000000000..bfb738cbab56d9b96673857a89ec6ebe231ec0b5 --- /dev/null +++ b/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 (soft modes) / is sampled (discrete), or at the +budget, where 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("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": "" + r["thinking"] + "" + 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) diff --git a/research/scripts/start_server.sh b/research/scripts/start_server.sh new file mode 100644 index 0000000000000000000000000000000000000000..77749a9b173a2b14f50b7914750fa8c30a67c750 --- /dev/null +++ b/research/scripts/start_server.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# usage: start_server.sh [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 diff --git a/research/scripts/stress.py b/research/scripts/stress.py new file mode 100644 index 0000000000000000000000000000000000000000..af58d11a1812bdba31216519175fc5820c398627 --- /dev/null +++ b/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("", 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 "" not in rec["completion"]: + return rec + final = rec["completion"].split("", 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"] = "(stress-tested)" + 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") diff --git a/research/scripts/svg_program.py b/research/scripts/svg_program.py new file mode 100644 index 0000000000000000000000000000000000000000..0c571c4b7c35e9bf2aad1e4f78d80453db67ca78 --- /dev/null +++ b/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("", 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"", out, re.S) + if not m: + return None, "It printed no ... 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": "(program)" + 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") diff --git a/research/scripts/teach_gen.py b/research/scripts/teach_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..2813d39237ec81869a870176b6529cba188cc683 --- /dev/null +++ b/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 (+ ) 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 = "" in text and correct(k, text.split("", 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") diff --git a/research/scripts/teach_probe.py b/research/scripts/teach_probe.py new file mode 100644 index 0000000000000000000000000000000000000000..cfb83d3d85155870f5eedbfd5a9e7c23c0ccc8ac --- /dev/null +++ b/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") diff --git a/research/scripts/teach_to_rl.py b/research/scripts/teach_to_rl.py new file mode 100644 index 0000000000000000000000000000000000000000..8c2270c91cf667d537e603772cdb01ede16950af --- /dev/null +++ b/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") diff --git a/research/scripts/tir.py b/research/scripts/tir.py new file mode 100644 index 0000000000000000000000000000000000000000..b591e4b24dd4b55f21a6ab0b0a1ff981a8d35ddc --- /dev/null +++ b/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 = "" + all turns (thinking, calls, outputs) + "" + 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 = "\n\n\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"\s*\s*\n?(.*?)\n?", 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]), "", "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("", 1)[1]) if "" 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 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": [""], + "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("", "\n") + FORCE_PREFIX + code + + "\n\n\n\n\n" + out + "\n\n") + ids = ids + gen[:k] + pre + tok(r2["content"].rstrip("\n") + "\n\n\n" + "<|im_end|>\n<|im_start|>tool\n\n" + out + + "\n<|im_end|>\n<|im_start|>assistant\n", + 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("", 1)[0] + "" + parts.append(turn.replace("", "\n", 1) + "\n\n" + out + "\n\n") + ids = ids + gen + tok("<|im_end|>\n<|im_start|>tool\n\n" + out + + "\n<|im_end|>\n<|im_start|>assistant\n", + add_special_tokens=False)["input_ids"] + last = parts[-1] + if "" in last: + think, final = last.split("", 1) + completion = "" + "".join(parts[:-1]) + think + "" + final + else: + completion = "" + "".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)) diff --git a/research/scripts/vis_autofix.py b/research/scripts/vis_autofix.py new file mode 100644 index 0000000000000000000000000000000000000000..7182bcb4b4c7b30a4ea8f14309ff73843b04b245 --- /dev/null +++ b/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 . + +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 = ('') +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"]*>", src, re.I) + if m: + return src[:m.end()] + "\n" + IMPORT_MAP + src[m.end():] + m = re.search(r"" in c: + head, fin = c.split("", 1) + new, fixes = autofix(fin, tasks[r["id"]]["kind"]) + if fixes: + r = {**r, "completion": head + "" + new, "autofix": fixes} + n += 1 + f.write(json.dumps(r, ensure_ascii=False) + "\n") + print(f"{n} answers changed -> {a.out}") diff --git a/research/scripts/visual_eval.py b/research/scripts/visual_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..5e33d69723089e57e10a5c239723a4265636dcf7 --- /dev/null +++ b/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//png/ for side-by-side review. +usage: visual_eval.py --tag vis_base (records results//vis.jsonl from gen.py / posthoc) + visual_eval.py --tag vis_claude --dir data/visual/claude (reference files .svg / .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("", 1)[1] if "" 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") diff --git a/research/scripts/visual_eyes.py b/research/scripts/visual_eyes.py new file mode 100644 index 0000000000000000000000000000000000000000..a97c33d751730fbc3fab82df4859ee7ec6728671 --- /dev/null +++ b/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("", 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".*?", "", 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("", 1)[1] if "" 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"] = "(eyes)" + 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") diff --git a/research/scripts/visual_loop.py b/research/scripts/visual_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..dd631e82b37b4056184899b96ac074ca523d6d93 --- /dev/null +++ b/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 , 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 ') +# 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 , 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 ') +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"]*>", src, re.I) + if m: + return src[:m.end()] + "\n" + IMPORT_MAP + src[m.end():] + m = re.search(r"" in c: + head, fin = c.split("", 1) + new, fixes = autofix(fin, tasks[r["id"]]["kind"]) + if fixes: + r = {**r, "completion": head + "" + new, "autofix": fixes} + n += 1 + f.write(json.dumps(r, ensure_ascii=False) + "\n") + print(f"{n} answers changed -> {a.out}") diff --git a/runtime/visual_loop.py b/runtime/visual_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..dd631e82b37b4056184899b96ac074ca523d6d93 --- /dev/null +++ b/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 , 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