| """`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) |
|
|
| |
| 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) |
|
|
| |
| 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)") |
|
|
| |
| 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)") |
|
|
| |
| 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 |
| 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() |
|
|
