release/run_train.py

#!/usr/bin/env python3
"""Portable launcher for the bind2_0 trainer.

WHY THIS EXISTS
---------------
`src/train_bind2_0_babylm.py` line 20 is:

    WORK = os.environ["BABYLM_WORK"]

a bare subscript at module scope with no default, so merely IMPORTING the module raises an
uncaught KeyError when the variable is unset — the failure a newcomer hits first, with no message
explaining what to set. It is read at import time, not call time, so it must be set BEFORE the
import; rebinding the attribute afterwards is too late. This launcher sets it, checks the two
files it derives (tokenizer.json and tokens_u16.bin), and reports clearly if they are absent.

The trainer takes bare positional sys.argv (line 26), not argparse, so arguments are passed
through verbatim in order.

    python run_train.py --work <dir> -- <positional args for train_bind2_0_babylm.py>
    python run_train.py --work <dir> --check

HONEST LIMITATION, not worked around: line 30 hardcodes dev = "cuda" with no CPU path, and
modeling_bind2_0.py imports flash-linear-attention, whose kernels are CUDA-only. This package
cannot run on CPU or on non-NVIDIA hardware. A launcher cannot fix that; only editing the frozen
source could, and that is deliberately not done here.
"""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--work", default=os.environ.get("BABYLM_WORK", str(HERE / "data")),
                    help="directory holding tokenizer.json and tokens_u16.bin")
    ap.add_argument("--check", action="store_true", help="resolve inputs and exit")
    args, passthrough = ap.parse_known_args()
    if passthrough and passthrough[0] == "--":
        passthrough = passthrough[1:]

    work = Path(args.work)
    os.environ["BABYLM_WORK"] = str(work)          # MUST precede the import
    sys.path.insert(0, str(HERE / "src"))

    need = {"tokenizer.json": work / "tokenizer.json", "tokens_u16.bin": work / "tokens_u16.bin"}
    missing = [k for k, v in need.items() if not v.exists()]
    if args.check or missing:
        print(f"  BABYLM_WORK = {work}")
        for k, v in need.items():
            print(f"  {k:16} {'OK ' if v.exists() else 'MISSING'}  {v}")
    if missing:
        print("\nERROR: tokens_u16.bin is not redistributed inside this package (it is large and "
              "regenerable). Build it with the bind1 package's tokenizer chain "
              "(train_tokenizer.py then make_tokens.py), or point --work at a directory that "
              "already has it. See BUILD.md.", file=sys.stderr)
        return 2
    if args.check:
        print("\nAll inputs resolve. Re-run without --check to train.")
        return 0

    import train_bind2_0_babylm as T
    sys.argv = ["train_bind2_0_babylm.py"] + passthrough
    T.main()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())