Spaces:
Running
Running
| """Curate a corpus: sample previews of the clean, calibrate, create, compose workflow, plus the recipe to run it in full.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import threading | |
| import time | |
| import uuid | |
| from fastapi import APIRouter, HTTPException | |
| from fastapi.responses import Response | |
| from pydantic import BaseModel | |
| from datainventor.clean import BENCHMARKS, DEFAULTS as CLEAN_RULES | |
| from datainventor.create import DOC_TYPES, STYLES, chunks, rephrase, styles_for | |
| from datainventor.llm import EndpointError | |
| from datainventor.preview import Preview, run_preview | |
| from datainventor.recipe import DEFAULT, dump, validate | |
| from datainventor.sources import KINDS, inspect | |
| router = APIRouter(prefix="/api/curate") | |
| PER_SOURCE = int(os.environ.get("CURATE_PER_SOURCE", "200")) | |
| GRADE_N = int(os.environ.get("CURATE_GRADE", "48")) | |
| REPHRASE_N = int(os.environ.get("CURATE_REPHRASE", "8")) | |
| MAX_SOURCES, MAX_TARGETS = 4, 2 | |
| INSTALL = "pip install 'datainventor[full] @ https://huggingface.co/spaces/codelion/data-inventor/resolve/main/dist/datainventor-0.2.2-py3-none-any.whl'" | |
| PREVIEWS: dict[str, Preview] = {} | |
| TRIES: dict[str, list] = {} | |
| GUARD = threading.Lock() | |
| def _ep(): | |
| from main import make_endpoint # the host's endpoint; visitors never supply one | |
| return make_endpoint() | |
| def _token() -> str: | |
| return os.environ.get("HF_TOKEN", "") | |
| def defaults(): | |
| from main import MAX_ROWS | |
| return {"recipe": DEFAULT, "benchmarks": list(BENCHMARKS), "kinds": list(KINDS), "clean_rules": CLEAN_RULES, | |
| "styles": {k: v[0] for k, v in STYLES.items()}, "styles_by_kind": {k: styles_for(k) for k in KINDS}, | |
| "doc_types": list(DOC_TYPES), "per_source": PER_SOURCE, "max_sources": MAX_SOURCES, "max_targets": MAX_TARGETS, | |
| "max_docs": MAX_ROWS, "install": INSTALL} | |
| def do_inspect(dataset: str): | |
| ds = dataset.strip().removeprefix("https://huggingface.co/datasets/").strip("/") | |
| if not ds or ds.count("/") != 1 or len(ds) > 120: | |
| raise HTTPException(400, "Use a Hugging Face dataset id like HuggingFaceFW/fineweb-edu.") | |
| try: | |
| return inspect(ds, _token()) | |
| except ValueError as e: | |
| raise HTTPException(404, str(e)) | |
| except Exception: | |
| raise HTTPException(404, "The dataset viewer could not open that dataset. It must be public and have the viewer enabled.") | |
| class RecipeIn(BaseModel): | |
| recipe: dict | |
| def _checked(rec: dict) -> dict: | |
| for s in rec.get("sources", []) + rec.get("targets", {}).get("datasets", []): | |
| if s.get("files"): | |
| raise HTTPException(400, "Local files only work in the runner on your own machines. Use Hugging Face datasets here.") | |
| if len(rec.get("sources", [])) > MAX_SOURCES or len(rec.get("targets", {}).get("datasets", [])) > MAX_TARGETS: | |
| raise HTTPException(400, f"The hosted preview takes up to {MAX_SOURCES} sources and {MAX_TARGETS} target datasets.") | |
| try: | |
| return validate(rec) | |
| except ValueError as e: | |
| raise HTTPException(400, str(e)) | |
| def recipe_file(body: RecipeIn): | |
| """The recipe as YAML, for the runner.""" | |
| r = _checked(body.recipe) | |
| return Response(dump(r), media_type="application/x-yaml", headers={"Content-Disposition": 'attachment; filename="recipe.yaml"'}) | |
| def start(body: RecipeIn): | |
| r = _checked(body.recipe) | |
| ep = _ep() | |
| with GUARD: | |
| running = sum(1 for p in PREVIEWS.values() if p.state in ("queued", "running")) | |
| if running >= 2: | |
| raise HTTPException(429, "Two previews are already running. Wait for one to finish.") | |
| pid = uuid.uuid4().hex[:10] | |
| p = Preview(r) | |
| PREVIEWS[pid] = p | |
| for old in sorted(PREVIEWS, key=lambda k: PREVIEWS[k].started)[:-20]: | |
| PREVIEWS.pop(old, None) | |
| from main import MAX_ROWS | |
| threading.Thread(target=run_preview, args=(p, ep, _token()), | |
| kwargs={"per_source": PER_SOURCE, "grade_n": GRADE_N, "rephrase_n": REPHRASE_N, "invent_n": 4, "max_docs": MAX_ROWS}, | |
| daemon=True).start() | |
| return {"id": pid} | |
| def _get(pid: str) -> Preview: | |
| if pid not in PREVIEWS: | |
| raise HTTPException(404, "That preview is no longer here. Run it again.") | |
| return PREVIEWS[pid] | |
| def status(pid: str): | |
| p = _get(pid) | |
| out = {"id": pid, "state": p.state, "stage": p.stage, "message": p.message, "elapsed": round((p.finished or time.time()) - p.started, 1)} | |
| if p.state in ("done", "stopped"): | |
| out["report"] = p.report | |
| return out | |
| def stop(pid: str): | |
| _get(pid).stop.set() | |
| return {"ok": True} | |
| def download(pid: str): | |
| p = _get(pid) | |
| if not p.docs: | |
| raise HTTPException(404, "There are no documents to download yet.") | |
| body = "\n".join(json.dumps(d, ensure_ascii=False) for d in p.docs) + "\n" | |
| return Response(body, media_type="application/jsonl", headers={"Content-Disposition": f'attachment; filename="curated-sample-{pid}.jsonl"'}) | |
| class TryIn(BaseModel): | |
| text: str | |
| kind: str = "web" | |
| styles: list[str] = [] | |
| def try_styles(body: TryIn): | |
| """Rewrite one document in up to four styles, side by side.""" | |
| text = body.text.strip() | |
| if len(text.split()) < 60: | |
| raise HTTPException(400, "Paste or pick a document of at least 60 words.") | |
| styles = [s for s in (body.styles or styles_for(body.kind)) if s in STYLES][:4] | |
| if not styles: | |
| raise HTTPException(400, "Choose at least one style.") | |
| now = time.time() | |
| with GUARD: | |
| recent = [t for t in TRIES.get("all", []) if now - t < 60] | |
| if len(recent) >= 6: | |
| raise HTTPException(429, "Too many rewrites in the last minute. Try again shortly.") | |
| TRIES["all"] = recent + [now] | |
| ep = _ep() | |
| piece = chunks(text)[0] | |
| from concurrent.futures import ThreadPoolExecutor | |
| try: | |
| with ThreadPoolExecutor(len(styles)) as ex: | |
| outs = list(ex.map(lambda s: rephrase(ep, piece, s), styles)) | |
| except EndpointError as e: | |
| raise HTTPException(502, str(e)) | |
| return {"source": piece, "results": [{"style": o["style"], "name": STYLES[o["style"]][0], "text": o["text"], "ok": o["ok"], | |
| "why": o.get("why", ""), "fidelity": o.get("fidelity")} for o in outs]} | |