dataset/generate.py

#!/usr/bin/env python3
"""Generate the code -> Mermaid-flowchart fine-tuning dataset.

Every emitted example is hard-validated:
  * engine.validate_example  -> Mermaid well-formedness, label rules, linemap accuracy
  * Python `compile()`        -> every Python source is syntactically valid
  * `node --check` (sampled)  -> JavaScript source is syntactically valid

Output is JSONL in chat ("messages") format:
    {"messages": [ {system}, {user: line-numbered code}, {assistant: thinking+graph+linemap} ]}

Usage:
    python generate.py --n 1500 --val-frac 0.08 --seed 7
    python generate.py --selftest          # exercise every template, exhaustive syntax check
"""

from __future__ import annotations

import argparse
import json
import os
import random
import subprocess
import sys
import tempfile
from collections import Counter

from engine import ValidationError, validate_example
from system_prompt import SYSTEM_PROMPT
from templates import TEMPLATES

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(HERE)

# "mainly Python and JavaScript", with C/C++ as a meaningful minority.
LANG_WEIGHT = {"python": 0.36, "javascript": 0.28, "cpp": 0.20, "c": 0.16}


# --------------------------------------------------------------------------- #
# Selection
# --------------------------------------------------------------------------- #


def pick_lang(rng, langs):
    if len(langs) == 1:
        return langs[0]
    weights = [LANG_WEIGHT[l] for l in langs]
    return rng.choices(langs, weights=weights, k=1)[0]


def gen_one(rng):
    fn, langs, _ = rng.choices(TEMPLATES, weights=[t[2] for t in TEMPLATES], k=1)[0]
    from pools import Lang
    lang = pick_lang(rng, langs)
    return fn(rng, Lang(lang))


# --------------------------------------------------------------------------- #
# Syntax checks
# --------------------------------------------------------------------------- #


def py_syntax_ok(source: str):
    try:
        compile(source, "<generated>", "exec")
        return True, ""
    except SyntaxError as e:
        return False, f"{e.msg} (line {e.lineno})"


def _run_check(cmd, source, suffix):
    with tempfile.NamedTemporaryFile("w", suffix=suffix, delete=False) as fh:
        fh.write(source)
        path = fh.name
    try:
        res = subprocess.run(cmd + [path], capture_output=True, text=True, timeout=30)
        err = res.stderr.strip().splitlines()[-1] if res.stderr.strip() else ""
        return res.returncode == 0, err
    finally:
        os.unlink(path)


def js_syntax_ok(source: str):
    return _run_check(["node", "--check"], source, ".js")


def _find_libcxx():
    """Locate a libc++ <vector> header so STL C++ can be syntax-checked."""
    roots = ["/Library/Developer/CommandLineTools/SDKs",
             "/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs"]
    for root in roots:
        if not os.path.isdir(root):
            continue
        for sdk in sorted(os.listdir(root), reverse=True):
            v1 = os.path.join(root, sdk, "usr/include/c++/v1")
            if os.path.isfile(os.path.join(v1, "vector")):
                return os.path.join(root, sdk), v1
    return None, None


_CPP_SDK, _CPP_V1 = _find_libcxx()


def c_syntax_ok(source: str):
    return _run_check(["clang", "-fsyntax-only", "-Wno-everything", "-x", "c"], source, ".c")


def cpp_syntax_ok(source: str):
    cmd = ["clang++", "-std=c++17", "-fsyntax-only", "-Wno-everything", "-x", "c++"]
    if _CPP_V1:
        cmd += ["-nostdinc++", "-isystem", _CPP_V1, "-isysroot", _CPP_SDK]
    return _run_check(cmd, source, ".cpp")


SYNTAX_CHECK = {
    "python": py_syntax_ok,
    "javascript": js_syntax_ok,
    "cpp": cpp_syntax_ok,
    "c": c_syntax_ok,
}


# --------------------------------------------------------------------------- #
# Self-test: every template x every supported language
# --------------------------------------------------------------------------- #


def selftest(per_template: int = 6, seed: int = 0):
    from pools import Lang
    if not _CPP_V1:
        print("WARNING: libc++ headers not found; C++ will be syntax-checked without STL "
              "(may give false failures).")
    rng = random.Random(seed)
    failures = 0
    checked = Counter()
    for fn, langs, _ in TEMPLATES:
        for lang in langs:
            for _ in range(per_template):
                ex = fn(rng, Lang(lang))
                try:
                    validate_example(ex)
                except ValidationError as e:
                    failures += 1
                    print(f"  STRUCT FAIL {fn.__name__}/{lang}: {e}")
                    print(ex.output)
                    continue
                if ex.template == "error_node":
                    continue  # intentionally unparseable
                ok, err = SYNTAX_CHECK[lang](ex.source)
                checked[lang] += 1
                if not ok:
                    failures += 1
                    print(f"  SYNTAX FAIL {fn.__name__}/{lang}: {err}")
                    print(ex.source)
    print(f"selftest: {len(TEMPLATES)} templates, syntax-checked {dict(checked)}, "
          f"{failures} failures")
    return failures


