Spaces:
Running
Running
File size: 4,953 Bytes
4f740eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """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()
|