Spaces:
Running
Running
| """Run a recipe on a sample: every stage, small enough for a hosted app. | |
| It reads a spread sample of each source through the dataset viewer, applies clean, calibrate, create and compose, | |
| and reports what each stage did, with estimates scaled up to the full sources. The kept and created documents | |
| come back too, so a sample run is also a small curated dataset. | |
| """ | |
| from __future__ import annotations | |
| import random | |
| import threading | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dataclasses import dataclass, field | |
| from typing import Callable | |
| 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, pick_style, rephrase, styles_for | |
| from .llm import Endpoint, EndpointError | |
| from .recipe import held_out_of, sources_of, targets_of | |
| from .sources import sample | |
| from .text import est_tokens | |
| _BENCH_CACHE: dict[str, list[str]] = {} | |
| _BENCH_LOCK = threading.Lock() | |
| def bench_texts(key: str, token: str = "") -> list[str]: | |
| with _BENCH_LOCK: | |
| if key not in _BENCH_CACHE: | |
| _BENCH_CACHE[key] = benchmark_texts(key, token) | |
| return _BENCH_CACHE[key] | |
| class Preview: | |
| recipe: dict | |
| state: str = "queued" # queued | running | done | failed | stopped | |
| message: str = "" | |
| stage: str = "" | |
| report: dict = field(default_factory=dict) | |
| docs: list = field(default_factory=list) | |
| started: float = field(default_factory=time.time) | |
| finished: float = 0.0 | |
| stop: threading.Event = field(default_factory=threading.Event) | |
| def say(self, stage: str, msg: str): | |
| self.stage, self.message = stage, msg | |
| def _pmap(fn: Callable, items: list, workers: int, stop: threading.Event, tick: Callable[[int], None] | None = None) -> list: | |
| done = [0] | |
| lock = threading.Lock() | |
| def guarded(x): | |
| if stop.is_set(): | |
| return None | |
| out = fn(x) | |
| if tick: | |
| with lock: | |
| done[0] += 1 | |
| tick(done[0]) | |
| return out | |
| with ThreadPoolExecutor(max_workers=max(1, workers)) as pool: | |
| return list(pool.map(guarded, items)) | |
| def run_preview(p: Preview, ep: Endpoint, token: str = "", per_source: int = 200, grade_n: int = 32, | |
| rephrase_n: int = 8, invent_n: int = 4, max_docs: int = 1000, seed: int = 7): | |
| r = p.recipe | |
| rng = random.Random(seed) | |
| try: | |
| p.state = "running" | |
| rep: dict = {"sources": [], "stages": {}} | |
| # ---- targets | |
| targets_text: list[str] = [] | |
| for t in targets_of(r): | |
| p.say("setup", f"Reading target {t.name}") | |
| targets_text += [d["text"] for d in sample(t, min(100, per_source), token)] | |
| calib = cal.Calibrator(float(r["calibrate"]["focus"]), r["targets"]["description"], targets_text) | |
| # ---- clean | |
| cleaner = Cleaner(r["clean"]["rules"]) | |
| dec = Decontaminator(int(r["clean"].get("ngram", 13))) | |
| for b in r["clean"]["decontaminate"]: | |
| p.say("clean", f"Loading the {b} test set for decontamination") | |
| try: | |
| dec.add(b, bench_texts(b, token)) | |
| except Exception as e: | |
| rep.setdefault("warnings", []).append(f"Could not load {b} for decontamination: {e}") | |
| for h in held_out_of(r): | |
| p.say("clean", f"Loading held-out set {h.name}") | |
| try: | |
| dec.add(h.name, [d["text"] for d in sample(h, 2000, token)]) | |
| except Exception as e: | |
| rep.setdefault("warnings", []).append(f"Could not load held-out set {h.name}: {e}") | |
| dedup = cal.Deduper(near=r["calibrate"]["dedup"]["near"], threshold=float(r["calibrate"]["dedup"]["threshold"])) | |
| survivors: list[dict] = [] | |
| for s in sources_of(r): | |
| if p.stop.is_set(): | |
| break | |
| p.say("clean", f"Sampling {s.name}") | |
| raw = sample(s, per_source, token) | |
| f = {"name": s.name, "dataset": s.dataset or s.files, "kind": s.kind, "num_rows": s.num_rows, "read": len(raw), | |
| "dropped": {}, "decontaminated": {}, "kept": 0} | |
| f["repeated_passages"] = cal.repeated_passages(raw) | |
| f["raw_tokens"] = sum(est_tokens(d["text"]) for d in raw) | |
| for d in raw: | |
| d2, why = cleaner.apply(d) | |
| if not d2: | |
| f["dropped"][why] = f["dropped"].get(why, 0) + 1 | |
| continue | |
| b = dec.hit(d2["text"]) | |
| if b: | |
| f["decontaminated"][b] = f["decontaminated"].get(b, 0) + 1 | |
| continue | |
| why = dedup.check(d2) | |
| if why: | |
| f["dropped"][why] = f["dropped"].get(why, 0) + 1 | |
| continue | |
| survivors.append(d2) | |
| f["kept"] += 1 | |
| rep["sources"].append(f) | |
| rep["stages"]["clean"] = {"benchmarks": r["clean"]["decontaminate"], "test_items": len(dec)} | |
| if not survivors: | |
| raise ValueError("Nothing survived cleaning. Loosen the rules, or check the text column of each source.") | |
| # ---- calibrate: grade, learn, score | |
| by_src: dict[str, list[dict]] = {} | |
| for d in survivors: | |
| by_src.setdefault(d["source"], []).append(d) | |
| to_grade = [] | |
| per = max(2, grade_n // max(1, len(by_src))) | |
| for docs in by_src.values(): | |
| to_grade += rng.sample(docs, min(per, len(docs))) | |
| p.say("calibrate", f"Grading {len(to_grade)} documents with the model") | |
| desc = r["targets"]["description"] | |
| grades = _pmap(lambda d: cal.grade(ep, d["text"], desc), to_grade, ep.workers, p.stop, | |
| lambda n: p.say("calibrate", f"Graded {n} of {len(to_grade)} documents with the model")) | |
| graded = [(d, g) for d, g in zip(to_grade, grades) if g] | |
| if len(graded) < 4: | |
| raise EndpointError("The model returned too few usable grades. Check the endpoint, or try again.") | |
| calib.fit([(d["text"], g) for d, g in graded], [d["text"] for d in survivors]) | |
| gmap = {d["id"]: g for d, g in graded} | |
| for d in survivors: | |
| g = gmap.get(d["id"]) | |
| if g: | |
| d["grade"] = g | |
| d.update(calib.score(d["text"], g)) | |
| d["tokens"] = est_tokens(d["text"]) | |
| keep_top = float(r["calibrate"]["keep_top"]) | |
| for f in rep["sources"]: | |
| docs = sorted(by_src.get(f["name"], []), key=lambda d: -d["score"]) | |
| k = max(1, round(len(docs) * keep_top)) if docs else 0 | |
| hi = max(1, round(len(docs) * 0.2)) if docs else 0 | |
| for i, d in enumerate(docs): | |
| d["kept"] = i < k | |
| d["tier"] = "high" if i < hi else "medium" if i < k else "low" | |
| kept = [d for d in docs if d["kept"]] | |
| f["selected"] = len(kept) | |
| f["dropped"]["below score cut"] = len(docs) - len(kept) | |
| rate = len(kept) / max(1, f["read"]) | |
| mean_tok = sum(d["tokens"] for d in kept) / max(1, len(kept)) | |
| f["mean_tokens"] = round(mean_tok) | |
| f["keep_rate"] = round(rate, 3) | |
| if f["num_rows"]: | |
| f["available_tokens"] = int(f["num_rows"] * rate * mean_tok) | |
| f["high_tokens"] = int(f["num_rows"] * (hi / max(1, f["read"])) * mean_tok) | |
| f["score_hist"] = _hist([d["score"] for d in docs]) | |
| f["examples"] = {"top": [_ex(d) for d in docs[:2]], "bottom": [_ex(d) for d in docs[-2:]]} | |
| rep["stages"]["calibrate"] = {"graded": len(graded), "agreement": calib.agreement, "focus": calib.focus, "keep_top": keep_top, | |
| "has_targets": calib.has_targets, "grades": [{"id": d["id"], **g} for d, g in graded][:12]} | |
| # ---- create | |
| created: list[dict] = [] | |
| rcfg, icfg = r["create"]["rephrase"], r["create"]["invent"] | |
| if rcfg["enabled"] and rephrase_n and not p.stop.is_set(): | |
| queues = [sorted([d for d in docs if d.get("tier") == "high"], key=lambda d: -d["score"]) for docs in by_src.values()] | |
| picks = [] | |
| while len(picks) < rephrase_n and any(queues): # round-robin so every source is represented | |
| for q in queues: | |
| if q and len(picks) < rephrase_n: | |
| picks.append(q.pop(0)) | |
| jobs = [({**d, "text": chunks(d["text"])[0]}, pick_style(d["kind"], rng, rcfg["styles"] or styles_for(d["kind"]))) for d in picks] | |
| p.say("create", f"Rephrasing {len(jobs)} of the best documents") | |
| outs = _pmap(lambda j: rephrase(ep, j[0]["text"], j[1], rcfg.get("language", ""), float(rcfg["max_new_share"])), jobs, ep.workers, p.stop, | |
| lambda n: p.say("create", f"Rephrased {n} of {len(jobs)} of the best documents")) | |
| for (d, style), o in zip(jobs, outs): | |
| if o and o["text"]: | |
| created.append({"id": f"{d['id']}:{style}", "text": o["text"], "source": d["source"], "kind": "synthetic", | |
| "origin": d["id"], "origin_text": d["text"], "style": style, "fidelity": o.get("fidelity"), "kept": o["ok"], | |
| "why": o.get("why", ""), "tokens": est_tokens(o["text"])}) | |
| if icfg["enabled"] and invent_n and icfg["request"].strip() and not p.stop.is_set(): | |
| p.say("create", "Writing new documents") | |
| topics = plan_topics(ep, icfg["request"], invent_n) | |
| jobs = [(t, rng.choice(icfg["doc_types"]), rng.choice(icfg["levels"])) for t in topics[:invent_n]] | |
| outs = _pmap(lambda j: invent(ep, j[0], j[1], j[2], icfg["request"], icfg.get("language", "English")), jobs, ep.workers, p.stop) | |
| for (t, dt, lv), o in zip(jobs, outs): | |
| if o and o["text"]: | |
| created.append({"id": f"invent:{len(created)}", "text": o["text"], "source": "invented", "kind": "synthetic", | |
| "topic": t, "doc_type": dt, "level": lv, "kept": True, "tokens": est_tokens(o["text"])}) | |
| rep["stages"]["create"] = {"made": len(created), "kept": sum(1 for c in created if c["kept"]), | |
| "flagged": [{"id": c["id"], "why": c["why"], "fidelity": c.get("fidelity")} for c in created if not c["kept"]]} | |
| # ---- compose | |
| src_rows = [{"name": f["name"], "kind": f["kind"], "available_tokens": f.get("available_tokens", 0), | |
| "high_tokens": f.get("high_tokens", 0)} for f in rep["sources"]] | |
| defaults = default_weights(src_rows) | |
| plans = [] | |
| for m in r["compose"]["mixes"]: | |
| w = {**defaults, **{k: float(v) for k, v in (m.get("weights") or {}).items()}} | |
| plan = plan_mix(int(m["budget_tokens"]), src_rows, w, float(r["compose"]["max_epochs"]), float(r["compose"]["max_synthetic_share"])) | |
| plan["name"] = m["name"] | |
| plan["budget"] = fmt(m["budget_tokens"]) | |
| plans.append(plan) | |
| if any(not f["num_rows"] for f in rep["sources"]): | |
| rep.setdefault("warnings", []).append("Some sources report no row count, so their available tokens are unknown and plans treat them as empty.") | |
| rep["stages"]["compose"] = {"default_weights": defaults, "plans": plans} | |
| # ---- output sample | |
| kept_docs = [d for d in survivors if d.get("kept")] | |
| kept_docs.sort(key=lambda d: -d["score"]) | |
| out = [c for c in created if c["kept"]] + kept_docs | |
| p.docs = [_out(d) for d in out[:max_docs]] | |
| n_created = min(sum(1 for c in created if c["kept"]), max_docs) | |
| rep["output"] = {"docs": len(p.docs), "curated": len(p.docs) - n_created, "created": n_created, | |
| "tokens": sum(d["tokens"] for d in out[:max_docs])} | |
| rep["created_examples"] = [{k: c.get(k) for k in ("id", "style", "doc_type", "topic", "kept", "why", "fidelity")} | {"text": c["text"][:1600], "origin_text": c.get("origin_text", "")[:1600]} for c in created[:6]] | |
| p.report = rep | |
| p.state = "stopped" if p.stop.is_set() else "done" | |
| p.say("done", f"{len(p.docs)} documents ready" if p.state == "done" else "Stopped") | |
| except EndpointError as e: | |
| p.state, p.message = "failed", str(e) | |
| except Exception as e: | |
| p.state, p.message = "failed", str(e) or e.__class__.__name__ | |
| finally: | |
| p.finished = time.time() | |
| def plan_topics(ep: Endpoint, request: str, n: int) -> list[str]: | |
| import json | |
| import re | |
| s = ep.chat(f"""List {n} specific concepts to write teaching documents about, for this collection: "{request}". | |
| Spread them across the whole subject and across difficulty. Return ONLY a JSON list of strings, each 2 to 6 words.""", | |
| max_tokens=600, temperature=0.8, validate=lambda t: bool(re.search(r"\[.*\]", t, re.S))) | |
| try: | |
| return [str(x) for x in json.loads(re.search(r"\[.*\]", s, re.S).group(0))][:n] | |
| except Exception: | |
| return [request] * n | |
| def _hist(vals: list[float], bins: int = 10) -> list[int]: | |
| h = [0] * bins | |
| for v in vals: | |
| h[min(bins - 1, int(v * bins))] += 1 | |
| return h | |
| def _ex(d: dict) -> dict: | |
| return {"id": d["id"], "score": d["score"], "quality": d["quality"], "relevance": d["relevance"], | |
| "reason": (d.get("grade") or {}).get("reason", ""), "text": d["text"][:700]} | |
| def _out(d: dict) -> dict: | |
| keep = ("id", "text", "source", "kind", "quality", "relevance", "score", "tier", "style", "origin", "doc_type", "topic", "level", "tokens") | |
| return {k: d[k] for k in keep if k in d} | |