Spaces:
Running
Running
| """Visualize: ask a question about a dataset, get a chart and the rows behind every mark. | |
| Built for small datasets first, the 25 to 200 rows a hosted run writes. Every row stays in | |
| memory, charts are exact counts over all rows (nothing is sampled), and when a question needs | |
| something no column has (tone, whether a reply gives steps), a model reads the rows in a few | |
| batched calls and the answer becomes a new column that later questions can reuse. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| import re | |
| import statistics | |
| import time | |
| from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait | |
| from typing import Callable, Optional | |
| ID_NAMES = {"id", "_id", "idx", "index", "row_idx", "uuid", "row_id"} | |
| CLASS_NAMES = {"label", "labels", "class", "category", "target", "sentiment", "rating", "stars", "grade"} | |
| MAX_COLUMNS = 30 | |
| MAX_CELL = 4000 # characters kept per cell | |
| LABEL_CHARS = 600 # characters of each row a labelling call reads | |
| LABEL_BATCH = 8 | |
| MAX_GROUPS = 12 # categories on one axis before the rest fold into "Other" | |
| MAX_COLORS = 6 | |
| KINDS = ("label", "score", "pattern", "count") | |
| HEDGE_AFTER = 8.0 # seconds before a slow call gets a twin; the first good answer wins | |
| _HEDGE = ThreadPoolExecutor(max_workers=24) | |
| CHARTS = ("bar", "stack", "box", "hist", "scatter") | |
| class PlanError(ValueError): | |
| pass | |
| # ------------------------------------------------------------------ rows and columns | |
| def _turns(v) -> Optional[list]: | |
| """Chat-style cells: [{"role", "content"}] or ShareGPT [{"from", "value"}].""" | |
| if isinstance(v, list) and v and all(isinstance(x, dict) for x in v): | |
| if all("content" in x for x in v): | |
| return [(str(x.get("role") or "turn"), x.get("content")) for x in v] | |
| if all("value" in x for x in v): | |
| return [(str(x.get("from") or "turn"), x.get("value")) for x in v] | |
| return None | |
| def _cell(v): | |
| if v is None: | |
| return "" | |
| if isinstance(v, bool): | |
| return "yes" if v else "no" | |
| if isinstance(v, (int, float)): | |
| return v if math.isfinite(v) else "" | |
| if isinstance(v, str): | |
| return v[:MAX_CELL] | |
| t = _turns(v) | |
| if t: | |
| return "\n".join(f"{r}: {c if isinstance(c, str) else json.dumps(c, ensure_ascii=False)}" for r, c in t)[:MAX_CELL] | |
| return json.dumps(v, ensure_ascii=False)[:MAX_CELL] | |
| def _num(v) -> Optional[float]: | |
| if isinstance(v, bool): | |
| return None | |
| if isinstance(v, (int, float)): | |
| return float(v) if math.isfinite(v) else None | |
| if isinstance(v, str) and v.strip(): | |
| try: | |
| f = float(v.strip().replace(",", "")) | |
| return f if math.isfinite(f) else None | |
| except ValueError: | |
| return None | |
| return None | |
| def words(text) -> int: | |
| return len(str(text or "").split()) | |
| def prepare(raw: list[dict]) -> tuple[list[dict], dict]: | |
| """Normalize raw rows and type every column: text, category or number. | |
| Returns (rows, columns). Each row gets `_id` (1-based). Long text columns get a derived | |
| "<name> words" number column, and chat columns a "<name> turns" one. | |
| """ | |
| names: list[str] = [] | |
| for r in raw: | |
| for k in r: | |
| if k not in names and str(k).lower() not in ID_NAMES and not str(k).startswith("_"): | |
| names.append(k) | |
| names = names[:MAX_COLUMNS] | |
| rows, turns = [], {} | |
| for i, r in enumerate(raw, 1): | |
| row = {"_id": i} | |
| for k in names: | |
| v = r.get(k) | |
| t = _turns(v) | |
| if t: | |
| turns.setdefault(k, {})[i] = len(t) | |
| row[k] = _cell(v) | |
| rows.append(row) | |
| cols: dict[str, dict] = {} | |
| n = len(rows) | |
| for k in names: | |
| vals = [row[k] for row in rows if row[k] != ""] | |
| if not vals: | |
| continue | |
| nums = [_num(v) for v in vals] | |
| few = all(x is not None for x in nums) and all(float(x).is_integer() for x in nums) and \ | |
| (len(set(nums)) <= 2 or (str(k).lower() in CLASS_NAMES and len(set(nums)) <= 12)) | |
| if few: # 0/1 flags and class ids are categories, not amounts | |
| for row in rows: | |
| row[k] = str(int(_num(row[k]))) if row[k] != "" else "" | |
| cols[k] = {"type": "cat"} | |
| continue | |
| if all(x is not None for x in nums) and not all(isinstance(v, str) and len(v) > 12 for v in vals): | |
| for row in rows: | |
| row[k] = _num(row[k]) if row[k] != "" else None | |
| cols[k] = {"type": "num"} | |
| continue | |
| strs = [str(v) for v in vals] | |
| avg = sum(len(s) for s in strs) / len(strs) | |
| distinct = len(set(strs)) | |
| if distinct == len(strs) and len(strs) > 1 and not any(" " in x for x in strs) and k not in turns: | |
| continue # hashes, ids, urls: unique and never a phrase | |
| if avg > 60 or k in turns: | |
| cols[k] = {"type": "text"} | |
| elif distinct <= 30 and (distinct <= max(6, 0.6 * len(strs)) if len(strs) > 6 else distinct < len(strs)): | |
| cols[k] = {"type": "cat"} | |
| elif avg >= 12: | |
| cols[k] = {"type": "text"} | |
| else: | |
| continue # short and unique: names, ids, codes | |
| for row in rows: | |
| if not isinstance(row[k], str): | |
| row[k] = str(row[k]) | |
| for k, c in list(cols.items()): | |
| if c["type"] != "text": | |
| continue | |
| if k in turns: | |
| name = f"{k} turns" | |
| for row in rows: | |
| row[name] = float(turns[k].get(row["_id"], 0)) | |
| cols[name] = {"type": "num", "derived": "turns", "from": k} | |
| avg_words = sum(words(row[k]) for row in rows) / max(1, n) | |
| if avg_words >= 8: | |
| name = f"{k} words" | |
| for row in rows: | |
| row[name] = float(words(row[k])) | |
| cols[name] = {"type": "num", "derived": "words", "from": k} | |
| order = {"text": 0, "cat": 1, "num": 2} | |
| cols = dict(sorted(cols.items(), key=lambda kv: (order[kv[1]["type"]], bool(kv[1].get("derived"))))) | |
| return rows, cols | |
| def describe(rows: list[dict], cols: dict) -> list[dict]: | |
| """Per-column facts for the UI and the planner.""" | |
| out = [] | |
| for k, c in cols.items(): | |
| d = {"name": k, "type": c["type"], "derived": c.get("derived", ""), "from": c.get("from", "")} | |
| vals = [r.get(k) for r in rows] | |
| if c["type"] == "cat": | |
| counts: dict[str, int] = {} | |
| for v in vals: | |
| counts[v or "(blank)"] = counts.get(v or "(blank)", 0) + 1 | |
| d["values"] = sorted(counts.items(), key=lambda kv: -kv[1])[:12] | |
| d["distinct"] = len(counts) | |
| elif c["type"] == "num": | |
| xs = [v for v in vals if isinstance(v, (int, float))] | |
| if xs: | |
| d.update(min=min(xs), max=max(xs), median=statistics.median(xs)) | |
| else: | |
| xs = [words(v) for v in vals if v] | |
| d["avg_words"] = round(sum(xs) / max(1, len(xs))) | |
| d["example"] = next((str(v)[:140] for v in vals if v), "") | |
| if c.get("labels"): | |
| d["labels"] = c["labels"] | |
| if c.get("guide"): | |
| d["guide"] = c["guide"] | |
| out.append(d) | |
| return out | |
| # ------------------------------------------------------------------ charts | |
| def _q(xs: list[float], p: float) -> float: | |
| s = sorted(xs) | |
| if len(s) == 1: | |
| return s[0] | |
| pos = (len(s) - 1) * p | |
| lo = int(pos) | |
| return s[lo] + (s[min(lo + 1, len(s) - 1)] - s[lo]) * (pos - lo) | |
| def _nice(x: float) -> float: | |
| if x <= 0: | |
| return 1.0 | |
| e = 10 ** math.floor(math.log10(x)) | |
| for m in (1, 2, 2.5, 5, 10): | |
| if x <= m * e: | |
| return m * e | |
| return 10 * e | |
| def _key(v) -> str: | |
| return str(v) if v not in ("", None) else "(blank)" | |
| def _fold(groups: dict[str, list[int]], limit: int) -> list[tuple[str, list[int]]]: | |
| items = sorted(groups.items(), key=lambda kv: -len(kv[1])) | |
| if len(items) <= limit: | |
| return items | |
| keep, rest = items[:limit - 1], items[limit - 1:] | |
| ids = [i for _, g in rest for i in g] | |
| return keep + [(f"Other ({len(rest)})", ids)] | |
| def chart(rows: list[dict], cols: dict, spec: dict) -> dict: | |
| """Exact aggregation of `rows` for one chart spec. Every mark carries the ids of its rows.""" | |
| t, x, y, color, agg = spec.get("type"), spec.get("x"), spec.get("y"), spec.get("color"), spec.get("agg") or "count" | |
| by = {r["_id"]: r for r in rows} | |
| out = {"type": t, "n": len(rows), "x": x, "y": y, "color": color, "agg": agg} | |
| def groups_of(col): | |
| g: dict[str, list[int]] = {} | |
| for r in rows: | |
| g.setdefault(_key(r.get(col)), []).append(r["_id"]) | |
| return g | |
| def nums(ids, col): | |
| return [(i, by[i][col]) for i in ids if isinstance(by[i].get(col), (int, float))] | |
| if t == "bar": | |
| marks = [] | |
| for k, ids in _fold(groups_of(x), MAX_GROUPS): | |
| if y and agg in ("mean", "median"): | |
| pts = nums(ids, y) | |
| if not pts: | |
| continue | |
| vs = [v for _, v in pts] | |
| val = statistics.fmean(vs) if agg == "mean" else statistics.median(vs) | |
| marks.append({"key": k, "value": round(val, 1 if abs(val) >= 10 else 2), "n": len(pts), "ids": [i for i, _ in pts]}) | |
| else: | |
| marks.append({"key": k, "value": len(ids), "n": len(ids), "ids": ids}) | |
| marks.sort(key=lambda m: -m["value"]) | |
| out["marks"] = marks | |
| elif t == "stack": | |
| colors = [k for k, _ in _fold(groups_of(color), MAX_COLORS)] | |
| color_ids = dict(_fold(groups_of(color), MAX_COLORS)) | |
| rev = {i: k for k, ids in color_ids.items() for i in ids} | |
| marks = [] | |
| for k, ids in _fold(groups_of(x), MAX_GROUPS): | |
| parts = [] | |
| for c in colors: | |
| pids = [i for i in ids if rev.get(i) == c] | |
| if pids: | |
| parts.append({"key": c, "value": len(pids), "ids": pids}) | |
| marks.append({"key": k, "n": len(ids), "ids": ids, "parts": parts}) | |
| marks.sort(key=lambda m: -m["n"]) | |
| out.update(marks=marks, colors=colors) | |
| elif t == "box": | |
| gs = _fold(groups_of(x), MAX_GROUPS) if x else [("All rows", [r["_id"] for r in rows])] | |
| marks = [] | |
| for k, ids in gs: | |
| pts = nums(ids, y) | |
| if not pts: | |
| continue | |
| vs = [v for _, v in pts] | |
| q1, med, q3 = _q(vs, .25), _q(vs, .5), _q(vs, .75) | |
| iqr = q3 - q1 | |
| lo = min(v for v in vs if v >= q1 - 1.5 * iqr) | |
| hi = max(v for v in vs if v <= q3 + 1.5 * iqr) | |
| marks.append({"key": k, "n": len(pts), "q1": q1, "median": med, "q3": q3, "lo": lo, "hi": hi, | |
| "ids": [i for i, _ in pts], "points": [[i, v] for i, v in pts]}) | |
| marks.sort(key=lambda m: -m["median"]) | |
| out["marks"] = marks | |
| elif t == "hist": | |
| pts = nums([r["_id"] for r in rows], x) | |
| vs = [v for _, v in pts] | |
| if vs and all(float(v).is_integer() for v in vs) and max(vs) - min(vs) <= 12: | |
| # small whole numbers (scores, counts): one bar per value | |
| out["discrete"] = True | |
| out["marks"] = [{"x0": v, "x1": v + 1, "value": len(ids), "ids": ids} | |
| for v in range(int(min(vs)), int(max(vs)) + 1) for ids in [[i for i, x in pts if x == v]]] | |
| elif vs: | |
| lo, hi = min(vs), max(vs) | |
| k = 6 if len(vs) <= 30 else min(14, int(math.sqrt(len(vs))) + 2) | |
| step = _nice((hi - lo) / k) if hi > lo else 1.0 | |
| start = math.floor(lo / step) * step | |
| nb = max(1, min(40, math.ceil((hi - start) / step + 1e-9))) | |
| if start + nb * step <= hi: | |
| nb += 1 | |
| bins = [{"x0": round(start + b * step, 4), "x1": round(start + (b + 1) * step, 4), "value": 0, "ids": []} for b in range(nb)] | |
| for i, v in pts: | |
| b = min(nb - 1, int((v - start) / step)) | |
| bins[b]["ids"].append(i) | |
| bins[b]["value"] += 1 | |
| out["marks"] = bins | |
| else: | |
| out["marks"] = [] | |
| elif t == "scatter": | |
| rev = {} | |
| if color: | |
| for k, ids in _fold(groups_of(color), MAX_COLORS): | |
| for i in ids: | |
| rev[i] = k | |
| out["colors"] = [k for k, _ in _fold(groups_of(color), MAX_COLORS)] | |
| marks = [{"id": r["_id"], "x": r[x], "y": r[y], "c": rev.get(r["_id"], "")} for r in rows | |
| if isinstance(r.get(x), (int, float)) and isinstance(r.get(y), (int, float))] | |
| out["marks"] = marks | |
| if len(marks) >= 3: | |
| xs, ys = [m["x"] for m in marks], [m["y"] for m in marks] | |
| try: | |
| out["r"] = round(statistics.correlation(xs, ys), 2) | |
| except statistics.StatisticsError: | |
| pass | |
| else: | |
| raise PlanError(f"Unknown chart type {t!r}.") | |
| return out | |
| def check_spec(spec: dict, cols: dict) -> dict: | |
| """Validate a chart spec against the columns, fixing the mistakes that have one obvious fix.""" | |
| s = {k: spec.get(k) for k in ("type", "x", "y", "color", "agg")} | |
| for k in ("x", "y", "color"): | |
| if s[k] in ("", "null", "none", "None"): | |
| s[k] = None | |
| typ = lambda c: (cols.get(c) or {}).get("type") | |
| for k in ("x", "y", "color"): | |
| if s[k] is not None and s[k] not in cols: | |
| raise PlanError(f"There is no column named {s[k]!r}. Columns: {', '.join(cols)}.") | |
| t = s["type"] | |
| if t not in CHARTS: | |
| raise PlanError(f"chart.type must be one of {', '.join(CHARTS)}.") | |
| if t == "bar" and typ(s["x"]) == "num": | |
| t, s["y"] = "hist", None | |
| if t == "box" and s["y"] is None and typ(s["x"]) == "num": | |
| s["x"], s["y"] = None, s["x"] | |
| if t == "stack" and not s["color"]: | |
| t = "bar" | |
| if t == "bar": | |
| if typ(s["x"]) not in ("cat",): | |
| raise PlanError("A bar chart needs x to be a category column.") | |
| if s["y"] and typ(s["y"]) != "num": | |
| raise PlanError("A bar chart's y must be a number column, or null to count rows.") | |
| s["agg"] = (s["agg"] if s["agg"] in ("mean", "median") else "mean") if s["y"] else "count" | |
| elif t == "stack": | |
| if typ(s["x"]) != "cat" or typ(s["color"]) != "cat": | |
| raise PlanError("A stacked chart needs x and color to be category columns.") | |
| s["y"], s["agg"] = None, s["agg"] if s["agg"] in ("share", "count") else "share" | |
| elif t == "box": | |
| if typ(s["y"]) != "num": | |
| raise PlanError("A box plot needs y to be a number column.") | |
| if s["x"] and typ(s["x"]) != "cat": | |
| raise PlanError("A box plot's x must be a category column or null.") | |
| elif t == "hist": | |
| if typ(s["x"]) != "num": | |
| raise PlanError("A histogram needs x to be a number column.") | |
| elif t == "scatter": | |
| if typ(s["x"]) != "num" or typ(s["y"]) != "num": | |
| raise PlanError("A scatter plot needs x and y to be number columns.") | |
| if s["color"] and typ(s["color"]) != "cat": | |
| s["color"] = None | |
| s["type"] = t | |
| return s | |
| # ------------------------------------------------------------------ instant views (no model) | |
| def _main_cat(cols: dict, rows: list[dict]) -> Optional[str]: | |
| prefer = ("area", "category", "label", "topic", "subtopic", "department", "language") | |
| cats = [k for k, c in cols.items() if c["type"] == "cat" and len({r.get(k) for r in rows}) > 1] | |
| for p in prefer: | |
| if p in cats: | |
| return p | |
| return min(cats, key=lambda k: abs(len({r.get(k) for r in rows}) - 5)) if cats else None | |
| def _main_len(cols: dict) -> Optional[str]: | |
| prefer = ("completion words", "chosen words", "response words", "output words", "answer words", "text words") | |
| nums = [k for k, c in cols.items() if c.get("derived") == "words"] | |
| for p in prefer: | |
| if p in nums: | |
| return p | |
| return nums[-1] if nums else None | |
| def glance(rows: list[dict], cols: dict) -> list[dict]: | |
| """Two or three charts that need no model, so a dataset shows something the moment it loads.""" | |
| out = [] | |
| cat, ln = _main_cat(cols, rows), _main_len(cols) | |
| if cat: | |
| out.append({"title": f"Rows by {cat}", **chart(rows, cols, {"type": "bar", "x": cat})}) | |
| if ln: | |
| spec = {"type": "box", "x": cat, "y": ln} if cat else {"type": "hist", "x": ln} | |
| out.append({"title": f"{ln.capitalize()}" + (f" by {cat}" if cat else ""), **chart(rows, cols, spec)}) | |
| if not cat: | |
| more = [k for k, c in cols.items() if c.get("derived") in ("words", "turns") and k != ln] | |
| if more: | |
| out.append({"title": more[0].capitalize(), **chart(rows, cols, {"type": "hist", "x": more[0]})}) | |
| second = [k for k, c in cols.items() if c["type"] == "cat" and k != cat and 1 < len({r.get(k) for r in rows}) <= 8] | |
| if second: | |
| k = "language" if "language" in second else "difficulty" if "difficulty" in second else second[0] | |
| out.append({"title": f"Rows by {k}", **chart(rows, cols, {"type": "bar", "x": k})}) | |
| return out[:3] | |
| def _pl(name: str) -> str: | |
| return name if name.endswith("s") else name + "s" | |
| def presets(rows: list[dict], cols: dict) -> dict[str, dict]: | |
| """Starter questions with their plans already written, so clicking one skips the planning call.""" | |
| cat, ln = _main_cat(cols, rows), _main_len(cols) | |
| texts = [k for k, c in cols.items() if c["type"] == "text"] | |
| reply = next((k for k in ("completion", "chosen", "response", "output", "answer", "messages", "conversations") if k in texts), None) | |
| ask = next((k for k in ("prompt", "instruction", "question", "input", "query") if k in texts), None) | |
| other = next((k for k in texts if k not in (reply, ask)), None) | |
| out: dict[str, dict] = {} | |
| def add(q, title, spec, derive=(), ups=()): | |
| known = dict(cols) | |
| for d in derive: | |
| known[d["name"]] = {"type": "num" if d["kind"] in ("score", "count") else "cat"} | |
| try: | |
| out[q] = {"title": title, "derive": [dict(d) for d in derive], "chart": check_spec(spec, known), "follow_ups": list(ups)} | |
| except PlanError: | |
| pass | |
| if reply: | |
| tone = {"name": "tone", "kind": "label", "from": reply, "guide": f"the overall tone of the {reply}", | |
| "labels": ["friendly", "formal", "empathetic", "technical", "blunt"]} | |
| add(f"What tone do the {_pl(reply)} take" + (f" in each {cat}?" if cat else "?"), f"Tone of the {_pl(reply)}" + (f" by {cat}" if cat else ""), | |
| {"type": "stack", "x": cat, "color": "tone", "agg": "share"} if cat else {"type": "bar", "x": "tone"}, [tone], | |
| [f"How long are the {_pl(reply)} for each tone?"] if ln else []) | |
| if ask: | |
| add(f"What do the {_pl(ask)} ask for? Group them into a few themes.", f"What the {_pl(ask)} ask for", | |
| {"type": "bar", "x": "request"}, [{"name": "request", "kind": "label", "from": ask, "guide": f"what the {ask} asks for", "labels": []}]) | |
| if not ask and not reply and other: | |
| add(f"What themes come up in the {_pl(other)}?", f"Themes in the {_pl(other)}", {"type": "bar", "x": "theme"}, | |
| [{"name": "theme", "kind": "label", "from": other, "guide": f"the main theme of the {other}", "labels": []}]) | |
| add(f"How positive or negative is each {other}?", f"Sentiment of the {_pl(other)}", {"type": "hist", "x": "sentiment"}, | |
| [{"name": "sentiment", "kind": "score", "from": other, "guide": "1 = very negative, 3 = neutral or mixed, 5 = very positive"}]) | |
| if ln and cat: | |
| add(f"Which {cat} gets the longest {_pl(ln.replace(' words', ''))}?", f"{ln.replace(' words', '').capitalize()} length by {cat}", | |
| {"type": "box", "x": cat, "y": ln}) | |
| elif ln: | |
| add(f"How long are the {_pl(ln.replace(' words', ''))}?", f"{ln.capitalize()}", {"type": "hist", "x": ln}) | |
| if reply: | |
| add(f"How many {_pl(reply)} give step-by-step instructions?", f"{_pl(reply).capitalize()} with step-by-step instructions", | |
| {"type": "bar", "x": "gives steps"}, | |
| [{"name": "gives steps", "kind": "label", "from": reply, "guide": "whether it walks the reader through numbered or ordered steps", | |
| "labels": ["yes", "no"]}]) | |
| seconds = [k for k, c in cols.items() if c["type"] == "cat" and k != cat and 1 < len({r.get(k) for r in rows}) <= 8] | |
| if cat and seconds: | |
| add(f"How does {seconds[0]} vary across {cat}?", f"{seconds[0].capitalize()} within each {cat}", | |
| {"type": "stack", "x": cat, "color": seconds[0], "agg": "share"}) | |
| return dict(list(out.items())[:4]) | |
| def _pl(name: str) -> str: | |
| return name if name.endswith("s") else name + "s" | |
| def suggestions(rows: list[dict], cols: dict) -> list[str]: | |
| return list(presets(rows, cols)) | |
| # ------------------------------------------------------------------ the model's part | |
| def first_json(text: str): | |
| """The first JSON object or list in a model reply (fences and chatter around it are ignored).""" | |
| dec = json.JSONDecoder() | |
| for i, ch in enumerate(text): | |
| if ch in "{[": | |
| try: | |
| return dec.raw_decode(text[i:])[0] | |
| except ValueError: | |
| continue | |
| raise ValueError("no JSON in the reply") | |
| def _snip(v, n=LABEL_CHARS) -> str: | |
| return re.sub(r"\s+", " ", str(v or "")).strip()[:n] | |
| def plan_prompt(question: str, name: str, rows: list[dict], cols: dict, scope: str = "", error: str = "") -> str: | |
| lines = [] | |
| for d in describe(rows, cols): | |
| if d["type"] == "cat": | |
| vals = ", ".join(f"{k} {v}" for k, v in d["values"][:8]) | |
| more = f" (+{d['distinct'] - 8} more)" if d["distinct"] > 8 else "" | |
| lines.append(f'- "{d["name"]}" (category, {d["distinct"]} values: {vals}{more})') | |
| elif d["type"] == "num": | |
| if "min" in d: | |
| lines.append(f'- "{d["name"]}" (number, {d["min"]:g} to {d["max"]:g}, median {d["median"]:g})') | |
| else: | |
| lines.append(f'- "{d["name"]}" (text, about {d["avg_words"]} words) e.g. "{_snip(d["example"], 100)}"') | |
| if d.get("derived") in ("label", "score", "pattern", "count"): | |
| lines[-1] += f' [made earlier by reading "{d["from"]}"]' | |
| texts = [k for k, c in cols.items() if c["type"] in ("text", "cat")][:5] | |
| ex = [] | |
| for r in rows[:5]: | |
| ex.append(f"#{r['_id']} " + " | ".join(f"{k}: {_snip(r.get(k), 110)}" for k in texts)) | |
| return f"""You turn a question about a dataset into a chart plan. | |
| Dataset: {name}, {len(rows)} rows{scope}. | |
| Columns: | |
| {chr(10).join(lines)} | |
| A few rows: | |
| {chr(10).join(ex)} | |
| Question: "{question}" | |
| Reply with only a JSON object: | |
| {{"title": "...", "derive": [], "chart": {{"type": "...", "x": "...", "y": null, "color": null, "agg": "count"}}, "follow_ups": ["...", "...", "..."]}} | |
| Rules: | |
| - Use the columns above when they can answer the question. They are instant. | |
| - Questions about what the text says, asks for, means or sounds like (themes, intent, tone, quality) need a derived column read from that text, unless a category column already names exactly that. | |
| - Chart a length or other number column only when the question is about length or amount. | |
| - Only when no column answers it, add one or two entries to "derive". Each becomes a new column a model fills by reading one column of every row: | |
| {{"name": "tone", "kind": "label", "from": "completion", "labels": ["formal", "friendly", "curt"], "guide": "the overall tone of the reply"}} | |
| kind "label": one of 2 to 7 short labels per row. Give "labels" yourself when the question or the rows above suggest them (faster); use [] only for open-ended themes, to have them discovered from the rows. | |
| kind "score": an integer from 1 to 5; "guide" must say what 1 and 5 mean. | |
| kind "pattern": yes or no, whether the case-insensitive regex in "pattern" matches. kind "count": how many times it matches. | |
| Use "pattern" or "count" only for literal marks (links, numbered lists, tables, code blocks, digits); judge meaning with "label" or "score". | |
| "name" is 1 to 3 lowercase words, "from" is a text or category column. | |
| - chart.type: | |
| "bar": x is a category; y null counts rows, or y is a number column with agg "mean" or "median". | |
| "stack": x and color are categories; agg "share" or "count". Use it to compare make-up across groups. | |
| "box": y is a number column; x is a category or null. Use it for how a number spreads in each group. | |
| "hist": x is a number column. | |
| "scatter": x and y are number columns; color is a category or null. | |
| A derived column can be used in the chart by its name. | |
| - "title": what the chart shows, under 60 characters, sentence case, no claims about the result. | |
| - "follow_ups": three short questions a person exploring this data might ask next.{chr(10) + chr(10) + "Your last plan was rejected: " + error + " Fix it." if error else ""}""" | |
| def _clean_name(raw: str) -> str: | |
| return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9 ]", "", re.sub(r"[_\-]+", " ", raw.lower()))).strip()[:24] | |
| def parse_plan(text: str, cols: dict) -> dict: | |
| try: | |
| p = first_json(text) | |
| except ValueError: | |
| raise PlanError("The reply was not JSON.") | |
| if not isinstance(p, dict) or not isinstance(p.get("chart"), dict): | |
| raise PlanError('The reply needs a "chart" object.') | |
| derive = [] | |
| known = dict(cols) | |
| for d in (p.get("derive") or [])[:2]: | |
| if not isinstance(d, dict): | |
| continue | |
| kind = d.get("kind") | |
| raw_name = str(d.get("name") or "") | |
| name = _clean_name(raw_name) | |
| for k in ("x", "y", "color"): # the chart may refer to the name as the model wrote it | |
| if p["chart"].get(k) is not None and p["chart"].get(k) not in cols and _clean_name(str(p["chart"][k])) == name: | |
| p["chart"][k] = name | |
| src = d.get("from") | |
| if kind not in KINDS or not name: | |
| raise PlanError(f"derive entries need a name and a kind from {', '.join(KINDS)}.") | |
| if src not in cols or cols[src]["type"] not in ("text", "cat"): | |
| raise PlanError(f'derive "from" must be a text or category column, not {src!r}.') | |
| item = {"name": name, "kind": kind, "from": src, "guide": str(d.get("guide") or "")[:200]} | |
| if kind == "label": | |
| labels = [str(x).strip()[:40] for x in (d.get("labels") or []) if str(x).strip()] | |
| item["labels"] = list(dict.fromkeys(labels))[:7] | |
| elif kind in ("pattern", "count"): | |
| pat = re.sub(r"^\(\?i\)", "", str(d.get("pattern") or "")) | |
| if not pat or len(pat) > 160 or re.search(r"\([^)]*[+*]\)[+*{]", pat): | |
| raise PlanError("pattern must be a short regex without nested repeats.") | |
| try: | |
| re.compile(pat, re.I) | |
| except re.error: | |
| raise PlanError(f"{pat!r} is not a valid regex.") | |
| item["pattern"] = pat | |
| if name in cols and cols[name].get("derived") == kind and cols[name].get("from") == src \ | |
| and (kind != "label" or not item["labels"] or set(item["labels"]) == set(cols[name].get("labels") or [])): | |
| continue # already made by an earlier question: reuse it | |
| if name in cols and cols[name].get("derived") not in KINDS: # an original column: never overwrite it | |
| new = f"{name} (model)" | |
| for k in ("x", "y", "color"): | |
| if p["chart"].get(k) == name: | |
| p["chart"][k] = new | |
| item["name"] = name = new | |
| derive.append(item) | |
| known[name] = {"type": "num" if kind in ("score", "count") else "cat"} | |
| spec = check_spec(p["chart"], known) | |
| title = re.sub(r"\s+", " ", str(p.get("title") or "")).strip()[:80] or "What the data shows" | |
| ups = [re.sub(r"\s+", " ", str(q)).strip()[:140] for q in (p.get("follow_ups") or []) if str(q).strip()][:3] | |
| return {"title": title, "derive": derive, "chart": spec, "follow_ups": ups} | |
| def fallback_plan(question: str, rows: list[dict], cols: dict) -> dict: | |
| """Used when the model cannot produce a valid plan: answer with the closest instant chart.""" | |
| q = question.lower() | |
| cat = next((k for k, c in cols.items() if c["type"] == "cat" and k.lower() in q), None) or _main_cat(cols, rows) | |
| num = next((k for k, c in cols.items() if c["type"] == "num" and k.lower().replace(" words", "") in q), None) | |
| if not num and re.search(r"\blong|length|short|words?\b", q): | |
| num = _main_len(cols) | |
| if num: | |
| spec = {"type": "box", "x": cat, "y": num} if cat else {"type": "hist", "x": num} | |
| title = f"{num.capitalize()}" + (f" by {cat}" if cat else "") | |
| elif cat: | |
| spec, title = {"type": "bar", "x": cat}, f"Rows by {cat}" | |
| else: | |
| raise PlanError("This data has no column a chart can use.") | |
| return {"title": title, "derive": [], "chart": check_spec(spec, cols), "follow_ups": [], "fallback": True} | |
| def discover_prompt(question: str, d: dict, rows: list[dict]) -> str: | |
| snips = "\n".join(f"- {_snip(r.get(d['from']), 220)}" for r in rows[:40]) | |
| return f"""Here are {min(40, len(rows))} entries from the "{d['from']}" column of a dataset. Someone asked: "{question}" | |
| We need a label for every row: "{d['name']}"{f" ({d['guide']})" if d['guide'] else ""}. | |
| Propose 3 to 7 short labels (1 to 4 words each, lowercase) that together cover these entries, most common first. Use labels a reader would understand without seeing the data. | |
| {snips} | |
| Reply with only a JSON list of strings.""" | |
| def label_prompt(d: dict, batch: list[dict]) -> str: | |
| items = "\n".join(f"[{k}] {_snip(r.get(d['from']))}" for k, r in enumerate(batch, 1)) | |
| if d["kind"] == "label": | |
| ask = f'Pick exactly one label for each item from: {", ".join(json.dumps(x) for x in d["labels"])}.' | |
| eg = ", ".join(f'"{k}": {json.dumps(d["labels"][k % len(d["labels"])])}' for k in (1, 2)) | |
| else: | |
| ask = "Give each item an integer score from 1 to 5." | |
| eg = '"1": 4, "2": 2' | |
| return f"""You are labelling rows of a dataset. | |
| {d['name'].capitalize()}: {d['guide'] or d['name']} | |
| {ask} | |
| {items} | |
| Reply with only a JSON object mapping each item number to its {"label" if d["kind"] == "label" else "score"}, like {{{eg}}}.""" | |
| def _read_labels(text: str, d: dict, k: int) -> Optional[dict]: | |
| try: | |
| obj = first_json(text) | |
| except ValueError: | |
| return None | |
| if isinstance(obj, list): | |
| obj = {str(i): v for i, v in enumerate(obj, 1)} | |
| if not isinstance(obj, dict): | |
| return None | |
| out = {} | |
| canon = {x.lower(): x for x in d.get("labels", [])} | |
| for key, v in obj.items(): | |
| m = re.search(r"\d+", str(key)) | |
| if not m or not 1 <= int(m.group()) <= k: | |
| continue | |
| i = int(m.group()) | |
| if d["kind"] == "label": | |
| s = str(v).strip().lower() | |
| hit = canon.get(s) or next((canon[c] for c in canon if s.startswith(c) or c.startswith(s)), None) | |
| if hit: | |
| out[i] = hit | |
| else: | |
| n = _num(v) | |
| if n is not None and 1 <= n <= 5: | |
| out[i] = float(round(n)) | |
| return out if len(out) >= max(1, int(0.8 * k)) else None | |
| def hedged(fn: Callable[[float], object], after: float = HEDGE_AFTER, tries: int = 2): | |
| """Call fn(temperature) and return its first truthy result. A call still running after `after` | |
| seconds gets a twin, so one slow provider does not hold up the chart. None if every try fails.""" | |
| for attempt in range(tries): | |
| futs = {_HEDGE.submit(fn, 0.1 if attempt == 0 else 0.4)} | |
| start, twin = time.time(), False | |
| while futs: | |
| done, futs = wait(futs, timeout=None if twin else max(0.05, after - (time.time() - start)), return_when=FIRST_COMPLETED) | |
| for f in done: | |
| try: | |
| got = f.result() | |
| except Exception: | |
| got = None | |
| if got: | |
| return got | |
| if not twin and (time.time() - start >= after or not futs): | |
| futs = set(futs) | {_HEDGE.submit(fn, 0.5)} | |
| twin = True | |
| return None | |
| def fill(d: dict, rows: list[dict], chat: Callable[..., str], question: str = "", | |
| progress: Optional[Callable[[int], None]] = None, workers: int = 6) -> dict: | |
| """Compute one derived column for `rows`. Returns {row id: value} and may set d["labels"].""" | |
| if d["kind"] in ("pattern", "count"): | |
| rx = re.compile(d["pattern"], re.I) | |
| vals = {} | |
| for r in rows: | |
| hits = len(rx.findall(str(r.get(d["from"]) or "")[:MAX_CELL])) | |
| vals[r["_id"]] = ("yes" if hits else "no") if d["kind"] == "pattern" else float(hits) | |
| if progress: | |
| progress(len(rows)) | |
| return vals | |
| if d["kind"] == "label" and len(d.get("labels") or []) < 2: | |
| def discover(temp): | |
| try: | |
| obj = first_json(chat(discover_prompt(question, d, rows), max_tokens=200, temperature=temp + 0.2)) | |
| except ValueError: | |
| return None | |
| if isinstance(obj, dict): # {"labels": [...]} and similar | |
| obj = next((v for v in obj.values() if isinstance(v, list)), []) | |
| labels = [str(x).strip().lower()[:40] for x in obj if isinstance(x, (str, int, float)) and str(x).strip()] | |
| labels = list(dict.fromkeys(labels))[:7] | |
| return labels if len(labels) >= 2 else None | |
| labels = hedged(discover) | |
| if not labels: | |
| raise PlanError(f'Could not find labels for "{d["name"]}". Try asking again, or name the groups you expect.') | |
| d["labels"] = labels | |
| batches = [rows[i:i + LABEL_BATCH] for i in range(0, len(rows), LABEL_BATCH)] | |
| vals: dict[int, object] = {} | |
| done = [0] | |
| def run(batch): | |
| got = hedged(lambda temp: _read_labels(chat(label_prompt(d, batch), max_tokens=12 * len(batch) + 40, temperature=temp), d, len(batch))) | |
| res = {} | |
| for k, r in enumerate(batch, 1): | |
| res[r["_id"]] = (got or {}).get(k, "unclear" if d["kind"] == "label" else None) | |
| done[0] += len(batch) | |
| if progress: | |
| progress(done[0]) | |
| return res | |
| with ThreadPoolExecutor(max_workers=max(1, min(workers, len(batches)))) as pool: | |
| for res in pool.map(run, batches): | |
| vals.update(res) | |
| return vals | |
| def _rows(k) -> str: | |
| return f"{k} row" if k == 1 else f"{k} rows" | |
| def _unit(col) -> str: | |
| col = str(col or "") | |
| return " words" if col.endswith(" words") else " turns" if col.endswith(" turns") else "" | |
| def facts(c: dict, limit: int = 14) -> list[str]: | |
| """The numbers behind a chart in plain lines, for the write-up (and nothing else).""" | |
| t, marks, n = c["type"], c.get("marks") or [], max(1, c["n"]) | |
| u = _unit(c.get("y") if t in ("bar", "box") else c.get("x")) | |
| out = [] | |
| if t == "bar": | |
| for m in marks[:limit]: | |
| if c["agg"] == "count": | |
| out.append(f"{m['key']}: {_rows(m['value'])} ({round(100 * m['value'] / n)}%)") | |
| else: | |
| out.append(f"{m['key']}: {c['agg']} {c['y']} {m['value']:g}{u} ({_rows(m['n'])})") | |
| elif t == "stack": | |
| for m in marks[:limit]: | |
| out.append(f"{m['key']} ({_rows(m['n'])}): " + ", ".join(f"{p['key']} {p['value']} ({round(100 * p['value'] / m['n'])}%)" for p in m["parts"])) | |
| elif t == "box": | |
| for m in marks[:limit]: | |
| out.append(f"{m['key']} ({_rows(m['n'])}): median {m['median']:g}{u}, middle half {m['q1']:g} to {m['q3']:g}{u}, range {m['lo']:g} to {m['hi']:g}{u}") | |
| elif t == "hist": | |
| for m in marks: | |
| if m["value"]: | |
| out.append(f"{c['x']} {m['x0']:g}{u}: {_rows(m['value'])}" if c.get("discrete") else f"{m['x0']:g} to {m['x1']:g}{u}: {_rows(m['value'])}") | |
| elif t == "scatter": | |
| out.append(f"{len(marks)} rows plotted") | |
| if "r" in c: | |
| out.append(f"correlation between {c['x']} and {c['y']}: r = {c['r']}") | |
| return out | |
| def summary_prompt(question: str, title: str, c: dict, derived: list[dict]) -> str: | |
| made = "".join(f'\n"{d["name"]}" was assigned by a model reading each row\'s "{d["from"]}" text.' for d in derived) | |
| return f"""Someone asked about a dataset: "{question}" | |
| The chart "{title}" covers {c['n']} rows. Its exact numbers: | |
| {chr(10).join(facts(c))}{made} | |
| Answer the question in 2 to 4 short sentences of plain prose. Use only these numbers and keep them exact, but write them into sentences; do not paste the lines above. Lead with the clearest difference. Mention groups with fewer than 5 rows at most once, as too small to read much into. No quotation marks, markdown, lists, headings or advice.""" | |
| def notes(c: dict, derived: list[dict], fallback: bool = False) -> list[str]: | |
| out = [] | |
| if fallback: | |
| out.append("The model could not plan a chart for that question, so this is the closest chart the columns give directly.") | |
| for d in derived: | |
| out.append(f'"{d["name"].capitalize()}" was assigned by a model reading each row\'s {d["from"]}. Open a few rows before relying on it.') | |
| if c["n"] < 20: | |
| out.append(f"Only {c['n']} rows, so a difference of one or two rows can change the picture.") | |
| thin = [m["key"] for m in (c.get("marks") or []) if isinstance(m, dict) and "key" in m and 0 < m.get("n", 99) < 3] | |
| if thin and c["type"] in ("bar", "box", "stack") and c.get("agg") != "count": | |
| out.append(f"{', '.join(thin[:3])} {'has' if len(thin) == 1 else 'have'} fewer than three rows.") | |
| return out | |