Spaces:
Running
Running
| """Visualize: load a dataset (a run from this page, a sample, a Hugging Face dataset or a file), ask questions, get charts.""" | |
| from __future__ import annotations | |
| import csv | |
| import io | |
| import json | |
| import math | |
| import os | |
| import re | |
| import threading | |
| import time | |
| import uuid | |
| from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| from fastapi import APIRouter, HTTPException | |
| from fastapi.responses import Response | |
| from pydantic import BaseModel | |
| from datainventor import visualize as V | |
| from datainventor.llm import EndpointError | |
| from datainventor.sources import _get, inspect | |
| router = APIRouter(prefix="/api/vis") | |
| MAX_ROWS = int(os.environ.get("VIS_MAX_ROWS", "200")) | |
| MAX_UPLOAD = 6_000_000 | |
| MAX_RUNNING = int(os.environ.get("VIS_MAX_RUNNING", "4")) | |
| WRITE_SECONDS = float(os.environ.get("VIS_WRITE_SECONDS", "20")) | |
| CALL_SECONDS = int(os.environ.get("VIS_CALL_SECONDS", "30")) | |
| WRITERS = ThreadPoolExecutor(max_workers=4) | |
| SAMPLES = [("support-general", "Customer support replies"), ("hotel-reviews", "Hotel reviews"), ("medical-qa", "Medical questions"), | |
| ("legal-qa", "Legal questions"), ("support-banking", "Banking support"), ("african-qa", "Questions about Africa"), | |
| ("hindi-news", "Hindi news requests"), ("serbian-car-ads", "Serbian car ads")] | |
| JOB_ORDER = ["prompt", "completion", "chosen", "rejected", "area", "subtopic", "topic", "persona", "shape", "difficulty", "language"] | |
| class VisSet: | |
| id: str | |
| name: str | |
| source: str | |
| rows: list | |
| cols: dict | |
| created: float = field(default_factory=time.time) | |
| lock: threading.Lock = field(default_factory=threading.Lock) | |
| busy: bool = False | |
| presets: dict = field(default_factory=dict) | |
| class Ask: | |
| id: str | |
| set_id: str | |
| question: str | |
| scope: list | |
| scope_label: str | |
| stage: str = "plan" # plan | label | chart | write | done | failed | |
| done: int = 0 | |
| total: int = 0 | |
| result: Optional[dict] = None | |
| error: str = "" | |
| started: float = field(default_factory=time.time) | |
| SETS: dict[str, VisSet] = {} | |
| ASKS: dict[str, Ask] = {} | |
| 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", "") | |
| # ------------------------------------------------------------------ loading | |
| def _hub_rows(dataset: str, n: int) -> list[dict]: | |
| """A spread of up to n full rows through the dataset viewer: no download, every column kept.""" | |
| if not re.fullmatch(r"[\w.-]+/[\w.-]+", dataset): | |
| raise HTTPException(400, "Give a dataset id like owner/name.") | |
| try: | |
| info = inspect(dataset, _token()) | |
| except Exception as e: | |
| raise HTTPException(400, f"Could not open {dataset}: {str(e)[:160]}") | |
| total = info.get("num_rows") or 0 | |
| chunks = 1 if not total or total <= n else min(4, max(3, math.ceil(n / 25))) # sorted datasets need several slices | |
| per = math.ceil(n / chunks) | |
| starts = [0] if chunks == 1 else [max(0, int(total * (i + 0.5) / chunks) - per // 2) for i in range(chunks)] | |
| out, seen = [], set() | |
| for s in starts: | |
| off = s | |
| while off < s + per and len(out) < n: | |
| try: | |
| got = _get("rows", {"dataset": dataset, "config": info["config"], "split": info["split"], "offset": off, "length": min(100, s + per - off)}, _token()) | |
| except Exception as e: | |
| if not out: | |
| raise HTTPException(502, f"The dataset viewer did not return rows for {dataset}: {str(e)[:160]}") | |
| break | |
| rows = got.get("rows") or [] | |
| for r in rows: | |
| idx = r.get("row_idx", off) | |
| if idx not in seen: | |
| seen.add(idx) | |
| out.append(r.get("row") or {}) | |
| if len(rows) < min(100, s + per - off): | |
| break | |
| off += len(rows) | |
| if not out: | |
| raise HTTPException(400, f"{dataset} has no rows the viewer can show.") | |
| return out[:n] | |
| def _parse_upload(text: str) -> list[dict]: | |
| t = text.lstrip("") | |
| s = t.lstrip() | |
| try: | |
| if s.startswith("["): | |
| data = json.loads(s) | |
| elif s.startswith("{"): | |
| try: | |
| one = json.loads(s) | |
| data = next((v for v in one.values() if isinstance(v, list)), [one]) if isinstance(one, dict) else [] | |
| except ValueError: | |
| data = [json.loads(line) for line in s.splitlines() if line.strip()] | |
| else: | |
| csv.field_size_limit(10_000_000) | |
| data = list(csv.DictReader(io.StringIO(t))) | |
| except (ValueError, csv.Error) as e: | |
| raise HTTPException(400, f"Could not read that file as CSV, JSON or JSONL: {str(e)[:120]}") | |
| rows = [r for r in data if isinstance(r, dict)] | |
| if not rows: | |
| raise HTTPException(400, "The file has no rows. Use CSV with a header row, a JSON list of objects, or JSONL.") | |
| return rows | |
| def _job_rows(jid: str) -> tuple[list[dict], str]: | |
| from main import JOBS | |
| j = JOBS.get(jid) | |
| if not j: | |
| raise HTTPException(404, "That run is no longer on this host. Download it and upload the file instead.") | |
| with j.lock: | |
| rows = list(j.rows) | |
| if not rows: | |
| raise HTTPException(400, "That run has no rows yet.") | |
| keys = [k for k in JOB_ORDER if k in rows[0]] | |
| req = (j.config.get("request") or "Your run").strip() | |
| return [{k: r.get(k) for k in keys} for r in rows], req if len(req) <= 80 else req[:78].rsplit(" ", 1)[0] + "…" | |
| class LoadIn(BaseModel): | |
| source: str # job | sample | hub | upload | |
| job: str = "" | |
| key: str = "" | |
| dataset: str = "" | |
| text: str = "" | |
| filename: str = "" | |
| n: int = 50 | |
| def samples(): | |
| return {"samples": [{"key": k, "title": t, "repo": f"LocalLLaMA/datainventor-{k}"} for k, t in SAMPLES], "max_rows": MAX_ROWS} | |
| def load(body: LoadIn): | |
| n = max(5, min(body.n or 50, MAX_ROWS)) | |
| note = "" | |
| if body.source == "job": | |
| raw, name = _job_rows(body.job) | |
| src = "This page's run" | |
| elif body.source == "sample": | |
| if body.key not in dict(SAMPLES): | |
| raise HTTPException(400, "Unknown sample.") | |
| raw, name, src = _hub_rows(f"LocalLLaMA/datainventor-{body.key}", n), dict(SAMPLES)[body.key], f"LocalLLaMA/datainventor-{body.key}" | |
| elif body.source == "hub": | |
| ds = body.dataset.strip().removeprefix("https://huggingface.co/datasets/").strip("/") | |
| raw, name, src = _hub_rows(ds, n), ds.split("/")[-1], ds | |
| elif body.source == "upload": | |
| if len(body.text) > MAX_UPLOAD: | |
| raise HTTPException(413, "That file is over 6 MB. Upload a smaller sample.") | |
| raw, name, src = _parse_upload(body.text), (body.filename or "Your file")[:80], "Uploaded file" | |
| else: | |
| raise HTTPException(400, "Unknown source.") | |
| if len(raw) > MAX_ROWS: | |
| note = f"Showing the first {MAX_ROWS} of {len(raw)} rows." | |
| raw = raw[:MAX_ROWS] | |
| rows, cols = V.prepare(raw) | |
| if not cols: | |
| raise HTTPException(400, "None of the columns hold text, categories or numbers this can chart.") | |
| sid = uuid.uuid4().hex[:10] | |
| vs = VisSet(id=sid, name=name, source=src, rows=rows, cols=cols, presets=V.presets(rows, cols)) | |
| with GUARD: | |
| SETS[sid] = vs | |
| for old in sorted(SETS, key=lambda k: SETS[k].created)[:-30]: | |
| SETS.pop(old, None) | |
| return {"id": sid, "name": name, "source": src, "n": len(rows), "note": note, "columns": V.describe(rows, cols), | |
| "glance": V.glance(rows, cols), "suggestions": list(vs.presets), "rows": rows} | |
| # ------------------------------------------------------------------ asking | |
| class AskIn(BaseModel): | |
| question: str | |
| scope: list[int] = [] | |
| scope_label: str = "" | |
| def _column(d: dict, vals: dict, filled: set, prev: Optional[dict] = None) -> dict: | |
| c = {"type": "num" if d["kind"] in ("score", "count") else "cat", "derived": d["kind"], "from": d["from"], | |
| "guide": d.get("guide", ""), "filled": set(filled)} | |
| if d.get("labels"): | |
| c["labels"] = d["labels"] | |
| if d.get("pattern"): | |
| c["pattern"] = d["pattern"] | |
| if prev: | |
| c["filled"] |= prev.get("filled", set()) | |
| return c | |
| def run_ask(a: Ask, vs: VisSet, ep): | |
| # short per-call timeouts: a stuck provider is retried instead of holding the chart for minutes | |
| chat = lambda prompt, **kw: ep.chat(prompt, retries=3, timeout=CALL_SECONDS, **kw) | |
| try: | |
| ids = set(a.scope) | |
| rows = [r for r in vs.rows if r["_id"] in ids] if ids else list(vs.rows) | |
| scope = f" (this question is about a selection of {len(rows)} of {len(vs.rows)} rows: {a.scope_label or 'chosen on a chart'})" if ids else "" | |
| plan, err = None, "" | |
| pre = vs.presets.get(a.question) | |
| if pre: # a starter question: its plan is already written, and labels found before are reused | |
| plan = {**pre, "derive": [], "chart": dict(pre["chart"])} | |
| for d in pre["derive"]: | |
| c = vs.cols.get(d["name"]) | |
| if not (c and c.get("derived") == d["kind"] and c.get("from") == d["from"]): | |
| plan["derive"].append(dict(d)) | |
| def plan_once(temp): | |
| text = chat(V.plan_prompt(a.question, vs.name, rows, vs.cols, scope, err), max_tokens=700, temperature=temp) | |
| try: | |
| return V.parse_plan(text, vs.cols) | |
| except V.PlanError as e: | |
| errs.append(str(e)) | |
| print(f"[visualize] plan rejected: {e} | reply: {text[:600]!r}", flush=True) | |
| return None | |
| errs: list = [] | |
| for _ in range(0 if plan else 2): | |
| plan = V.hedged(plan_once, tries=1) | |
| if plan: | |
| break | |
| err = errs[-1] if errs else "" | |
| if not plan: | |
| plan = V.fallback_plan(a.question, rows, vs.cols) | |
| # new columns, plus earlier model-made columns this chart needs that do not cover these rows yet | |
| todo = [(d, rows, False) for d in plan["derive"]] | |
| for k in ("x", "y", "color"): | |
| c = vs.cols.get(plan["chart"].get(k) or "") | |
| if c and c.get("derived") in V.KINDS: | |
| missing = [r for r in rows if r["_id"] not in c.get("filled", set())] | |
| if missing: | |
| todo.append(({"name": plan["chart"][k], "kind": c["derived"], "from": c["from"], "guide": c.get("guide", ""), | |
| "labels": c.get("labels", []), "pattern": c.get("pattern", "")}, missing, True)) | |
| touched = {} | |
| if todo: | |
| a.stage, a.total, a.done = "label", sum(len(rs) for _, rs, _ in todo), 0 | |
| for d, rs, refill in todo: | |
| base = a.done | |
| vals = V.fill(d, rs, chat, a.question, progress=lambda k, base=base: setattr(a, "done", base + k)) | |
| with vs.lock: | |
| for r in vs.rows: | |
| if r["_id"] in vals: | |
| r[d["name"]] = vals[r["_id"]] | |
| elif not refill or d["name"] not in r: # a new definition: clear values left by an old one | |
| r[d["name"]] = None if d["kind"] in ("score", "count") else "" | |
| vs.cols[d["name"]] = _column(d, vals, set(vals), vs.cols.get(d["name"]) if refill else None) | |
| touched[d["name"]] = {str(i): v for i, v in vals.items()} | |
| a.stage = "chart" | |
| with vs.lock: | |
| c = V.chart(rows, vs.cols, V.check_spec(plan["chart"], vs.cols)) | |
| columns = V.describe(vs.rows, vs.cols) | |
| made = [{"name": d["name"], "from": d["from"]} for d, _, _ in todo] | |
| made = list({m["name"]: m for m in made}.values()) | |
| a.result = {"title": plan["title"], "chart": c, "summary": "", "notes": V.notes(c, made, plan.get("fallback", False)), | |
| "facts": V.facts(c), "follow_ups": plan["follow_ups"], "derived": touched, "columns": columns, | |
| "scope": len(rows) if ids else 0} | |
| a.stage = "write" | |
| try: | |
| # the chart is already on screen; a slow write-up gives way to the exact numbers after a while | |
| fut = WRITERS.submit(chat, V.summary_prompt(a.question, plan["title"], c, made), max_tokens=320, temperature=0.3) | |
| text = fut.result(timeout=WRITE_SECONDS) | |
| text = re.sub(r"[*#_`]+", "", re.sub(r"^\s*[-•]\s+", "", text, flags=re.M)).strip() | |
| text = re.sub(r"\s*\n\s*", " ", text) | |
| a.result["summary"] = " ".join(re.split(r"(?<=[.!?])\s+(?=[A-Z0-9])", text)[:4])[:700] | |
| except (EndpointError, FutureTimeout): | |
| pass | |
| a.stage = "done" | |
| except (EndpointError, V.PlanError) as e: | |
| a.stage, a.error = "failed", str(e)[:300] | |
| except Exception as e: # never leave the page polling a dead ask | |
| a.stage, a.error = "failed", f"Something went wrong drawing that: {str(e)[:200]}" | |
| finally: | |
| vs.busy = False | |
| def ask(sid: str, body: AskIn): | |
| vs = SETS.get(sid) | |
| if not vs: | |
| raise HTTPException(404, "That dataset is no longer loaded here. Load it again.") | |
| q = re.sub(r"\s+", " ", body.question).strip() | |
| if not q: | |
| raise HTTPException(400, "Type a question first.") | |
| ep = _ep() | |
| with GUARD: | |
| if vs.busy: | |
| raise HTTPException(409, "This dataset is still answering the last question.") | |
| if sum(1 for a in ASKS.values() if a.stage not in ("done", "failed")) >= MAX_RUNNING: | |
| raise HTTPException(429, "Several charts are being drawn right now. Try again in a moment.") | |
| aid = uuid.uuid4().hex[:10] | |
| a = Ask(id=aid, set_id=sid, question=q[:400], scope=[i for i in body.scope if isinstance(i, int)][:MAX_ROWS], scope_label=body.scope_label[:120]) | |
| ASKS[aid] = a | |
| vs.busy = True | |
| for old in sorted(ASKS, key=lambda k: ASKS[k].started)[:-200]: | |
| ASKS.pop(old, None) | |
| threading.Thread(target=run_ask, args=(a, vs, ep), daemon=True).start() | |
| return {"ask": aid} | |
| def ask_status(aid: str): | |
| a = ASKS.get(aid) | |
| if not a: | |
| raise HTTPException(404, "That question is no longer here. Ask it again.") | |
| out = {"stage": a.stage, "done": a.done, "total": a.total, "error": a.error, "elapsed": round(time.time() - a.started, 1)} | |
| if a.result and a.stage in ("write", "done"): | |
| out["result"] = a.result | |
| return out | |
| def download(sid: str, ids: str = "", format: str = "csv"): | |
| vs = SETS.get(sid) | |
| if not vs: | |
| raise HTTPException(404, "That dataset is no longer loaded here.") | |
| want = {int(x) for x in ids.split(",") if x.strip().isdigit()} | |
| rows = [r for r in vs.rows if not want or r["_id"] in want] | |
| keys = ["_id"] + [k for k, c in vs.cols.items() if c.get("derived") not in ("words", "turns")] | |
| name = re.sub(r"[^\w-]+", "-", vs.name.lower()).strip("-")[:40] or "rows" | |
| if format == "jsonl": | |
| body = "\n".join(json.dumps({("row" if k == "_id" else k): r.get(k) for k in keys}, ensure_ascii=False) for r in rows) | |
| return Response(body, media_type="application/jsonl", headers={"Content-Disposition": f'attachment; filename="{name}-{len(rows)}.jsonl"'}) | |
| buf = io.StringIO() | |
| w = csv.writer(buf) | |
| w.writerow(["row" if k == "_id" else k for k in keys]) | |
| for r in rows: | |
| w.writerow(["" if r.get(k) is None else r.get(k) for k in keys]) | |
| return Response(buf.getvalue().encode("utf-8-sig"), media_type="text/csv", headers={"Content-Disposition": f'attachment; filename="{name}-{len(rows)}.csv"'}) | |