src/xscript/cli.py

"""`xscript` command-line entry point.

Pipeline order (see README):
  flores-download -> byte-premium
  tok-corpus -> tok-train -> tok-analyze            (the tokenizer gate)
  pool -> pack                                       (per language x chosen tok)
  train                                              (one run of the matrix)
  eval-bpb / eval-align -> bts                       (headline analysis)

Heavy steps are meant to run inside Slurm jobs (see slurm/); the CLI is the
single interface those jobs call, so behaviour is identical locally and on the
compute nodes.
"""
import argparse

from .langs import (LANGS, TOK_FLAVORS, TOK_CONDITIONS, MODEL_FLAVORS,
                    tok_name, tok_conditions)


def _add(sub, name, help):
    p = sub.add_parser(name, help=help)
    return p


def main(argv=None):
    ap = argparse.ArgumentParser(prog="xscript", description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    # ---- data prep ----
    p = _add(sub, "flores-download", "download FLORES+ dev/devtest (needs HF_TOKEN)")
    p.add_argument("--langs", nargs="*", default=list(LANGS))

    _add(sub, "byte-premium", "compute FLORES+ byte premiums (+ compare Arnett)")

    p = _add(sub, "tok-corpus", "build raw FineWeb/FineWeb2 tokenizer-training corpora")
    p.add_argument("condition", choices=TOK_CONDITIONS + ["both"])
    p.add_argument("--gb", type=float, default=4.0, help="target corpus size (GB)")

    p = _add(sub, "tok-train", "train tokenizer(s): unigram/bpe/pa")
    p.add_argument("--flavor", choices=TOK_FLAVORS + ["all"], default="all")
    p.add_argument("--condition", choices=TOK_CONDITIONS + ["both"], default="both")

    p = _add(sub, "tok-analyze", "fertility / allocation gate on FLORES+")
    p.add_argument("--toks", nargs="*", default=None)

    # ---- model-data prep ----
    p = _add(sub, "pool", "build FineWeb(-2)-HQ text pool for a language")
    p.add_argument("--lang", required=True, choices=list(LANGS))
    p.add_argument("--gb", type=float, default=None, help="override byte budget (GB)")

    p = _add(sub, "pack", "tokenize a pool into uint16 shards")
    p.add_argument("--lang", required=True, choices=list(LANGS))
    p.add_argument("--tok", required=True)
    p.add_argument("--workers", type=int, default=8)

    _add(sub, "plan", "print per-language pool budgets and the run matrix")

    p = _add(sub, "runs", "list generated run names")
    p.add_argument("--base", default="configs/base_main.yaml")
    p.add_argument("--flavor", default="unigram", choices=MODEL_FLAVORS)
    p.add_argument("--only-30b", action="store_true",
                   help="list 18 independent 30B runs (no extension trunks)")

    # ---- training ----
    p = _add(sub, "train", "train one run of the matrix")
    p.add_argument("name", help="run name (see `xscript runs`)")
    p.add_argument("--base", default="configs/base_main.yaml")
    p.add_argument("--flavor", default="unigram", choices=MODEL_FLAVORS)
    p.add_argument("--only-30b", action="store_true",
                   help="use a self-contained 30B WSD config, never a trunk branch")
    p.add_argument("--output-name", default=None,
                   help="store an independent diagnostic replicate under this run name")
    p.add_argument("--seed", type=int, default=None,
                   help="override model/optimizer RNG seed for a diagnostic replicate")
    p.add_argument("--data-seed", type=int, default=None,
                   help="override packed-stream order seed for a diagnostic replicate")
    p.add_argument("--wandb-id", default=None,
                   help="override the stable W&B run ID (useful for a clean replacement run)")

    # ---- eval ----
    p = _add(sub, "eval-bpb", "re-evaluate a checkpoint's BPB")
    p.add_argument("name"); p.add_argument("--tok", required=True)
    p.add_argument("--tag", default="final")

    p = _add(sub, "eval-align", "MEXA alignment for a run")
    p.add_argument("name"); p.add_argument("--tok", required=True)
    p.add_argument("--split", default="dev")

    p = _add(sub, "eval-bench", "downstream benchmarks (Global-MMLU/Belebele/XNLI) via lm-eval-harness")
    p.add_argument("name"); p.add_argument("--tok", required=True)
    p.add_argument("--tag", default="final")
    p.add_argument("--tasks", nargs="*", default=None,
                   help="override tasks; default is all three benchmarks for the run's languages")
    p.add_argument("--num-fewshot", type=int, default=0)
    p.add_argument("--limit", type=float, default=None,
                   help="cap examples/task (for quick smoke checks)")
    p.add_argument("--batch-size", type=int, default=4,
                   help="likelihood requests per GPU batch")
    p.add_argument("--no-wandb", action="store_true")

    p = _add(sub, "bts", "compute BTS + interaction across runs")
    p.add_argument("--flavor", default="unigram", choices=MODEL_FLAVORS)
    p.add_argument("--source", default="flores", choices=["flores", "holdout"])

    args = ap.parse_args(argv)
    return _dispatch(args)


def _dispatch(args):
    cmd = args.cmd
    if cmd == "flores-download":
        from . import flores
        flores.download(args.langs)

    elif cmd == "byte-premium":
        from . import byte_premium
        byte_premium.run()

    elif cmd == "tok-corpus":
        from .data import tokcorpus
        conds = TOK_CONDITIONS if args.condition == "both" else [args.condition]
        for c in conds:
            if c == "starved":
                tokcorpus.build_starved(total_bytes=args.gb * 1e9)
            else:
                tokcorpus.build_destarved(total_bytes=args.gb * 1e9)

    elif cmd == "tok-train":
        from .tok import train as toktrain
        flavors = TOK_FLAVORS if args.flavor == "all" else [args.flavor]
        want = TOK_CONDITIONS if args.condition == "both" else [args.condition]
        for f in flavors:
            for c in want:
                if c not in tok_conditions(f):
                    continue  # pa has no starved condition
                print(f"[tok-train] {tok_name(f, c)}")
                toktrain.train(f, c)

    elif cmd == "tok-analyze":
        from .tok import analyze
        analyze.run(args.toks)

    elif cmd == "pool":
        from .data import fineweb
        budget = (args.gb * 1e9) if args.gb else fineweb.plan_budgets()[args.lang]
        fineweb.build_pool(args.lang, budget)

    elif cmd == "pack":
        from .data import pack
        pack.pack(args.lang, args.tok, workers=args.workers)

    elif cmd == "plan":
        _plan()

    elif cmd == "runs":
        from . import runmatrix
        for n in runmatrix.list_runs(args.base, args.flavor, args.only_30b):
            print(n)

    elif cmd == "train":
        from . import runmatrix, train
        cfg = runmatrix.get_run(args.base, args.flavor, args.name, args.only_30b)
        if args.output_name is not None:
            cfg["name"] = args.output_name
        if args.seed is not None:
            cfg["seed"] = args.seed
        if args.data_seed is not None:
            cfg["data_seed"] = args.data_seed
        if args.wandb_id is not None:
            cfg["wandb_id"] = args.wandb_id
        train.run_from_config(cfg)

    elif cmd == "eval-bpb":
        from .eval import bpb
        bpb.run(args.name, args.tok, args.tag)

    elif cmd == "eval-align":
        from .eval import alignment
        alignment.run(args.name, args.tok, args.split)

    elif cmd == "eval-bench":
        from .eval import bench
        bench.run(args.name, args.tok, args.tag, tasks=args.tasks,
                  num_fewshot=args.num_fewshot, limit=args.limit,
                  log_wandb=not args.no_wandb, batch_size=args.batch_size)

    elif cmd == "bts":
        from .eval import bts
        bts.run(args.flavor, args.source)


def _plan():
    from .data.fineweb import plan_budgets
    from . import runmatrix
    b = plan_budgets()
    print("Per-language pool byte budgets (worst-case, destarved tokenizer):")
    for l, v in b.items():
        print(f"  {l}: {v/1e9:.1f} GB")
    print("\nRun matrix (flavor=unigram):")
    from . import _yaml
    base = _yaml.load("configs/base_main.yaml")
    for n in sorted(runmatrix.all_runs(base, "unigram")):
        print(f"  {n}")


if __name__ == "__main__":
    main()