"""Data access for the NayanaOCRBench explorer. Everything is read from DATA_DIR (the mounted bucket).""" import difflib, html, json, os from functools import lru_cache import duckdb import polars as pl DATA_DIR = os.environ.get("NAYANA_DATA", "/data") if not os.path.exists(os.path.join(DATA_DIR, "index", "samples.parquet")): # local dev fallback _local = os.path.join(os.path.dirname(__file__), "..", "nayanaocr_corpus_2026_preview", ".work", "bucket_mirror") if os.path.exists(os.path.join(_local, "index", "samples.parquet")): DATA_DIR = os.path.abspath(_local) IDX = os.path.join(DATA_DIR, "index") IMG = os.path.join(DATA_DIR, "images") LANG_NAME = {"en": "English", "zh": "Chinese", "hi": "Hindi", "bn": "Bengali", "mr": "Marathi", "gu": "Gujarati", "pa": "Punjabi", "or": "Odia", "kn": "Kannada", "ta": "Tamil", "te": "Telugu", "ml": "Malayalam", "sa": "Sanskrit", "ja": "Japanese", "ko": "Korean", "de": "German", "fr": "French", "es": "Spanish", "it": "Italian", "ru": "Russian", "ar": "Arabic", "th": "Thai"} LANG_ORDER = ["en", "hi", "bn", "mr", "gu", "pa", "or", "kn", "ta", "te", "ml", "sa", "zh", "ja", "ko", "de", "fr", "es", "it", "ru", "ar", "th"] GROUPS = [("Indic", ["hi", "bn", "mr", "gu", "pa", "or", "kn", "ta", "te", "ml", "sa"]), ("East Asian", ["zh", "ja", "ko"]), ("European", ["en", "de", "fr", "es", "it", "ru"]), ("Other", ["ar", "th"])] INK, INK2, RULE, EMPTY = "var(--body-text-color)", "var(--body-text-color-subdued)", "var(--border-color-primary)", "var(--background-fill-secondary)" METRICS = { # column -> (label, higher_is_better, max for colour scale) "accuracy": ("Accuracy % (100 − CER)", True, 100), "composite": ("Composite %", True, 100), "f1": ("F1 %", True, 100), "rougeL": ("ROUGE-L", True, 1), "cer": ("CER (lower is better)", False, 100), "wer": ("WER (lower is better)", False, 3), } RAMP = ["#e3ecf9", "#c7d9f3", "#a9c5ec", "#8bb0e4", "#6d9adb", "#4f83d0", "#3a6bbf", "#2a539f", "#1a3b78"] TYPE_COLOR = {"title": "#eb6834", "text": "#2a78d6", "text_block": "#2a78d6", "list": "#1baf7a", "table": "#eda100", "table_body": "#eda100", "table_caption": "#eda100", "table_footnote": "#eda100", "figure": "#e87ba4", "image": "#e87ba4", "image_body": "#e87ba4", "image_caption": "#e87ba4", "index": "#4a3aa7", "equation": "#e34948", "header": "#008300", "footer": "#008300", "page_number": "#008300", "abandon": "#8c8a93", "code": "#e34948"} _con = duckdb.connect(database=":memory:") _con.execute(f"CREATE VIEW samples AS SELECT * FROM read_parquet('{os.path.join(IDX, 'samples.parquet')}')") _con.execute(f"CREATE VIEW runs AS SELECT * FROM read_parquet('{os.path.join(IDX, 'runs.parquet')}')") _con.execute(f"CREATE VIEW gt AS SELECT * FROM read_parquet('{os.path.join(IDX, 'gt.parquet')}')") def q(sql, params=None) -> pl.DataFrame: # DuckDB connections are not thread-safe; Gradio runs handlers in a thread pool, so use a cursor per query. cur = _con.cursor() try: return cur.execute(sql, params or []).pl() finally: cur.close() @lru_cache(maxsize=None) def models(split: str) -> list[str]: return q("SELECT DISTINCT model_short FROM runs WHERE split=? ORDER BY model_short", [split])["model_short"].to_list() @lru_cache(maxsize=None) def langs(split: str) -> list[str]: ls = set(q("SELECT DISTINCT lang FROM runs WHERE split=?", [split])["lang"].to_list()) return [l for l in LANG_ORDER if l in ls] def stats() -> dict: r = q("SELECT count(*) AS runs, sum(n) AS preds, count(DISTINCT model) AS models, count(DISTINCT lang) AS langs, sum(spend_usd) AS spend FROM runs").row(0, named=True) g = q("SELECT split, count(*) AS n FROM gt GROUP BY split").rows() r["gt"] = dict(g) return r # ---------------------------------------------------------------- leaderboard def matrix(split: str, metric: str) -> pl.DataFrame: """rows = models, columns = langs, values = run-level mean of metric.""" df = q(f"SELECT model_short, family, lang, {metric} AS v, n FROM runs WHERE split=?", [split]) return df def leaderboard_table(split: str, metric: str, sort: str) -> pl.DataFrame: df = matrix(split, metric) ls = langs(split) higher = METRICS[metric][1] wide = df.pivot(on="lang", index=["model_short", "family"], values="v") for l in ls: if l not in wide.columns: wide = wide.with_columns(pl.lit(None).cast(pl.Float64).alias(l)) wide = wide.select(["model_short", "family"] + ls) fill = 0.0 if higher else None vals = wide.select(ls) wide = wide.with_columns([ pl.mean_horizontal([pl.col(l).fill_null(fill) for l in ls]).alias("avg_all") if higher else pl.mean_horizontal([pl.col(l) for l in ls]).alias("avg_all"), pl.mean_horizontal([pl.col(l) for l in ls]).alias("avg_eval"), pl.sum_horizontal([pl.col(l).is_not_null().cast(pl.Int32) for l in ls]).alias("n_langs"), ]) key = {"Average over all 22 languages": "avg_all", "Average over evaluated languages": "avg_eval", "Indic average": None, "Family": "family"}[sort] if sort == "Indic average": indic = [l for l in GROUPS[1][1] if l in ls] wide = wide.with_columns(pl.mean_horizontal([pl.col(l) for l in indic]).alias("_k")) key = "_k" wide = wide.sort(key, descending=higher if key != "family" else False, nulls_last=True) return wide.drop([c for c in ["_k"] if c in wide.columns]) def heatmap_svg(split: str, metric: str, sort: str) -> str: wide = leaderboard_table(split, metric, sort) ls = langs(split) label, higher, vmax = METRICS[metric] groups = [] small = [] for g, gl in GROUPS: kept = [l for l in gl if l in ls] if g == "Other": small.extend(kept) elif len(kept) >= 2: groups.append((g, kept)) else: small.extend(kept) if small: groups.append(("Other", small)) order = [l for g in groups for l in g[1]] left, top, cw, ch = 150, 50, 44, 28 W, H = left + len(order) * cw + 130, top + wide.height * ch + 6 s = [f'
'] x = left for g, gl in groups: w = len(gl) * cw s.append(f'{g.upper()}') x += w for j, l in enumerate(order): s.append(f'{l}') s.append(f'avg/{len(ls)} · eval-only') for i, row in enumerate(wide.iter_rows(named=True)): y = top + i * ch s.append(f'{html.escape(row["model_short"])}') for j, l in enumerate(order): v = row.get(l) cx = left + j * cw if v is None: s.append(f'—') else: t = (v / vmax) if higher else 1 - min(v, vmax) / vmax step = max(0, min(8, int(t * 9))) txt = f"{v:.0f}" if vmax >= 10 else f"{v:.2f}" s.append(f'{html.escape(row["model_short"])} · {LANG_NAME.get(l, l)} · {label}: {v:.2f}' f'= 5 else "#15140f"}" style="pointer-events:none">{txt}') ax = left + len(order) * cw + 10 a1, a2 = row["avg_all"], row["avg_eval"] f = (lambda v: "—" if v is None else (f"{v:.1f}" if vmax >= 10 else f"{v:.3f}")) s.append(f'{f(a1)}{f(a2)}') s.append("
") return "".join(s) # ---------------------------------------------------------------- explore def sample_list(split: str, lang: str, model: str, sort: str, lo: float, hi: float, search: str, limit: int = 120) -> pl.DataFrame: order = {"Worst first": "accuracy ASC", "Best first": "accuracy DESC", "Longest output": "pred_chars DESC", "Sample id": "sample_id"}[sort] where = "s.split=? AND s.lang=? AND s.model_short=? AND s.accuracy BETWEEN ? AND ?" params = [split, lang, model, lo, hi] if search: where += " AND s.sample_id ILIKE ?" params.append(f"%{search}%") return q(f"""SELECT s.sample_id, s.accuracy, s.cer, s.wer, s.f1, s.composite, s.pred_chars, s.gt_chars, g.has_image FROM samples s LEFT JOIN gt g USING (split, lang, sample_id) WHERE {where} ORDER BY {order} NULLS LAST LIMIT {limit}""", params) def image_path(split: str, lang: str, sid: str, thumb=False) -> str | None: p = os.path.join(IMG, "thumbs", split, lang, sid + ".webp") if thumb else os.path.join(IMG, split, lang, sid + ".webp") return p if os.path.exists(p) else None def sample_detail(split: str, lang: str, sid: str) -> dict: g = q("SELECT * FROM gt WHERE split=? AND lang=? AND sample_id=?", [split, lang, sid]) preds = q("SELECT model_short, family, prediction, gt_trace, accuracy, cer, wer, f1, precision, recall, composite, rougeL, pred_chars, gt_chars, in_tok, out_tok FROM samples WHERE split=? AND lang=? AND sample_id=? ORDER BY accuracy DESC NULLS LAST", [split, lang, sid]) return {"gt": g.row(0, named=True) if g.height else None, "preds": preds} def overlay_html(img_path: str, layout_json: str, show_boxes: bool, width: int | None, height: int | None) -> str: """Image with SVG layout overlay (boxes are in a 1000x1000 normalised space).""" if not img_path: return "
No image for this sample in the bucket.
" import base64 with open(img_path, "rb") as f: b64 = base64.b64encode(f.read()).decode() boxes = "" legend = "" if show_boxes: try: layout = json.loads(layout_json or "[]") except Exception: # noqa: BLE001 layout = [] seen = {} for r in layout: bb = r.get("bbox") or [] if len(bb) != 4: continue c = TYPE_COLOR.get(r.get("type"), "#8c8a93") seen[r.get("type")] = c boxes += f'{html.escape(str(r.get("type")))}' legend = " ".join(f'{html.escape(str(t))}' for t, c in seen.items()) return f"""
{boxes}
{legend}
""" def diff_html(gt: str, pred: str, max_chars: int = 6000, rtl: bool = False) -> str: """Side-by-side GT vs prediction with character-level diff highlighting.""" gt = (gt or "")[:max_chars] pred = (pred or "")[:max_chars] sm = difflib.SequenceMatcher(None, gt, pred, autojunk=False) L, R = [], [] for op, i1, i2, j1, j2 in sm.get_opcodes(): a, b = html.escape(gt[i1:i2]), html.escape(pred[j1:j2]) if op == "equal": L.append(a); R.append(b) elif op == "delete": L.append(f'{a}') elif op == "insert": R.append(f'{b}') else: L.append(f'{a}'); R.append(f'{b}') ratio = sm.ratio() css = ("") d = ' dir="rtl" style="text-align:right"' if rtl else "" return (f'{css}
' f'

Ground truth · {len(gt):,} chars · red = missing in prediction

{"".join(L)}
' f'

Prediction · {len(pred):,} chars · similarity {ratio:.2f} · green = extra / changed

{"".join(R)}
') # ---------------------------------------------------------------- compare / language def compare(split: str, lang: str, a: str, b: str) -> pl.DataFrame: return q("""SELECT x.sample_id, x.accuracy AS acc_a, y.accuracy AS acc_b, x.accuracy - y.accuracy AS delta, x.pred_chars AS chars_a, y.pred_chars AS chars_b FROM samples x JOIN samples y USING (split, lang, sample_id) WHERE x.split=? AND x.lang=? AND x.model_short=? AND y.model_short=? ORDER BY abs(delta) DESC""", [split, lang, a, b]) def language_view(split: str, lang: str) -> pl.DataFrame: return q("""SELECT model_short AS model, family, count(*) AS n, round(avg(accuracy),1) AS acc_mean, round(median(accuracy),1) AS acc_median, round(quantile_cont(accuracy, 0.1),1) AS acc_p10, round(avg(cer),1) AS cer, round(avg(wer),2) AS wer, round(avg(f1),1) AS f1, sum(CASE WHEN pred_chars IS NULL OR pred_chars < 5 THEN 1 ELSE 0 END) AS empty_out, sum(CASE WHEN accuracy < 10 THEN 1 ELSE 0 END) AS below_10, sum(CASE WHEN pred_chars > 3*gt_chars THEN 1 ELSE 0 END) AS runaway_out FROM samples WHERE split=? AND lang=? GROUP BY 1,2 ORDER BY acc_mean DESC""", [split, lang]) def hardest_pages(split: str, lang: str, k: int = 24) -> pl.DataFrame: return q("""SELECT sample_id, round(avg(accuracy),1) AS mean_acc, round(max(accuracy),1) AS best_acc, count(*) AS n_models FROM samples WHERE split=? AND lang=? GROUP BY sample_id HAVING count(*) >= 3 ORDER BY mean_acc ASC LIMIT ?""", [split, lang, k])