Spaces:
Running
Running
| """datainventor: curate a pretraining corpus on your own machines. | |
| datainventor init recipe.yaml write a starter recipe to edit | |
| datainventor preview recipe.yaml run every stage on a small sample and print what each did | |
| datainventor run recipe.yaml --out ./curated run it in full (resumable) | |
| The model endpoint is any OpenAI-compatible server: vLLM, SGLang, Ollama, a hosted API. Set its key in the | |
| environment variable named by endpoint.api_key_env (DI_API_KEY by default). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| from .compose import fmt | |
| from .recipe import DEFAULT, dump, endpoint_of, load | |
| def main(argv=None): | |
| ap = argparse.ArgumentParser(prog="datainventor", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| sub = ap.add_subparsers(dest="cmd", required=True) | |
| a = sub.add_parser("init", help="write a starter recipe") | |
| a.add_argument("path") | |
| a = sub.add_parser("preview", help="run every stage on a sample") | |
| a.add_argument("recipe") | |
| a.add_argument("--per-source", type=int, default=300) | |
| a.add_argument("--grade", type=int, default=64) | |
| a.add_argument("--rephrase", type=int, default=6) | |
| a.add_argument("--workers", type=int, default=None) | |
| a.add_argument("--json", help="also write the full report here") | |
| a = sub.add_parser("run", help="run the recipe in full") | |
| a.add_argument("recipe") | |
| a.add_argument("--out", default=None, help="output folder (default: output.dir in the recipe)") | |
| a.add_argument("--procs", type=int, default=0, help="CPU processes for cleaning and scoring (default: all but one)") | |
| a.add_argument("--workers", type=int, default=None, help="parallel model calls (default: endpoint.workers)") | |
| a.add_argument("--stages", default="grade,curate,create,compose") | |
| a.add_argument("--limit", type=int, default=0, help="read at most this many documents per source (for a trial run)") | |
| args = ap.parse_args(argv) | |
| if args.cmd == "init": | |
| if os.path.exists(args.path): | |
| sys.exit(f"{args.path} exists; not overwriting it.") | |
| r = json.loads(json.dumps(DEFAULT)) | |
| r["sources"] = [ | |
| {"name": "finepdfs", "dataset": "codelion/finepdfs-1B", "kind": "textbook"}, | |
| {"name": "dclm", "dataset": "codelion/dclm-baseline-1B", "kind": "web"}, | |
| {"name": "fineweb-edu", "dataset": "codelion/fineweb-edu-1B", "kind": "educational"}, | |
| ] | |
| r["targets"]["description"] = "grade-school to university math and science reasoning" | |
| r["compose"]["mixes"][0]["budget_tokens"] = 2_000_000_000 | |
| with open(args.path, "w") as f: | |
| f.write(dump(r)) | |
| print(f"Wrote {args.path}. Edit the sources, targets and budget, then: datainventor preview {args.path}") | |
| return | |
| r = load(args.recipe) | |
| if args.cmd == "preview": | |
| from .preview import Preview, run_preview | |
| from .sources import inspect | |
| token = os.environ.get("HF_TOKEN", "") | |
| for s in r["sources"]: # row counts make the token estimates possible | |
| if s.get("dataset") and not s.get("num_rows"): | |
| try: | |
| s["num_rows"] = inspect(s["dataset"], token)["num_rows"] if (s.get("config", "default") == "default") else 0 | |
| except Exception: | |
| pass | |
| p = Preview(r) | |
| ep = endpoint_of(r, args.workers) | |
| run_preview(p, ep, token, per_source=args.per_source, grade_n=args.grade, rephrase_n=args.rephrase, invent_n=2) | |
| if p.state != "done": | |
| sys.exit(f"Preview {p.state}: {p.message}") | |
| rep = p.report | |
| for f in rep["sources"]: | |
| drops = ", ".join(f"{k} {v}" for k, v in sorted(f["dropped"].items(), key=lambda x: -x[1])) | |
| print(f"{f['name']}: read {f['read']}, selected {f['selected']}; dropped: {drops}; decontaminated: {f['decontaminated'] or 'none'}") | |
| if f.get("available_tokens"): | |
| print(f" about {fmt(f['available_tokens'])} tokens available after curation") | |
| c = rep["stages"]["calibrate"] | |
| print(f"graded {c['graded']}; scorer agreement {c['agreement']}") | |
| for plan in rep["stages"]["compose"]["plans"]: | |
| print(f"mix {plan['name']} ({plan['budget']}): " + "; ".join(f"{x['name']} {int(x['share'] * 100)}% x{x['epochs']}" for x in plan["rows"])) | |
| for w in plan["warnings"]: | |
| print(" - " + w) | |
| if args.json: | |
| with open(args.json, "w") as fh: | |
| json.dump({"report": rep, "docs": p.docs}, fh, indent=1, ensure_ascii=False) | |
| return | |
| from .runner import Run | |
| out = args.out or r["output"]["dir"] | |
| Run(r, out, procs=args.procs, workers=args.workers, limit=args.limit).go(tuple(args.stages.split(","))) | |
| print(f"Done. See {os.path.join(out, 'REPORT.md')}") | |
| if __name__ == "__main__": | |
| main() | |