# --------------------------------------------------------------------------- #
# Generation
# --------------------------------------------------------------------------- #


def to_record(ex):
    return {"messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": ex.code},
        {"role": "assistant", "content": ex.output},
    ]}


def generate(n, seed, check_sample):
    rng = random.Random(seed)
    seen = set()
    examples = []
    pools = {"javascript": [], "cpp": [], "c": []}  # sampled compiler checks
    attempts = 0
    max_attempts = n * 60

    while len(examples) < n and attempts < max_attempts:
        attempts += 1
        ex = gen_one(rng)
        if ex.code in seen:
            continue
        try:
            validate_example(ex)
        except ValidationError as e:
            raise SystemExit(f"FATAL: invalid example from {ex.template}/{ex.language}: {e}\n{ex.output}")
        if ex.template != "error_node" and ex.language == "python":
            ok, err = py_syntax_ok(ex.source)  # in-process, every example
            if not ok:
                raise SystemExit(f"FATAL: invalid python from {ex.template}: {err}\n{ex.source}")
        seen.add(ex.code)
        examples.append(ex)
        if ex.template != "error_node" and ex.language in pools:
            pools[ex.language].append(ex.source)

    if len(examples) < n:
        print(f"WARNING: only produced {len(examples)} unique examples "
              f"(requested {n}) after {attempts} attempts.")

    # sampled compiler syntax checks for the brace languages
    for lang, srcs in pools.items():
        rng.shuffle(srcs)
        sample = srcs[:check_sample]
        fails = 0
        for src in sample:
            ok, err = SYNTAX_CHECK[lang](src)
            if not ok:
                fails += 1
                print(f"  {lang} SYNTAX FAIL: {err}\n{src}")
        if fails:
            raise SystemExit(f"FATAL: {fails}/{len(sample)} sampled {lang} examples failed syntax check")
        print(f"{lang} sampled syntax check: {len(sample)} ok")

    return examples


def write_jsonl(path, records):
    with open(path, "w", encoding="utf-8") as fh:
        for r in records:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")


def write_preview(path, examples, k=6):
    rng = random.Random(123)
    picks = rng.sample(examples, min(k, len(examples)))
    out = ["# Dataset preview\n",
           f"Random sample of {len(picks)} examples (system prompt omitted for brevity).\n"]
    for i, ex in enumerate(picks, 1):
        out.append(f"## Example {i} — `{ex.template}` ({ex.language}), {ex.n_nodes} nodes\n")
        out.append("**User (input):**\n\n```\n" + ex.code + "\n```\n")
        out.append("**Assistant (target):**\n\n```\n" + ex.output + "\n```\n")
    with open(path, "w", encoding="utf-8") as fh:
        fh.write("\n".join(out))


def print_stats(examples):
    by_lang = Counter(e.language for e in examples)
    by_tpl = Counter(e.template for e in examples)
    nodes = [e.n_nodes for e in examples]
    out_chars = [len(e.output) for e in examples]
    print("\n--- dataset stats ---")
    print(f"total examples : {len(examples)}")
    print(f"by language    : {dict(by_lang)}")
    print(f"avg nodes/graph: {sum(nodes)/len(nodes):.1f}  (min {min(nodes)}, max {max(nodes)})")
    print(f"avg output len : {sum(out_chars)//len(out_chars)} chars "
          f"(~{sum(out_chars)//len(out_chars)//4} tokens)")
    print("by template    :")
    for tpl, c in sorted(by_tpl.items(), key=lambda kv: -kv[1]):
        print(f"    {tpl:18s} {c}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=1500, help="total examples (train+val)")
    ap.add_argument("--val-frac", type=float, default=0.08)
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--out-dir", default=os.path.join(REPO, "data"))
    ap.add_argument("--check-sample", type=int, default=60,
                    help="examples per brace-language to compiler syntax-check")
    ap.add_argument("--selftest", action="store_true")
    args = ap.parse_args()

    if args.selftest:
        sys.exit(1 if selftest() else 0)

    examples = generate(args.n, args.seed, args.check_sample)
    print_stats(examples)

    rng = random.Random(args.seed + 1)
    rng.shuffle(examples)
    n_val = max(1, int(len(examples) * args.val_frac))
    val, train = examples[:n_val], examples[n_val:]

    os.makedirs(args.out_dir, exist_ok=True)
    write_jsonl(os.path.join(args.out_dir, "train.jsonl"), [to_record(e) for e in train])
    write_jsonl(os.path.join(args.out_dir, "val.jsonl"), [to_record(e) for e in val])
    write_preview(os.path.join(args.out_dir, "preview.md"), examples)
    print(f"\nwrote {len(train)} -> {args.out_dir}/train.jsonl")
    print(f"wrote {len(val)} -> {args.out_dir}/val.jsonl")
    print(f"wrote preview -> {args.out_dir}/preview.md")


if __name__ == "__main__":
    main()