Spaces:
Running
Running
| """DataInventor: open-source training data creation and curation. Bring any OpenAI-compatible endpoint.""" | |
| from __future__ import annotations | |
| import base64 | |
| import csv | |
| import io | |
| import json | |
| import os | |
| import secrets | |
| import threading | |
| import urllib.parse | |
| import urllib.request | |
| import time | |
| import uuid | |
| import sys | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # the datainventor package sits next to app/ | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel | |
| from publish import LICENSES, publish as hub_publish, slug | |
| from engine import HF_BASE, HF_MODEL, Endpoint, EndpointError, Job, make_blueprint, run_job, stats | |
| ROOT = Path(__file__).resolve().parent.parent | |
| STATIC = ROOT / "static" | |
| MAX_ROWS = int(os.environ.get("MAX_ROWS", "200")) | |
| MAX_RUNNING = int(os.environ.get("MAX_RUNNING", "2")) | |
| SERVER_KEY = os.environ.get("HF_TOKEN", "") | |
| app = FastAPI(title="DataInventor") | |
| JOBS: dict[str, Job] = {} | |
| ENDPOINTS: dict[str, Endpoint] = {} # kept apart from Job so a key never lands in a response | |
| GUARD = threading.Lock() | |
| class SpecIn(BaseModel): | |
| format: str = "instruction" | |
| request: str | |
| domains: list[str] = [] | |
| subdomains: list[str] = [] | |
| languages: list[str] = ["English"] | |
| constraints: str = "" | |
| rows: int = 100 | |
| blueprint: dict | None = None | |
| def make_endpoint() -> Endpoint: | |
| """The endpoint is fixed by the host (environment variables). Visitors never supply a URL or a key.""" | |
| key = os.environ.get("DI_API_KEY") or SERVER_KEY | |
| if not key: | |
| raise HTTPException(503, "This host has no model configured. Set DI_API_KEY (or HF_TOKEN) and restart.") | |
| return Endpoint(base_url=os.environ.get("DI_BASE_URL", HF_BASE).rstrip("/"), api_key=key, model=os.environ.get("DI_MODEL", HF_MODEL), | |
| disable_thinking=os.environ.get("DI_DISABLE_THINKING", "0") != "0", workers=max(1, min(int(os.environ.get("DI_WORKERS", "4")), 8)), | |
| min_max_tokens=int(os.environ.get("DI_MIN_MAX_TOKENS", "0"))) | |
| def defaults(): | |
| return {"model": os.environ.get("DI_MODEL", HF_MODEL), "configured": bool(os.environ.get("DI_API_KEY") or SERVER_KEY), "max_rows": MAX_ROWS} | |
| def domains(): | |
| return JSONResponse(json.loads((STATIC / "domains.json").read_text())) | |
| def blueprint(s: SpecIn): | |
| ep = make_endpoint() | |
| if not s.request.strip(): | |
| raise HTTPException(400, "Describe the dataset first.") | |
| try: | |
| return make_blueprint(ep, s.request, s.domains, s.subdomains, s.languages, s.constraints, s.rows) | |
| except EndpointError as err: | |
| raise HTTPException(502, str(err)) | |
| def start(s: SpecIn): | |
| ep = make_endpoint() | |
| if not s.request.strip(): | |
| raise HTTPException(400, "Describe the dataset first.") | |
| if not 1 <= s.rows <= MAX_ROWS: | |
| raise HTTPException(400, f"Rows must be between 1 and {MAX_ROWS} on this host.") | |
| with GUARD: | |
| running = sum(1 for j in JOBS.values() if j.state in ("queued", "planning", "generating")) | |
| if running >= MAX_RUNNING: | |
| raise HTTPException(429, "Two datasets are already being written. Wait for one to finish.") | |
| jid = uuid.uuid4().hex[:10] | |
| cfg = s.model_dump(exclude={"engine"}) | |
| job = Job(id=jid, config=cfg) | |
| JOBS[jid] = job | |
| ENDPOINTS[jid] = ep | |
| for old in sorted(JOBS, key=lambda k: JOBS[k].started)[:-20]: | |
| JOBS.pop(old, None); ENDPOINTS.pop(old, None) | |
| threading.Thread(target=run_job, args=(job, ep), daemon=True).start() | |
| return {"id": jid} | |
| def get_job(jid: str) -> Job: | |
| if jid not in JOBS: | |
| raise HTTPException(404, "That dataset is no longer here. Start a new run.") | |
| return JOBS[jid] | |
| def status(jid: str, since: int = 0): | |
| j = get_job(jid) | |
| with j.lock: | |
| rows = list(j.rows) | |
| out = {"id": j.id, "state": j.state, "message": j.message, "target": j.target, "count": len(rows), "attempts": j.attempts, | |
| "rejected": dict(j.rejected), "config": {k: j.config.get(k) for k in ("format", "request", "domains", "languages", "constraints", "rows")}, "slug": slug(j.config.get("request", "")), "elapsed": round((j.finished or time.time()) - j.started, 1), "new_rows": rows[since:since + 50]} | |
| if j.state in ("done", "stopped", "failed"): | |
| out["stats"] = stats(rows) | |
| out["blueprint"] = j.blueprint | |
| return out | |
| def stop(jid: str): | |
| get_job(jid).stop.set() | |
| return {"ok": True} | |
| def download(jid: str, format: str = "csv"): | |
| j = get_job(jid) | |
| rows = list(j.rows) | |
| if not rows: | |
| raise HTTPException(404, "There are no rows to download yet.") | |
| pref = "chosen" in rows[0] | |
| cols = ["prompt", "chosen", "rejected"] if pref else ["prompt", "completion"] | |
| meta = ["area", "subtopic", "topic", "persona", "shape", "difficulty", "language"] | |
| name = f"datainventor-{jid}" | |
| if format == "jsonl": | |
| body = "\n".join(json.dumps({k: r.get(k) for k in cols + meta}, ensure_ascii=False) for r in rows) | |
| return Response(body, media_type="application/jsonl", headers={"Content-Disposition": f'attachment; filename="{name}.jsonl"'}) | |
| buf = io.StringIO() | |
| w = csv.writer(buf) | |
| w.writerow(cols + meta) | |
| for r in rows: | |
| w.writerow([r.get(k, "") for k in cols + meta]) | |
| return Response(buf.getvalue().encode("utf-8-sig"), media_type="text/csv", headers={"Content-Disposition": f'attachment; filename="{name}.csv"'}) | |
| # ---------------------------------------------------------------- sign in with Hugging Face (publishing) | |
| OAUTH_ID = os.environ.get("OAUTH_CLIENT_ID", "") | |
| OAUTH_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "") | |
| OAUTH_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co").rstrip("/") | |
| SPACE_HOST = os.environ.get("SPACE_HOST", "") | |
| DEV_TOKEN = os.environ.get("DI_DEV_USER_TOKEN", "") # local testing only | |
| SESSIONS: dict[str, dict] = {} | |
| PUBLISHES: dict[str, list] = {} | |
| SESSION_SECONDS = 2 * 3600 | |
| def current_session(request: Request) -> dict | None: | |
| sid = request.cookies.get("di_sid") | |
| s = SESSIONS.get(sid or "") | |
| if s and s["exp"] > time.time(): | |
| return s | |
| if DEV_TOKEN: | |
| from huggingface_hub import HfApi | |
| return {"user": HfApi(token=DEV_TOKEN).whoami()["name"], "token": DEV_TOKEN, "exp": time.time() + 600, "sid": "dev"} | |
| return None | |
| def redirect_uri() -> str: | |
| return f"https://{SPACE_HOST}/oauth/callback" | |
| def me(request: Request): | |
| s = current_session(request) | |
| return {"user": s["user"] if s else None, "can_sign_in": bool(OAUTH_ID and SPACE_HOST) or bool(DEV_TOKEN), "host": SPACE_HOST} | |
| def login(job: str = ""): | |
| if not (OAUTH_ID and SPACE_HOST): | |
| raise HTTPException(404, "Sign-in is not available on this host.") | |
| state = secrets.token_urlsafe(24) | |
| q = urllib.parse.urlencode({"client_id": OAUTH_ID, "redirect_uri": redirect_uri(), "response_type": "code", "scope": "openid profile contribute-repos", "state": state}) | |
| resp = RedirectResponse(f"{OAUTH_URL}/oauth/authorize?{q}") | |
| resp.set_cookie("di_state", state + "|" + (job if job.isalnum() else ""), max_age=600, httponly=True, secure=True, samesite="lax") | |
| return resp | |
| def callback(request: Request, code: str = "", state: str = "", error: str = ""): | |
| saved = request.cookies.get("di_state", "") | |
| want, _, job = saved.partition("|") | |
| if error or not code or not want or not secrets.compare_digest(want, state): | |
| return RedirectResponse("/?signin=failed") | |
| basic = base64.b64encode(f"{OAUTH_ID}:{OAUTH_SECRET}".encode()).decode() | |
| try: | |
| data = urllib.parse.urlencode({"grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri(), "client_id": OAUTH_ID}).encode() | |
| tok = json.load(urllib.request.urlopen(urllib.request.Request(f"{OAUTH_URL}/oauth/token", data, {"Authorization": "Basic " + basic, "Content-Type": "application/x-www-form-urlencoded"}), timeout=30)) | |
| info = json.load(urllib.request.urlopen(urllib.request.Request(f"{OAUTH_URL}/oauth/userinfo", headers={"Authorization": "Bearer " + tok["access_token"]}), timeout=30)) | |
| except Exception: | |
| return RedirectResponse("/?signin=failed") | |
| sid = secrets.token_urlsafe(24) | |
| SESSIONS[sid] = {"user": info.get("preferred_username") or info.get("name"), "token": tok["access_token"], "exp": time.time() + SESSION_SECONDS, "sid": sid} | |
| for k in [k for k, v in SESSIONS.items() if v["exp"] < time.time()]: | |
| SESSIONS.pop(k, None) | |
| resp = RedirectResponse("/?signin=ok" + (f"&job={job}" if job else "")) | |
| resp.set_cookie("di_sid", sid, max_age=SESSION_SECONDS, httponly=True, secure=True, samesite="lax") | |
| resp.delete_cookie("di_state") | |
| return resp | |
| def signout(request: Request): | |
| SESSIONS.pop(request.cookies.get("di_sid", ""), None) | |
| resp = JSONResponse({"ok": True}) | |
| resp.delete_cookie("di_sid") | |
| return resp | |
| class PublishIn(BaseModel): | |
| name: str | |
| private: bool = True | |
| license: str = "apache-2.0" | |
| confirm: bool = False | |
| def publish(jid: str, body: PublishIn, request: Request): | |
| s = current_session(request) | |
| if not s: | |
| raise HTTPException(401, "Sign in with Hugging Face to publish.") | |
| if not body.confirm: | |
| raise HTTPException(400, "Confirm that you have the right to publish this data.") | |
| if body.license not in LICENSES: | |
| raise HTTPException(400, "Choose one of the offered licenses.") | |
| j = get_job(jid) | |
| if j.state not in ("done", "stopped") or not j.rows: | |
| raise HTTPException(400, "Wait for the run to finish, or stop it, before publishing.") | |
| now = time.time() | |
| recent = [t for t in PUBLISHES.get(s["user"], []) if now - t < 3600] | |
| if len(recent) >= 5: | |
| raise HTTPException(429, "You have published five datasets in the last hour. Try again later.") | |
| try: | |
| url = hub_publish(s["token"], s["user"], body.name.strip(), body.private, body.license, j.config, list(j.rows), stats(list(j.rows)), ENDPOINTS[jid].model if jid in ENDPOINTS else "") | |
| except ValueError as e: | |
| raise HTTPException(400, str(e)) | |
| except Exception as e: | |
| msg = str(e) | |
| if "409" in msg or "already" in msg.lower(): | |
| raise HTTPException(409, "You already have a dataset with that name. Pick another name.") | |
| if "403" in msg or "401" in msg: | |
| raise HTTPException(403, "Hugging Face refused the request. Sign in again and allow DataInventor to create datasets.") | |
| raise HTTPException(502, "Publishing failed. Try again in a moment.") | |
| PUBLISHES[s["user"]] = recent + [now] | |
| return {"url": url} | |
| from curate_api import router as curate_router # noqa: E402 (imports make_endpoint from this module) | |
| app.include_router(curate_router) | |
| from visualize_api import router as vis_router # noqa: E402 | |
| app.include_router(vis_router) | |
| def _asset_version() -> str: | |
| import hashlib | |
| h = hashlib.sha1() | |
| for f in ("app.js", "curate.js", "visualize.js", "app.css"): | |
| h.update((STATIC / f).read_bytes()) | |
| return h.hexdigest()[:10] | |
| INDEX_HTML = (STATIC / "index.html").read_text().replace(".js\"", f".js?v={_asset_version()}\"").replace(".css\"", f".css?v={_asset_version()}\"") | |
| def index(): | |
| # never cache the page itself (embeds kept showing an old build); assets carry a content hash instead | |
| return Response(INDEX_HTML, media_type="text/html", headers={"Cache-Control": "no-cache"}) | |
| app.mount("/static", StaticFiles(directory=STATIC), name="static") | |