codelion's picture
Curate a corpus: clean, calibrate, create, compose on a sample, plus a recipe and runner for your own machines
4f740eb verified
Raw History Blame Contribute Delete
24.7 kB
"""Run a recipe in full on your own machines.
datainventor run recipe.yaml --out ./curated --procs 16 --workers 32
Stages, each resumable (re-run the same command after an interruption):
grade grade a sample of every source with your model, fit the scorers -> calibrator.json, grades.jsonl
curate stream every document: clean, decontaminate, dedup, score -> curated/<source>/part-*.jsonl
create rephrase the best documents where a mix needs more tokens, invent -> synthetic/<source>/part-*.jsonl
compose build each mix to its token budget -> mixes/<mix>/part-*.jsonl
Then REPORT.md and report.json summarize what every stage did.
Token counts are estimates (four characters per token). Calibrate them against your tokenizer if the budget is tight.
"""
from __future__ import annotations
import itertools
import json
import os
import random
import shutil
import time
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import get_context
from . import calibrate as cal
from .clean import Cleaner, Decontaminator, benchmark_texts
from .compose import default_weights, fmt, plan_mix
from .create import chunks, invent, rephrase, styles_for
from .preview import plan_topics
from .recipe import dump, endpoint_of, held_out_of, sources_of, targets_of
from .sources import Source, stream
from .text import MinHasher, NearDupIndex, doc_hash, est_tokens, h64
STAGES = ("grade", "curate", "create", "compose")
class Shards:
"""Append-only JSONL shards of a fixed number of documents."""
def __init__(self, folder: str, per: int):
os.makedirs(folder, exist_ok=True)
self.folder, self.per, self.n, self.f, self.count = folder, per, 0, None, 0
def write(self, doc: dict):
if self.f is None or self.count >= self.per:
if self.f:
self.f.close()
self.f = open(os.path.join(self.folder, f"part-{self.n:05d}.jsonl"), "w", encoding="utf-8")
self.n += 1
self.count = 0
self.f.write(json.dumps(doc, ensure_ascii=False) + "\n")
self.count += 1
def close(self):
if self.f:
self.f.close()
self.f = None
def read_dir(folder: str):
if not os.path.isdir(folder):
return
for name in sorted(os.listdir(folder)):
if name.endswith(".jsonl"):
with open(os.path.join(folder, name), encoding="utf-8") as f:
for line in f:
if line.strip():
yield json.loads(line)
class Run:
def __init__(self, recipe: dict, out: str, procs: int = 0, workers: int | None = None, limit: int = 0, log=print):
self.r, self.out, self.limit, self.log = recipe, os.path.abspath(out), limit, log
self.procs = procs or max(1, (os.cpu_count() or 2) - 1)
self.ep = endpoint_of(recipe, workers)
os.makedirs(self.out, exist_ok=True)
self.state_path = os.path.join(self.out, "state.json")
self.state = json.load(open(self.state_path)) if os.path.exists(self.state_path) else {"done": [], "sources": {}}
with open(os.path.join(self.out, "recipe.yaml"), "w", encoding="utf-8") as f:
f.write(dump(recipe))
def save(self):
tmp = self.state_path + ".tmp"
with open(tmp, "w") as f:
json.dump(self.state, f, indent=1)
os.replace(tmp, self.state_path)
def path(self, *p) -> str:
return os.path.join(self.out, *p)
def go(self, stages=STAGES):
if not self.ep.api_key and self.r["endpoint"]["base_url"].startswith("https://router.huggingface.co"):
raise SystemExit("Set DI_API_KEY (or HF_TOKEN) for the model endpoint.")
for st in STAGES:
if st not in stages:
continue
if st in self.state["done"]:
self.log(f"[{st}] already done, skipping (delete {self.state_path} to start over)")
continue
t = time.time()
getattr(self, "stage_" + st)()
self.state["done"].append(st)
self.save()
self.log(f"[{st}] done in {time.time() - t:.0f}s")
self.write_report()
# ------------------------------------------------------------ grade
def stage_grade(self):
r = self.r
n_total = int(r["calibrate"]["grade_docs"])
srcs = sources_of(r)
per = max(4, n_total // len(srcs))
cleaner = Cleaner(r["clean"]["rules"])
pool: list[dict] = []
for s in srcs:
got = []
for d in stream(s):
d2, why = cleaner.apply(d)
if d2:
got.append(d2)
if len(got) >= per * 3:
break
pool += random.Random(1).sample(got, min(per, len(got)))
self.log(f"[grade] {s.name}: sampled {min(per, len(got))} documents")
gpath = self.path("grades.jsonl")
done = {}
if os.path.exists(gpath):
for line in open(gpath, encoding="utf-8"):
g = json.loads(line)
done[g["id"]] = g
todo = [d for d in pool if d["id"] not in done]
desc = r["targets"]["description"]
self.log(f"[grade] grading {len(todo)} documents ({len(done)} already graded) with {self.ep.model}")
with open(gpath, "a", encoding="utf-8") as f, ThreadPoolExecutor(self.ep.workers) as ex:
for i, (d, g) in enumerate(zip(todo, ex.map(lambda d: cal.grade(self.ep, d["text"], desc), todo))):
if g:
rec = {"id": d["id"], "source": d["source"], "text": d["text"][:12000], **g}
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
done[d["id"]] = rec
if (i + 1) % 200 == 0:
f.flush()
self.log(f"[grade] {i + 1}/{len(todo)}")
graded = [(g["text"], g) for g in done.values()]
target_texts = []
for t in targets_of(r):
for i, d in enumerate(stream(t)):
if i >= int(r["targets"]["max_docs"]):
break
target_texts.append(d["text"])
calib = cal.Calibrator(float(r["calibrate"]["focus"]), desc, target_texts).fit(graded, [d["text"] for d in pool])
with open(self.path("calibrator.json"), "w") as f:
json.dump(calib.to_dict(), f)
self.state["grade"] = {"graded": len(graded), "agreement": calib.agreement}
self.log(f"[grade] scorer agreement with the model on held-out documents: {calib.agreement}")
# ------------------------------------------------------------ curate
def stage_curate(self):
r = self.r
dec_keys = r["clean"]["decontaminate"]
index = NearDupIndex(MinHasher(), threshold=float(r["calibrate"]["dedup"]["threshold"])) if r["calibrate"]["dedup"]["near"] else None
exact: set = set()
ctx = get_context("fork") if hasattr(os, "fork") else get_context("spawn")
held = {}
for h in held_out_of(r):
held[h.name] = [d["text"] for _, d in zip(range(200_000), stream(h))]
self.log(f"[curate] held-out set {h.name}: {len(held[h.name]):,} items")
init = (r["clean"]["rules"], dec_keys, held, int(r["clean"].get("ngram", 13)), self.path("calibrator.json"), bool(index))
with ctx.Pool(self.procs, initializer=_worker_init, initargs=init) as pool:
for s in sources_of(r):
st = self.state["sources"].get(s.name, {})
folder = self.path("curated", s.name)
if st.get("done"):
self.log(f"[curate] {s.name}: done earlier, loading it into the dedup index")
for d in read_dir(folder):
exact.add(d["h"])
if index is not None:
for sig, d in zip(pool.imap(_sig_only, (d["text"] for d in read_dir(folder)), chunksize=64), read_dir(folder)):
index.check_add_sig(h64(d["h"]), sig)
continue
if os.path.isdir(folder):
shutil.rmtree(folder)
sh = Shards(folder, int(r["output"]["shard_docs"]))
c = {"read": 0, "kept": 0, "dropped": {}, "decontaminated": {}, "tokens": 0}
scores: list[float] = []
t0 = time.time()
llm_all = r["calibrate"].get("scorer") == "llm"
calib = cal.Calibrator.from_dict(json.load(open(self.path("calibrator.json")))) if llm_all else None
ex = ThreadPoolExecutor(self.ep.workers) if llm_all else None
desc = r["targets"]["description"]
for batch in pool.imap(_worker_batch, _batches(stream(s), 256, self.limit), chunksize=1):
if llm_all: # grade every surviving document with the model instead of the learned scorer
live = [d for d, why, _ in batch if d]
for d, g in zip(live, ex.map(lambda d: cal.grade(self.ep, d["text"], desc), live)):
if g:
d.update(calib.score(d["text"], g))
for d, why, sig in batch:
c["read"] += 1
if why:
key = "decontaminated" if why.startswith("bench:") else "dropped"
name = why[6:] if key == "decontaminated" else why
c[key][name] = c[key].get(name, 0) + 1
continue
if d["h"] in exact:
c["dropped"]["exact duplicate"] = c["dropped"].get("exact duplicate", 0) + 1
continue
exact.add(d["h"])
if index is not None and index.check_add_sig(h64(d["h"]), sig) is not None:
c["dropped"]["near duplicate"] = c["dropped"].get("near duplicate", 0) + 1
continue
sh.write(d)
c["kept"] += 1
c["tokens"] += d["tokens"]
scores.append(d["score"])
if c["read"] % 50_000 < 256:
rate = c["read"] / max(1e-6, time.time() - t0)
self.log(f"[curate] {s.name}: {c['read']:,} read, {c['kept']:,} kept, {rate:,.0f} docs/s")
sh.close()
if ex:
ex.shutdown()
scores.sort(reverse=True)
keep_top, n = float(r["calibrate"]["keep_top"]), len(scores)
c["cut"] = scores[max(0, int(n * keep_top) - 1)] if n else 1.0
c["high_cut"] = scores[max(0, int(n * 0.2) - 1)] if n else 1.0
c["selected_tokens"], c["high_tokens"] = 0, 0
for d in read_dir(folder):
if d["score"] >= c["cut"]:
c["selected_tokens"] += d["tokens"]
if d["score"] >= c["high_cut"]:
c["high_tokens"] += d["tokens"]
c["done"] = True
self.state["sources"][s.name] = c
self.save()
self.log(f"[curate] {s.name}: kept {c['kept']:,} of {c['read']:,}; {fmt(c['selected_tokens'])} tokens above the score cut")
# ------------------------------------------------------------ create
def plans(self, synthetic: dict | None = None) -> list[dict]:
rows = [{"name": s.name, "kind": s.kind, "available_tokens": self.state["sources"].get(s.name, {}).get("selected_tokens", 0),
"high_tokens": self.state["sources"].get(s.name, {}).get("high_tokens", 0)} for s in sources_of(self.r)]
defaults = default_weights(rows)
out = []
for m in self.r["compose"]["mixes"]:
w = {**defaults, **{k: float(v) for k, v in (m.get("weights") or {}).items()}}
inv = float(m.get("invented_share", 0))
p = plan_mix(int(int(m["budget_tokens"]) * (1 - inv)), rows, w, float(self.r["compose"]["max_epochs"]), float(self.r["compose"]["max_synthetic_share"]))
p["name"], p["weights"], p["invented_share"], p["budget_tokens"] = m["name"], w, inv, int(m["budget_tokens"])
out.append(p)
return out
def stage_create(self):
r = self.r
rc, ic = r["create"]["rephrase"], r["create"]["invent"]
need: dict[str, int] = {}
for p in self.plans():
for row in p["rows"]:
need[row["name"]] = max(need.get(row["name"], 0), row["synthetic_tokens"])
made = self.state.setdefault("synthetic", {})
rng = random.Random(3)
if rc["enabled"]:
for s in sources_of(r):
want = need.get(s.name, 0)
have = made.get(s.name, {}).get("tokens", 0)
if want <= have:
continue
self.log(f"[create] {s.name}: rephrasing to add {fmt(want - have)} tokens")
self._rephrase_source(s, want, rc, rng)
if ic["enabled"] and int(ic["docs"]) > 0 and ic["request"].strip():
self._invent(ic, rng)
def _rephrase_source(self, s: Source, want: int, rc: dict, rng: random.Random):
folder = self.path("synthetic", s.name)
st = self.state["synthetic"].setdefault(s.name, {"tokens": 0, "docs": 0, "flagged": 0, "seen": 0})
cut = self.state["sources"][s.name]["high_cut"]
styles = rc["styles"] or styles_for(s.kind)
def gen():
# each pass over the best documents uses a different style, so later passes add variety, not copies
k = 0
for pass_ in range(len(styles)):
for d in read_dir(self.path("curated", s.name)):
if d["score"] >= cut:
for i, ch in enumerate(chunks(d["text"])):
yield d["id"], i, ch, styles[(pass_ + k) % len(styles)]
k += 1
jobs = itertools.islice(gen(), st["seen"], None)
sh = Shards(folder, int(self.r["output"]["shard_docs"]))
sh.n = len([f for f in os.listdir(folder)]) if os.path.isdir(folder) else 0
step = self.ep.workers * 4
with ThreadPoolExecutor(self.ep.workers) as ex:
while st["tokens"] < want:
part = list(itertools.islice(jobs, step))
if not part:
break
for (oid, i, ch, style), o in zip(part, ex.map(lambda j: rephrase(self.ep, j[2], j[3], rc.get("language", ""), float(rc["max_new_share"])), part)):
st["seen"] += 1
if not o["text"]:
continue
if not o["ok"]:
st["flagged"] += 1
continue
t = est_tokens(o["text"])
sh.write({"id": f"{oid}:{i}:{style}", "text": o["text"], "source": s.name, "kind": "synthetic", "origin": oid,
"style": style, "tokens": t, "fidelity": o["fidelity"]})
st["tokens"] += t
st["docs"] += 1
sh.close()
self.save()
self.log(f"[create] {s.name}: {fmt(st['tokens'])} of {fmt(want)} tokens, {st['flagged']} flagged")
if st["tokens"] < want:
self.log(f"[create] {s.name}: ran out of documents to rephrase at {fmt(st['tokens'])} tokens")
def _invent(self, ic: dict, rng: random.Random):
st = self.state["synthetic"].setdefault("invented", {"tokens": 0, "docs": 0, "flagged": 0, "seen": 0})
n = int(ic["docs"]) - st["docs"]
if n <= 0:
return
topics = plan_topics(self.ep, ic["request"], min(200, max(10, n // 5)))
jobs = [(rng.choice(topics), rng.choice(ic["doc_types"]), rng.choice(ic["levels"])) for _ in range(n)]
sh = Shards(self.path("synthetic", "invented"), int(self.r["output"]["shard_docs"]))
sh.n = len(os.listdir(self.path("synthetic", "invented")))
self.log(f"[create] inventing {n} documents over {len(topics)} topics")
with ThreadPoolExecutor(self.ep.workers) as ex:
for (t, dt, lv), o in zip(jobs, ex.map(lambda j: invent(self.ep, j[0], j[1], j[2], ic["request"], ic.get("language", "English")), jobs)):
if o.get("text"):
tk = est_tokens(o["text"])
sh.write({"id": f"invented:{st['docs']}", "text": o["text"], "source": "invented", "kind": "synthetic",
"topic": t, "doc_type": dt, "level": lv, "tokens": tk})
st["docs"] += 1
st["tokens"] += tk
sh.close()
# ------------------------------------------------------------ compose
def stage_compose(self):
r = self.r
synth = self.state.get("synthetic", {})
results = []
for p in self.plans():
name = p["name"]
folder = self.path("mixes", name)
if os.path.isdir(folder):
shutil.rmtree(folder)
sh = Shards(folder, int(r["output"]["shard_docs"]))
buckets = []
for row in p["rows"]:
s = row["name"]
cut = self.state["sources"].get(s, {}).get("cut", 1.0)
real_target = row["want_tokens"] - min(row["synthetic_tokens"], synth.get(s, {}).get("tokens", 0))
buckets.append(_Bucket(s, self.path("curated", s), cut, real_target, float(r["compose"]["max_epochs"]),
self.state["sources"].get(s, {}).get("selected_tokens", 0)))
if row["synthetic_tokens"] and synth.get(s, {}).get("tokens"):
buckets.append(_Bucket(s + ":synthetic", self.path("synthetic", s), -1, min(row["synthetic_tokens"], synth[s]["tokens"]), 1.0,
synth[s]["tokens"]))
share = float(p.get("invented_share", 0))
if share and synth.get("invented", {}).get("tokens"):
buckets.append(_Bucket("invented", self.path("synthetic", "invented"), -1, int(p["budget_tokens"] * share), 1.0,
synth["invented"]["tokens"]))
rng = random.Random(11)
buf: list[dict] = []
live = [b for b in buckets if b.target > 0]
while live:
# draw the next bucket in proportion to the tokens it still owes, which interleaves sources evenly
b = rng.choices(live, weights=[x.target - x.taken for x in live], k=1)[0]
d = b.next()
if d is None:
live.remove(b)
continue
buf.append({"id": d["id"], "text": d["text"], "source": d["source"], "kind": d.get("kind", "")})
if len(buf) >= 20_000:
rng.shuffle(buf)
for x in buf[:10_000]:
sh.write(x)
buf = buf[10_000:]
if b.taken >= b.target:
live.remove(b)
rng.shuffle(buf)
for x in buf:
sh.write(x)
sh.close()
manifest = {"name": name, "budget_tokens": p["budget_tokens"], "weights": p["weights"],
"buckets": [{"name": b.name, "target_tokens": b.target, "tokens": b.taken, "docs": b.docs, "epochs": b.epochs_used()} for b in buckets],
"tokens": sum(b.taken for b in buckets), "warnings": p["warnings"]}
with open(os.path.join(folder, "manifest.json"), "w") as f:
json.dump(manifest, f, indent=1)
results.append(manifest)
self.log(f"[compose] {name}: {fmt(manifest['tokens'])} tokens of {fmt(p['budget_tokens'])}")
self.state["mixes"] = results
# ------------------------------------------------------------ report
def write_report(self):
st = self.state
lines = [f"# Curation report: {self.r['name']}", "", f"Model: `{self.ep.model}`. Stages done: {', '.join(st['done']) or 'none'}.", ""]
if st.get("grade"):
a = st["grade"]["agreement"]
lines += [f"Graded {st['grade']['graded']:,} documents. Scorer rank agreement with the model on held-out documents: "
f"quality {a.get('quality')}, relevance {a.get('relevance')}.", ""]
if st["sources"]:
lines += ["## Sources", "", "| Source | Read | Kept | Decontaminated | Top drop reasons | Tokens above cut |", "|---|---|---|---|---|---|"]
for n, c in st["sources"].items():
top = ", ".join(f"{k} {v:,}" for k, v in sorted(c["dropped"].items(), key=lambda x: -x[1])[:3])
dec = sum(c["decontaminated"].values())
lines.append(f"| {n} | {c['read']:,} | {c['kept']:,} | {dec:,} | {top} | {fmt(c.get('selected_tokens', 0))} |")
lines.append("")
if st.get("synthetic"):
lines += ["## Created", "", "| From | Documents | Tokens | Flagged |", "|---|---|---|---|"]
for n, c in st["synthetic"].items():
lines.append(f"| {n} | {c['docs']:,} | {fmt(c['tokens'])} | {c['flagged']:,} |")
lines.append("")
for m in st.get("mixes", []):
lines += [f"## Mix `{m['name']}`: {fmt(m['tokens'])} of {fmt(m['budget_tokens'])} tokens", "",
"| Bucket | Target | Written | Epochs |", "|---|---|---|---|"]
for b in m["buckets"]:
lines.append(f"| {b['name']} | {fmt(b['target_tokens'])} | {fmt(b['tokens'])} | {b['epochs']} |")
lines += [""] + [f"- {w}" for w in m["warnings"]] + [""]
with open(self.path("REPORT.md"), "w") as f:
f.write("\n".join(lines))
with open(self.path("report.json"), "w") as f:
json.dump(st, f, indent=1)
class _Bucket:
"""Documents of one source above its score cut, cycled up to max_epochs times."""
def __init__(self, name: str, folder: str, cut: float, target: int, max_epochs: float, pass_tokens: int = 0):
self.name, self.folder, self.cut, self.target, self.max_epochs = name, folder, cut, int(target), max_epochs
self.pass_tokens = pass_tokens
self.taken, self.docs, self.epoch, self.it, self.epoch_tokens = 0, 0, 0, None, 0
def next(self):
while True:
if self.it is None:
if self.epoch >= max(1, int(self.max_epochs + 0.999)):
return None
self.it = iter(read_dir(self.folder))
self.epoch += 1
for d in self.it:
if d.get("score", 1.0) >= self.cut:
self.taken += d.get("tokens") or est_tokens(d["text"])
self.docs += 1
return d
self.it = None
if self.taken == 0:
return None
def epochs_used(self) -> float:
return round(self.taken / self.pass_tokens, 2) if self.pass_tokens else float(self.epoch)
# ---------------------------------------------------------------- worker processes
_W: dict = {}
def _worker_init(rules, dec_keys, held, ngram, calib_path, near):
_W["clean"] = Cleaner(rules)
dec = Decontaminator(ngram)
for k in dec_keys:
dec.add(k, benchmark_texts(k))
for k, texts in held.items():
dec.add(k, texts)
_W["dec"] = dec
_W["calib"] = cal.Calibrator.from_dict(json.load(open(calib_path)))
_W["mh"] = MinHasher() if near else None
def _worker_batch(batch: list[dict]) -> list[tuple]:
out = []
for d in batch:
d2, why = _W["clean"].apply(d)
if not d2:
out.append((None, why, None))
continue
b = _W["dec"].hit(d2["text"])
if b:
out.append((None, "bench:" + b, None))
continue
sc = _W["calib"].score(d2["text"])
rec = {**d2, **sc, "tokens": est_tokens(d2["text"]), "h": doc_hash(d2["text"])}
out.append((rec, "", _W["mh"].signature(d2["text"]) if _W["mh"] else None))
return out
def _sig_only(text: str):
return _W["mh"].signature(text) if _W.get("mh") else None
def _batches(it, n: int, limit: int = 0):
buf, seen = [], 0
for d in it:
buf.append(d)
seen += 1
if len(buf) >= n:
yield buf
buf = []
if limit and seen >= limit:
break
if buf:
yield buf