Spaces:
Running
Running
Download datainventor/runner.py from codelion/data-inventor: direct link, hf CLI and curl.
- Browser
- Download file 24.7 kB
-
https://huggingface.co/spaces/codelion/data-inventor/resolve/main/datainventor/runner.py
- Command line
-
hf download hf://spaces/codelion/data-inventor/datainventor/runner.py
-
curl -L -o runner.py https://huggingface.co/spaces/codelion/data-inventor/resolve/main/datainventor/runner.py
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 | |