| |
| """Task evaluation for a trained FLAME Rust model: does its output actually work? |
| |
| Perplexity (`flame eval`) measures how well a model predicts text. This script measures |
| two tasks, checked automatically with the Rust compiler: |
| |
| complete The prompt is a doc comment + function signature. The model writes the body. |
| Score: does the function compile, and does it pass the task's asserts? |
| fix The prompt is a real rustc error for a broken program, cut off at `help:`. |
| Score: does the model's help text contain the known fix (e.g. `.clone()`)? |
| |
| Every task is sampled with several seeds. The report gives, per task and overall: |
| rate share of all samples that passed (how often the model is right) |
| any@k share of tasks where at least one of k seeds passed |
| |
| Usage (from a folder that holds the model package: runs/, data/tokenizers/, target/): |
| python3 tools/task_eval.py |
| python3 tools/task_eval.py --root ~/flame-model --seeds 5 --out report.json |
| |
| Needs `rustc` on PATH (or --rustc). Generation runs on the CPU unless --gpu is given. |
| The "fix" score is a substring match: it can miss a correct fix worded differently, |
| and a match is not proof the model understood. Read the saved samples too. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| from pathlib import Path |
|
|
| DEFAULT_RUN = "runs/rust-30m-dgx-20261102-stage2" |
| DEFAULT_TASKS = Path(__file__).resolve().parent / "eval" / "rust_tasks.json" |
|
|
|
|
| |
|
|
| def generate(args, prompt: str, seed: int) -> str: |
| """Run `flame generate` once and return only the continuation after the prompt.""" |
| cmd = [args.flame, "generate", "--run", args.run, "--prompt", prompt, |
| "--max-tokens", str(args.max_tokens), "--temperature", str(args.temperature), |
| "--seed", str(seed)] |
| if not args.gpu: |
| cmd.append("--cpu") |
| res = subprocess.run(cmd, cwd=args.root, capture_output=True, text=True, timeout=600) |
| if res.returncode != 0: |
| raise RuntimeError(f"flame generate failed:\n{res.stderr.strip()}") |
| out = res.stdout |
| if out.endswith("\n"): |
| out = out[:-1] |
| |
| |
| return out[len(prompt):] if out.startswith(prompt) else out |
|
|
|
|
| |
|
|
| def run_rustc(args, source: str, workdir: Path, binary: bool) -> subprocess.CompletedProcess: |
| src = workdir / "src" / "main.rs" |
| src.parent.mkdir(parents=True, exist_ok=True) |
| src.write_text(source) |
| out = workdir / "prog" |
| cmd = [args.rustc, "--edition", "2021", "--color", "never", "--crate-name", "eval", |
| "-o", str(out) if binary else os.devnull, "src/main.rs"] |
| return subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, timeout=120) |
|
|
|
|
| def close_function(continuation: str) -> str | None: |
| """Cut the continuation where the function body (opened by the prompt's `{`) closes. |
| |
| Tracks braces outside strings and `//` comments. Returns None if it never closes. |
| """ |
| depth, i, n = 1, 0, len(continuation) |
| while i < n: |
| c = continuation[i] |
| if c == "/" and continuation.startswith("//", i): |
| j = continuation.find("\n", i) |
| i = n if j < 0 else j |
| continue |
| if c == '"': |
| i += 1 |
| while i < n and continuation[i] != '"': |
| i += 2 if continuation[i] == "\\" else 1 |
| elif c == "'" and i + 2 < n and continuation[i + 2] == "'": |
| i += 2 |
| elif c == "{": |
| depth += 1 |
| elif c == "}": |
| depth -= 1 |
| if depth == 0: |
| return continuation[: i + 1] |
| i += 1 |
| return None |
|
|
|
|
| def error_prompt(stderr: str) -> str: |
| """First rustc error block, cut right after its first `help:` (or with `help:` added).""" |
| lines = stderr.splitlines() |
| start = next(i for i, l in enumerate(lines) if l.startswith("error")) |
| block = [] |
| for line in lines[start:]: |
| if block and not line.strip(): |
| break |
| block.append(line) |
| text = "\n".join(block) |
| cut = text.find("help:") |
| if cut >= 0: |
| return text[: cut + len("help:")] |
| return text.rstrip() + "\nhelp:" |
|
|
|
|
| |
|
|
| def eval_complete(args, task, seed, workdir): |
| prompt = task["prompt"] |
| cont = generate(args, prompt, seed) |
| body = close_function(cont) |
| rec = {"output": cont} |
| if body is None: |
| rec.update(passed=False, stage="unclosed", detail="function body never closed") |
| return rec |
| program = "#![allow(warnings)]\n" + prompt + body + "\n\n" + task["test"] + "\n" |
| rec["program"] = program |
| comp = run_rustc(args, program, workdir, binary=True) |
| if comp.returncode != 0: |
| first = next((l for l in comp.stderr.splitlines() if l.startswith("error")), "") |
| rec.update(passed=False, stage="compile", detail=first) |
| return rec |
| try: |
| run = subprocess.run([str(workdir / "prog")], capture_output=True, text=True, |
| timeout=10) |
| except subprocess.TimeoutExpired: |
| rec.update(passed=False, stage="test", detail="timed out") |
| return rec |
| ok = run.returncode == 0 |
| err = run.stderr.strip().splitlines() |
| why = next((l for l in err if "assert" in l or "left:" in l), err[0] if err else "") |
| rec.update(passed=ok, stage="pass" if ok else "test", detail="" if ok else why.strip()) |
| return rec |
|
|
|
|
| def eval_fix(args, task, seed, prompt): |
| cont = generate(args, prompt, seed) |
| |
| head = "\n".join(cont.splitlines()[: args.fix_lines]) |
| hit = next((s for s in task["accept"] if s in head), None) |
| return {"output": cont, "judged": head, "passed": hit is not None, |
| "stage": "pass" if hit else "wrong", "detail": hit or ""} |
|
|
|
|
| def build_fix_prompt(args, task, workdir) -> str: |
| comp = run_rustc(args, task["source"], workdir, binary=False) |
| code = task.get("code") |
| if comp.returncode == 0 or (code and f"error[{code}]" not in comp.stderr): |
| raise RuntimeError(f"task {task['id']}: broken program did not give {code}:\n" |
| f"{comp.stderr}") |
| return error_prompt(comp.stderr) |
|
|
|
|
| |
|
|
| def main() -> int: |
| p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| p.add_argument("--root", default=".", help="folder holding runs/ and data/ (default: .)") |
| p.add_argument("--run", default=DEFAULT_RUN, help=f"run directory (default: {DEFAULT_RUN})") |
| p.add_argument("--flame", help="flame binary (default: <root>/target/release/flame)") |
| p.add_argument("--rustc", default=shutil.which("rustc") or |
| str(Path.home() / ".cargo/bin/rustc")) |
| p.add_argument("--tasks", default=str(DEFAULT_TASKS)) |
| p.add_argument("--only", choices=["complete", "fix"], help="run one task type") |
| p.add_argument("--seeds", type=int, default=3, help="samples per task (default: 3)") |
| p.add_argument("--max-tokens", type=int, default=120) |
| p.add_argument("--temperature", type=float, default=0.5) |
| p.add_argument("--fix-lines", type=int, default=4, |
| help="lines after `help:` that are judged (default: 4)") |
| p.add_argument("--gpu", action="store_true", help="do not pass --cpu to flame") |
| p.add_argument("--out", help="save every sample and the scores as JSON") |
| args = p.parse_args() |
|
|
| args.root = str(Path(args.root).expanduser().resolve()) |
| args.flame = str(Path(args.flame).expanduser().resolve()) if args.flame else \ |
| str(Path(args.root) / "target/release/flame") |
| for path, what in [(args.flame, "flame binary"), (args.rustc, "rustc"), |
| (Path(args.root) / args.run / "config.toml", "trained run")]: |
| if not Path(path).exists(): |
| sys.exit(f"error: {what} not found at {path}") |
|
|
| tasks = json.loads(Path(args.tasks).read_text())["tasks"] |
| if args.only: |
| tasks = [t for t in tasks if t["kind"] == args.only] |
| seeds = list(range(1, args.seeds + 1)) |
| print(f"FLAME task eval: {len(tasks)} tasks x {len(seeds)} seeds, run {args.run}") |
|
|
| results, started = [], time.time() |
| with tempfile.TemporaryDirectory(prefix="flame-eval-") as tmp: |
| workdir = Path(tmp) |
| for task in tasks: |
| prompt = build_fix_prompt(args, task, workdir) if task["kind"] == "fix" else None |
| samples = [] |
| for seed in seeds: |
| if task["kind"] == "complete": |
| rec = eval_complete(args, task, seed, workdir) |
| else: |
| rec = eval_fix(args, task, seed, prompt) |
| rec["seed"] = seed |
| samples.append(rec) |
| passed = sum(s["passed"] for s in samples) |
| stages = ", ".join(s["stage"] for s in samples) |
| print(f" [{task['kind']:8}] {task['id']:22} {passed}/{len(samples)} ({stages})", flush=True) |
| results.append({"id": task["id"], "kind": task["kind"], |
| "prompt": prompt or task["prompt"], "samples": samples}) |
|
|
| print(f"\nScores ({time.time() - started:.0f} s):") |
| summary = {} |
| for kind in ["complete", "fix"]: |
| rs = [r for r in results if r["kind"] == kind] |
| if not rs: |
| continue |
| total = sum(len(r["samples"]) for r in rs) |
| good = sum(s["passed"] for r in rs for s in r["samples"]) |
| anyk = sum(any(s["passed"] for s in r["samples"]) for r in rs) |
| summary[kind] = {"tasks": len(rs), "samples": total, "passed": good, |
| "rate": good / total, f"any@{len(seeds)}": anyk / len(rs)} |
| label = "code completion (compiles + tests pass)" if kind == "complete" \ |
| else "error fixing (known fix in help text)" |
| print(f" {label}: {good}/{total} samples = {good / total:.0%}, " |
| f"any@{len(seeds)} = {anyk}/{len(rs)} tasks") |
| if kind == "complete": |
| comp = sum(s["stage"] in ("pass", "test") for r in rs for s in r["samples"]) |
| print(f" compiled at all: {comp}/{total} = {comp / total:.0%}") |
|
|
| if args.out: |
| report = {"run": args.run, "seeds": seeds, "max_tokens": args.max_tokens, |
| "temperature": args.temperature, "summary": summary, "tasks": results} |
| Path(args.out).write_text(json.dumps(report, indent=2)) |
| print(f"\nAll samples saved to {args.out}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
