diff --git a/CHANGELOG.md b/CHANGELOG.md index a42079b..9ba1ca8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **Catalog dedup, curation + Forge Studio (Phase 3).** Find near-duplicate + episodes from the embeddings and curate a clean, labeled training set. + + ```bash + forge catalog dedup -c ./forge-catalog --threshold 0.97 + forge curate -c ./forge-catalog --where "overall_score > 6" \ + --dedup 0.97 --dedup-policy keep-higher-quality --label approved + forge studio -c ./forge-catalog -o studio.html + ``` + + - Two new tables (bumps catalog `SCHEMA_VERSION` 2 → 3, additive): + `dedup_edges` (near-dup pairs as *facts* — similarity, not verdicts) and + `curation_labels` (an append-log of approve/reject/hold decisions, + latest-row-wins). + - `forge catalog dedup` computes near-dup pairs (cosine over episode + embeddings, max over shared cameras; idempotent). `forge curate` applies a + WHERE filter + a dedup **policy** (`keep-higher-quality` / `keep-longer` / + `keep-first`) and labels survivors approved, dedup losers rejected. + - DuckDB views/macros: `v_curation` (latest label per episode) and + `v_dup_losers(threshold, policy)`. + - **`forge studio`** generates a self-contained, themed HTML app (Overview · + Corpus · Dedup review · Snapshot) from real catalog data and embedded video + thumbnails — one shareable file, no server. + - Reuses the Phase 2 vectors, the readers for thumbnails, and the catalog + writer/commit machinery. The per-dataset `forge dedup` (perceptual-hash) is + unchanged; catalog dedup lives under `forge catalog dedup`. + - See [forge/catalog/README.md](forge/catalog/README.md). + - **Catalog embeddings + semantic search (Phase 2).** Embed the episodes in a catalog and search them by natural language. diff --git a/README.md b/README.md index d489659..7325df4 100644 --- a/README.md +++ b/README.md @@ -308,6 +308,32 @@ forge search --like -c ./forge-catalog Vectors are versioned per model (`model_id = siglip-so400m@`) and stored in the same append-only catalog. Brute-force cosine in DuckDB is sub-second at lab scale. See [forge/embed/README.md](forge/embed/README.md) for models, device selection, and reproducibility. +### Dedup & curation + +Use the embeddings to find near-duplicate episodes (re-encodes, near-identical retakes) and curate a clean, labeled training set. Near-dup pairs are recorded as **facts** (`dedup_edges`); which episode wins is decided at curation time by **policy**. + +```bash +# Find near-duplicate pairs (cosine over episode embeddings) +forge catalog dedup -c ./forge-catalog --threshold 0.97 + +# Approve a high-quality selection, dropping dedup losers by policy +forge curate -c ./forge-catalog \ + --where "overall_score > 6 AND task = 'pick_place'" \ + --dedup 0.97 --dedup-policy keep-higher-quality --label approved +``` + +Curation is an append-log (`curation_labels`, latest-row-wins); nothing is ever deleted. Policies: `keep-higher-quality`, `keep-longer`, `keep-first`. + +### Forge Studio + +Generate a self-contained, themed HTML app to explore the catalog visually — Overview, Corpus (with thumbnails + quality rings), Dedup review (keep/reject pairs), and a Snapshot preview: + +```bash +forge studio -c ./forge-catalog -o studio.html && open studio.html +``` + +Everything is embedded (real data + video thumbnails as data URIs) — one shareable file, no server. See [forge/catalog/README.md](forge/catalog/README.md). + ## Dataset Registry A curated catalog of 23+ prominent robotics datasets — browse, search, and download by name instead of memorizing URIs. **[Browse the registry online](https://arpitg1304.github.io/forge/registry.html)** diff --git a/forge/catalog/README.md b/forge/catalog/README.md index 53c3669..0fcb3cc 100644 --- a/forge/catalog/README.md +++ b/forge/catalog/README.md @@ -108,16 +108,41 @@ for teleop data). ## Query surface -Views registered on every connection: +Views/macros registered on every connection: -- `episodes`, `quality_scores` — the raw tables. +- `episodes`, `quality_scores`, `embeddings`, `dedup_edges`, `curation_labels` + — the raw tables. - `v_latest_quality` — the quality row for the newest `scorer_version` per episode (then most recent `computed_at`). +- `v_curation` — the latest label per episode (append-log, newest wins). +- `v_dup_losers(threshold, policy)` — a table macro returning the episodes that + *lose* a near-dup pairing under a policy. ```python cat.sql("SELECT task, count(*) FROM episodes GROUP BY task") # -> pyarrow.Table ``` +## Dedup, curation & Studio (Phase 3) + +- **`dedup_edges`** — near-duplicate pairs as *facts*: `(episode_a, episode_b, + similarity, model_id, method)`, canonical order `a < b`. Which episode "wins" + is decided at curation, never baked into the edge. `forge catalog dedup` + computes them from the Phase 2 embeddings (cosine, max over shared cameras; + idempotent). Brute-force is sub-second at lab scale; past ~20k episodes an + ANN pass is the scale path (skipped-with-warning until then). +- **`curation_labels`** — an append-log of decisions (`approved` / `rejected` / + `held`); latest row wins per episode, history preserved. `forge curate` + selects with a WHERE filter, resolves near-dup losers under a **policy** + (`keep-higher-quality` / `keep-longer` / `keep-first`), and labels survivors + approved, losers rejected. +- **`forge studio`** ([studio.py](studio.py)) renders a self-contained, themed + HTML app (Overview · Corpus · Dedup review · Snapshot) from real catalog data + and embedded video thumbnails — the design system matches the Forge Studio + mockup and the `forge visualize` dark theme. One shareable file, no server. + +The per-dataset `forge dedup` (perceptual-hash → cleaned dataset) is a separate, +unchanged command; catalog-level dedup lives under `forge catalog dedup`. + ## Example [`catalog_example_droid_100/`](catalog_example_droid_100/) is a small, ready-to-query @@ -130,8 +155,9 @@ commands to reproduce it (from a local MinIO `s3://` source, or straight from forge catalog stats -c forge/catalog/catalog_example_droid_100 ``` -## Not in Phase 1 +## Not yet -Embeddings, dedup edges, curation labels, and snapshots (Phases 2–4); +Snapshots + training-set export (Phase 4: `forge snapshot create/materialize`); +frame/segment-level embeddings and Lance/ANN for >100k episodes (Phase 5); distributed execution, locking, compaction, and gc. The manifest convention is in place so gc becomes possible later, but it is not implemented. diff --git a/forge/catalog/_studio_template.py b/forge/catalog/_studio_template.py new file mode 100644 index 0000000..29ca81a --- /dev/null +++ b/forge/catalog/_studio_template.py @@ -0,0 +1,342 @@ +"""Forge Studio HTML template (design system adapted from the Studio mockup). + +`generate_studio` injects real catalog JSON at the `/*__DATA__*/` and +`/*__THUMBS__*/` placeholders and the catalog name at `__CATALOG__`. +""" + +from __future__ import annotations + +HTML = r""" + + + + +Forge Studio — __CATALOG__ + + + + + +
+
+ +
__CATALOG__
+ catalog · static export +
+ +
+
+
+
+
+
+
+ + + +""" diff --git a/forge/catalog/catalog.py b/forge/catalog/catalog.py index e34362d..cd6a782 100644 --- a/forge/catalog/catalog.py +++ b/forge/catalog/catalog.py @@ -125,6 +125,8 @@ def commit_batch( episodes: list[dict] | None = None, quality_scores: list[dict] | None = None, embeddings: list[dict] | None = None, + dedup_edges: list[dict] | None = None, + curation_labels: list[dict] | None = None, batch_id: str | None = None, ingest_date: str | None = None, ) -> str: @@ -142,6 +144,8 @@ def commit_batch( "episodes": episodes, "quality_scores": quality_scores, "embeddings": embeddings, + "dedup_edges": dedup_edges, + "curation_labels": curation_labels, } tables: dict[str, pa.Table] = {} for name, rows in staged.items(): @@ -166,6 +170,16 @@ def add_embeddings(self, rows: list[dict], *, batch_id: str | None = None) -> st """Commit a batch of ``embeddings`` rows (from an embed backfill).""" return self.commit_batch(embeddings=rows, batch_id=batch_id) + def add_dedup_edges(self, rows: list[dict], *, batch_id: str | None = None) -> str: + """Commit a batch of ``dedup_edges`` rows (near-duplicate pairs).""" + return self.commit_batch(dedup_edges=rows, batch_id=batch_id) + + def add_curation_labels( + self, rows: list[dict], *, batch_id: str | None = None + ) -> str: + """Commit a batch of ``curation_labels`` rows (append-log decisions).""" + return self.commit_batch(curation_labels=rows, batch_id=batch_id) + # -- reads -------------------------------------------------------------- def _read_base(self) -> str: @@ -264,6 +278,49 @@ def _connect(self): ) = 1 """ ) + # Latest curation label per episode (append-log; newest wins). + con.execute( + """ + CREATE VIEW v_curation AS + SELECT * FROM curation_labels + QUALIFY row_number() OVER ( + PARTITION BY episode_id ORDER BY labeled_at DESC + ) = 1 + """ + ) + # Episodes that LOSE a near-dup pairing under a policy, at a threshold. + # A table macro so it takes parameters: SELECT * FROM v_dup_losers(0.97, + # 'keep-higher-quality'). Loser = the weaker side per policy; ties break + # on the greater episode_id (deterministic). An episode losing any + # pairing is a loser (union). + con.execute( + """ + CREATE MACRO v_dup_losers(thr, policy) AS TABLE ( + WITH e AS (SELECT * FROM dedup_edges WHERE similarity >= thr), + q AS (SELECT episode_id, overall_score FROM v_latest_quality), + ep AS (SELECT episode_id, num_frames, ingested_at FROM episodes) + SELECT DISTINCT CASE + WHEN policy = 'keep-longer' THEN + CASE WHEN coalesce(pa.num_frames,0) < coalesce(pb.num_frames,0) THEN e.episode_a + WHEN coalesce(pa.num_frames,0) > coalesce(pb.num_frames,0) THEN e.episode_b + ELSE greatest(e.episode_a, e.episode_b) END + WHEN policy = 'keep-first' THEN + CASE WHEN pa.ingested_at < pb.ingested_at THEN e.episode_b + WHEN pa.ingested_at > pb.ingested_at THEN e.episode_a + ELSE greatest(e.episode_a, e.episode_b) END + ELSE -- keep-higher-quality (default) + CASE WHEN coalesce(qa.overall_score,-1) < coalesce(qb.overall_score,-1) THEN e.episode_a + WHEN coalesce(qa.overall_score,-1) > coalesce(qb.overall_score,-1) THEN e.episode_b + ELSE greatest(e.episode_a, e.episode_b) END + END AS episode_id + FROM e + LEFT JOIN q qa ON qa.episode_id = e.episode_a + LEFT JOIN q qb ON qb.episode_id = e.episode_b + LEFT JOIN ep pa ON pa.episode_id = e.episode_a + LEFT JOIN ep pb ON pb.episode_id = e.episode_b + ) + """ + ) return con def sql(self, query: str) -> pa.Table: diff --git a/forge/catalog/catalog_example_droid_100/README.md b/forge/catalog/catalog_example_droid_100/README.md index 43ae468..1abae2b 100644 --- a/forge/catalog/catalog_example_droid_100/README.md +++ b/forge/catalog/catalog_example_droid_100/README.md @@ -64,6 +64,17 @@ pip install "forge-robotics[embed]" forge search "close the drawer" -c forge/catalog/catalog_example_droid_100 --top 5 ``` +### Dedup, curation & Studio + +It also ships `dedup_edges` (24 near-dup pairs at cosine ≥ 0.9) and +`curation_labels` (84 approved / 16 dedup-rejected), so curation and the visual +app work immediately: + +```bash +forge curate -c forge/catalog/catalog_example_droid_100 --dedup 0.9 --label approved +forge studio -c forge/catalog/catalog_example_droid_100 -o studio.html && open studio.html +``` + Because it's plain Parquet, you can also read it with anything else — no Forge: ```bash diff --git a/forge/catalog/catalog_example_droid_100/_batches/2480f501d1174b1b89def3ba00b2c687/manifest-430912475719423dbaadb5277318d6f3.json b/forge/catalog/catalog_example_droid_100/_batches/2480f501d1174b1b89def3ba00b2c687/manifest-430912475719423dbaadb5277318d6f3.json new file mode 100644 index 0000000..05db80e --- /dev/null +++ b/forge/catalog/catalog_example_droid_100/_batches/2480f501d1174b1b89def3ba00b2c687/manifest-430912475719423dbaadb5277318d6f3.json @@ -0,0 +1,12 @@ +{ + "batch_id": "2480f501d1174b1b89def3ba00b2c687", + "files": [ + { + "path": "curation_labels/ingest_date=2026-07-19/part-9bbb958f3df2425ea9ba634325023438.parquet", + "rows": 100, + "table": "curation_labels" + } + ], + "ingest_date": "2026-07-19", + "schema_version": 3 +} \ No newline at end of file diff --git a/forge/catalog/catalog_example_droid_100/_batches/64d8d3c210bd4ebcb0f62f2be9a5ab4a/manifest-d6afc50dfba34121909f9b1030cad8e9.json b/forge/catalog/catalog_example_droid_100/_batches/64d8d3c210bd4ebcb0f62f2be9a5ab4a/manifest-d6afc50dfba34121909f9b1030cad8e9.json new file mode 100644 index 0000000..b03256e --- /dev/null +++ b/forge/catalog/catalog_example_droid_100/_batches/64d8d3c210bd4ebcb0f62f2be9a5ab4a/manifest-d6afc50dfba34121909f9b1030cad8e9.json @@ -0,0 +1,12 @@ +{ + "batch_id": "64d8d3c210bd4ebcb0f62f2be9a5ab4a", + "files": [ + { + "path": "dedup_edges/ingest_date=2026-07-19/part-862f6e26c4204e5ba776615aeaa012ee.parquet", + "rows": 24, + "table": "dedup_edges" + } + ], + "ingest_date": "2026-07-19", + "schema_version": 3 +} \ No newline at end of file diff --git a/forge/catalog/catalog_example_droid_100/catalog.json b/forge/catalog/catalog_example_droid_100/catalog.json index de3fc32..b5527a9 100644 --- a/forge/catalog/catalog_example_droid_100/catalog.json +++ b/forge/catalog/catalog_example_droid_100/catalog.json @@ -1,8 +1,10 @@ { "created_at": "2026-07-18T23:17:30.208440+00:00", "forge_version": "0.1.0", - "schema_version": 2, + "schema_version": 3, "tables": [ + "curation_labels", + "dedup_edges", "embeddings", "episodes", "quality_scores" diff --git a/forge/catalog/catalog_example_droid_100/curation_labels/ingest_date=2026-07-19/part-9bbb958f3df2425ea9ba634325023438.parquet b/forge/catalog/catalog_example_droid_100/curation_labels/ingest_date=2026-07-19/part-9bbb958f3df2425ea9ba634325023438.parquet new file mode 100644 index 0000000..d436ce6 Binary files /dev/null and b/forge/catalog/catalog_example_droid_100/curation_labels/ingest_date=2026-07-19/part-9bbb958f3df2425ea9ba634325023438.parquet differ diff --git a/forge/catalog/catalog_example_droid_100/dedup_edges/ingest_date=2026-07-19/part-862f6e26c4204e5ba776615aeaa012ee.parquet b/forge/catalog/catalog_example_droid_100/dedup_edges/ingest_date=2026-07-19/part-862f6e26c4204e5ba776615aeaa012ee.parquet new file mode 100644 index 0000000..c4eb662 Binary files /dev/null and b/forge/catalog/catalog_example_droid_100/dedup_edges/ingest_date=2026-07-19/part-862f6e26c4204e5ba776615aeaa012ee.parquet differ diff --git a/forge/catalog/cli.py b/forge/catalog/cli.py index b3a5e25..3980fe6 100644 --- a/forge/catalog/cli.py +++ b/forge/catalog/cli.py @@ -399,6 +399,133 @@ def search_cmd( _render_result(result, output_format) +@catalog_app.command("dedup") +def catalog_dedup_cmd( + catalog: str = typer.Option(..., "--catalog", "-c", help="Catalog root"), + threshold: float = typer.Option( + 0.97, "--threshold", "-t", help="Cosine similarity threshold (default 0.97)" + ), + model: str = typer.Option(None, "--model", "-m", help="Embedding model_id"), +) -> None: + """Find near-duplicate episode pairs from embeddings; store as dedup_edges. + + Records similarity *facts* (not verdicts) — which episode wins is decided at + curation time. Re-running is idempotent. + + Examples: + forge catalog dedup -c ./forge-catalog + forge catalog dedup -c ./forge-catalog --threshold 0.9 + """ + _require_catalog_deps() + from forge.catalog import Catalog + from forge.catalog.dedup import compute_dedup_edges + from forge.core.exceptions import ForgeError + + try: + cat = Catalog.open(catalog) + stats = compute_dedup_edges(cat, model_id=model, threshold=threshold) + except ForgeError as e: + console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(1) + console.print( + f"[green]Dedup complete[/green] ({stats.model_id}, ≥ {threshold}): " + f"[cyan]{stats.pairs_found}[/cyan] pairs found over {stats.episodes} episodes, " + f"[cyan]{stats.pairs_added}[/cyan] new" + ) + + +def curate_cmd( + catalog: str = typer.Option(..., "--catalog", "-c", help="Catalog root"), + where: str = typer.Option( + None, "--where", help="SQL filter over episodes + quality (e.g. \"overall_score > 6\")" + ), + label: str = typer.Option("approved", "--label", help="Label for survivors"), + reason: str = typer.Option(None, "--reason", help="Free-text reason"), + dedup: float = typer.Option( + None, "--dedup", help="Also drop near-dup losers at this cosine threshold" + ), + dedup_policy: str = typer.Option( + "keep-higher-quality", "--dedup-policy", + help="keep-higher-quality | keep-longer | keep-first", + ), + by: str = typer.Option("cli", "--by", help="Recorded in labeled_by"), +) -> None: + """Label a selection of episodes, optionally dropping near-duplicates. + + Appends to the curation_labels log (latest-row-wins); never deletes history. + + Examples: + forge curate -c ./cat --where "overall_score > 6 AND task='pick_place'" --label approved + forge curate -c ./cat --dedup 0.97 --dedup-policy keep-higher-quality --label approved + """ + _require_catalog_deps() + from forge.catalog import Catalog + from forge.catalog.dedup import compute_dedup_edges, curate + from forge.core.exceptions import ForgeError + + try: + cat = Catalog.open(catalog) + if dedup is not None: + # Ensure edges exist at this threshold (idempotent). + compute_dedup_edges(cat, threshold=dedup) + stats = curate( + cat, where=where, label=label, reason=reason, labeled_by=f"user:{by}", + dedup_threshold=dedup, dedup_policy=dedup_policy, + ) + except (ForgeError, ValueError) as e: + console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(1) + msg = ( + f"[green]Curated[/green]: {stats.selected} selected → " + f"[cyan]{stats.approved}[/cyan] {label}" + ) + if stats.rejected: + msg += f", [red]{stats.rejected}[/red] rejected ({stats.dedup_policy})" + console.print(msg) + + +def studio_cmd( + catalog: str = typer.Option(..., "--catalog", "-c", help="Catalog root"), + output: str = typer.Option("studio.html", "--output", "-o", help="Output HTML file"), + threshold: float = typer.Option( + 0.97, "--threshold", "-t", help="Near-dup cosine threshold for the review tab" + ), + max_thumbnails: int = typer.Option( + 120, "--max-thumbnails", help="Cap on thumbnails extracted (0 = none)" + ), +) -> None: + """Generate a self-contained Forge Studio HTML app for a catalog. + + Tabs: Overview · Corpus · Dedup review · Snapshot, with real data and video + thumbnails. Open the file in a browser. + + Examples: + forge studio -c ./forge-catalog -o studio.html + """ + _require_catalog_deps() + from forge.catalog import Catalog + from forge.catalog.dedup import compute_dedup_edges + from forge.catalog.studio import generate_studio + from forge.core.exceptions import ForgeError + + try: + cat = Catalog.open(catalog) + if cat.embedding_model_ids(): + compute_dedup_edges(cat, threshold=threshold) # idempotent + stats = generate_studio( + cat, output, threshold=threshold, + max_thumbnails=max_thumbnails, console=console, + ) + except ForgeError as e: + console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(1) + console.print( + f"[green]Studio written[/green] → [cyan]{stats.out_path}[/cyan] " + f"({stats.episodes} episodes, {stats.pairs} pairs, {stats.thumbnails} thumbnails)" + ) + console.print(f"[dim]open it:[/dim] open {stats.out_path}") + + def register_catalog_cli(app: typer.Typer) -> None: """Attach catalog commands to the main forge Typer app.""" app.add_typer(catalog_app, name="catalog") @@ -406,3 +533,5 @@ def register_catalog_cli(app: typer.Typer) -> None: app.command("query")(query_cmd) app.command("embed")(embed_cmd) app.command("search")(search_cmd) + app.command("curate")(curate_cmd) + app.command("studio")(studio_cmd) diff --git a/forge/catalog/dedup.py b/forge/catalog/dedup.py new file mode 100644 index 0000000..850a741 --- /dev/null +++ b/forge/catalog/dedup.py @@ -0,0 +1,204 @@ +"""Near-duplicate detection and policy-based curation over a catalog. + +`compute_dedup_edges` finds near-duplicate episode pairs from the Phase 2 +embeddings (cosine over per-camera vision vectors) and records them as +``dedup_edges`` — *facts*, not verdicts. `curate` then applies a WHERE filter +plus a dedup policy and appends ``curation_labels`` (approved survivors, +rejected dedup losers). + +Brute-force all-pairs cosine in numpy is sub-second at lab scale; past tens of +thousands of episodes a blocked/ANN pass is needed (logged, not silently +truncated). +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from forge.catalog.catalog import Catalog + +# Above this episode count, all-pairs O(N^2) cosine is skipped with a warning +# rather than blowing up memory. (ANN is the documented scale path.) +_MAX_BRUTE_FORCE = 20000 + +DEFAULT_THRESHOLD = 0.97 +DEFAULT_POLICY = "keep-higher-quality" +_POLICIES = ("keep-higher-quality", "keep-longer", "keep-first") + + +@dataclass +class DedupStats: + model_id: str + threshold: float + episodes: int = 0 + pairs_found: int = 0 + pairs_added: int = 0 + skipped_reason: str | None = None + + +@dataclass +class CurateStats: + label: str + selected: int = 0 + approved: int = 0 + rejected: int = 0 + dedup_policy: str | None = None + + +def compute_dedup_edges( + catalog: Catalog, + *, + model_id: str | None = None, + threshold: float = DEFAULT_THRESHOLD, + console=None, +) -> DedupStats: + """Find near-duplicate pairs from episode embeddings; append ``dedup_edges``. + + Similarity between two episodes is the **max cosine over their shared + cameras** (a dup on any view counts). Idempotent: pairs already recorded for + this ``model_id`` are not re-added. + """ + model_id = catalog._resolve_model_id(model_id) + stats = DedupStats(model_id=model_id, threshold=threshold) + + rows = catalog.sql( + "SELECT episode_id, camera, vector FROM embeddings " + f"WHERE model_id = '{model_id.replace(chr(39), chr(39) * 2)}' " + "AND level = 'episode'" + ).to_pylist() + + by_cam: dict[str, tuple[list, list]] = defaultdict(lambda: ([], [])) + all_eps: set[str] = set() + for r in rows: + ids, vecs = by_cam[r["camera"]] + ids.append(r["episode_id"]) + vecs.append(r["vector"]) + all_eps.add(r["episode_id"]) + stats.episodes = len(all_eps) + + if stats.episodes > _MAX_BRUTE_FORCE: + stats.skipped_reason = ( + f"{stats.episodes} episodes exceeds the {_MAX_BRUTE_FORCE} brute-force " + "cap; ANN-based dedup is a future scale path" + ) + if console is not None: + console.print(f"[yellow]Skipped:[/yellow] {stats.skipped_reason}") + return stats + + # (a, b) canonical -> max cosine across shared cameras. + pair_sim: dict[tuple[str, str], float] = {} + for _cam, (ids, vecs) in by_cam.items(): + if len(ids) < 2: + continue + v = np.nan_to_num(np.asarray(vecs, dtype=np.float32)) # already L2-normalized + with np.errstate(all="ignore"): # silence spurious float32 BLAS warnings + sim = v @ v.T + ii, jj = np.where(np.triu(sim >= threshold, k=1)) + for i, j in zip(ii.tolist(), jj.tolist()): + a, b = ids[i], ids[j] + if a > b: + a, b = b, a + key = (a, b) + pair_sim[key] = max(pair_sim.get(key, -1.0), float(sim[i, j])) + stats.pairs_found = len(pair_sim) + + # Skip pairs already recorded for this model. + existing = { + (r["episode_a"], r["episode_b"]) + for r in catalog.sql( + "SELECT episode_a, episode_b FROM dedup_edges " + f"WHERE model_id = '{model_id.replace(chr(39), chr(39) * 2)}'" + ).to_pylist() + } + + now = datetime.now(timezone.utc) + new_rows = [ + { + "episode_a": a, + "episode_b": b, + "similarity": sim, + "model_id": model_id, + "method": "ann-cosine", + "computed_at": now, + } + for (a, b), sim in pair_sim.items() + if (a, b) not in existing + ] + if new_rows: + catalog.add_dedup_edges(new_rows) + stats.pairs_added = len(new_rows) + return stats + + +def curate( + catalog: Catalog, + *, + where: str | None = None, + label: str = "approved", + reason: str | None = None, + labeled_by: str = "policy:curate", + dedup_threshold: float | None = None, + dedup_policy: str = DEFAULT_POLICY, +) -> CurateStats: + """Label a WHERE-selected set, dropping near-dup losers under a policy. + + Survivors get ``label`` (default ``approved``); if ``dedup_threshold`` is + given, the dedup losers within the selection get ``rejected``. Appends to the + ``curation_labels`` log (latest-row-wins), never deleting history. + """ + if dedup_policy not in _POLICIES: + raise ValueError(f"unknown dedup policy {dedup_policy!r}; choose {_POLICIES}") + + stats = CurateStats(label=label) + + base = ( + "SELECT e.episode_id FROM episodes e " + "LEFT JOIN v_latest_quality q USING(episode_id)" + ) + if where: + base += f" WHERE {where}" + selected = [r["episode_id"] for r in catalog.sql(base).to_pylist()] + stats.selected = len(selected) + if not selected: + return stats + + losers: set[str] = set() + if dedup_threshold is not None: + stats.dedup_policy = dedup_policy + loser_rows = catalog.sql( + f"SELECT episode_id FROM v_dup_losers({float(dedup_threshold)}, " + f"'{dedup_policy}')" + ).to_pylist() + all_losers = {r["episode_id"] for r in loser_rows} + losers = all_losers & set(selected) + + now = datetime.now(timezone.utc) + rows: list[dict] = [] + for ep in selected: + if ep in losers: + rows.append({ + "episode_id": ep, + "label": "rejected", + "reason": reason or f"near-duplicate ({dedup_policy})", + "labeled_by": f"policy:{dedup_policy}", + "labeled_at": now, + }) + stats.rejected += 1 + else: + rows.append({ + "episode_id": ep, + "label": label, + "reason": reason, + "labeled_by": labeled_by, + "labeled_at": now, + }) + stats.approved += 1 + + catalog.add_curation_labels(rows) + return stats diff --git a/forge/catalog/schema.py b/forge/catalog/schema.py index c72b073..652035b 100644 --- a/forge/catalog/schema.py +++ b/forge/catalog/schema.py @@ -5,9 +5,9 @@ current quality scorer version. Every write goes through :func:`validate_rows`, which builds a validated ``pyarrow.Table`` and fails loudly on any mismatch. -Phase 1 defines exactly two tables: ``episodes`` (the registry) and -``quality_scores`` (versioned derived metrics). Later phases (embeddings, dedup, -curation, snapshots) are intentionally absent. +Tables by phase: ``episodes`` + ``quality_scores`` (1), ``embeddings`` (2), +``dedup_edges`` + ``curation_labels`` (3). Snapshots (4) are JSON manifests, not +a table. """ from __future__ import annotations @@ -17,10 +17,11 @@ from forge.core.exceptions import ForgeError # Bump when the catalog table schemas change in a breaking way. catalog.json -# records this; Forge refuses to open a catalog written by a newer version. -# v2 (Phase 2) adds the additive `embeddings` table — backward compatible, so -# v1 catalogs open unchanged and gain the table on first embed. -SCHEMA_VERSION = 2 +# records this; Forge refuses to open a catalog written by a newer version. Each +# bump so far is purely additive (new tables), so older catalogs open unchanged +# and gain the new tables on first use. v2 added `embeddings`; v3 adds +# `dedup_edges` and `curation_labels`. +SCHEMA_VERSION = 3 # Version tag stamped on every quality_scores row. The quality module has no # version of its own today, so the catalog owns this constant. Bump it when the @@ -107,11 +108,39 @@ ] ) +# One row per detected near-duplicate pair. Edges record *similarity* — which +# episode "wins" is decided at curation time by policy, never baked in here +# (that would freeze the policy). Canonical order: episode_a < episode_b. +DEDUP_EDGES_SCHEMA = pa.schema( + [ + pa.field("episode_a", pa.string(), nullable=False), + pa.field("episode_b", pa.string(), nullable=False), + pa.field("similarity", pa.float32(), nullable=False), + pa.field("model_id", pa.string(), nullable=False), + pa.field("method", pa.string(), nullable=False), # ann-cosine | exact-hash + pa.field("computed_at", pa.timestamp("us", tz="UTC"), nullable=False), + ] +) + +# Append-log of human/policy decisions. Latest row wins per episode (resolved +# with a window function at query time); appending never deletes history. +CURATION_LABELS_SCHEMA = pa.schema( + [ + pa.field("episode_id", pa.string(), nullable=False), + pa.field("label", pa.string(), nullable=False), # approved | rejected | held + pa.field("reason", pa.string()), + pa.field("labeled_by", pa.string(), nullable=False), # user or policy: + pa.field("labeled_at", pa.timestamp("us", tz="UTC"), nullable=False), + ] +) + # Logical table name -> (pyarrow schema, partition column). TABLES: dict[str, tuple[pa.Schema, str]] = { "episodes": (EPISODES_SCHEMA, "ingest_date"), "quality_scores": (QUALITY_SCORES_SCHEMA, "ingest_date"), "embeddings": (EMBEDDINGS_SCHEMA, "ingest_date"), + "dedup_edges": (DEDUP_EDGES_SCHEMA, "ingest_date"), + "curation_labels": (CURATION_LABELS_SCHEMA, "ingest_date"), } diff --git a/forge/catalog/studio.py b/forge/catalog/studio.py new file mode 100644 index 0000000..c1c2848 --- /dev/null +++ b/forge/catalog/studio.py @@ -0,0 +1,243 @@ +"""Forge Studio — a self-contained, themed HTML app for a catalog. + +Generates a single shareable HTML file with tabs (Overview · Corpus · Dedup +review · Snapshot) rendered from **real catalog data** and **real video +thumbnails** (extracted via the existing readers, embedded as data URIs). The +design system matches the Forge Studio mockup and the dark theme of +`forge visualize`. + +The file is static: the Dedup tab lets you make keep/reject decisions and export +them as `forge curate`-ready commands (a static page can't write back to the +catalog). Thumbnails are capped so generation stays quick. +""" + +from __future__ import annotations + +import base64 +import html +import io +import json +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from forge.catalog.dedup import DEFAULT_THRESHOLD + +if TYPE_CHECKING: + from rich.console import Console + + from forge.catalog.catalog import Catalog + + +@dataclass +class StudioStats: + out_path: str + episodes: int = 0 + pairs: int = 0 + thumbnails: int = 0 + capped: bool = False + + +def _thumbnail_data_uri(img, size: int = 240) -> str | None: + """Downscale an HWC uint8 frame to a small JPEG data URI.""" + try: + import numpy as np + from PIL import Image + + arr = np.asarray(img) + if arr.dtype != np.uint8: + arr = arr.clip(0, 255).astype("uint8") + if arr.ndim == 2: + arr = np.stack([arr] * 3, -1) + if arr.shape[-1] == 4: + arr = arr[..., :3] + im = Image.fromarray(arr) + im.thumbnail((size, size)) + buf = io.BytesIO() + im.convert("RGB").save(buf, format="JPEG", quality=72) + return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode() + except Exception: + return None + + +def _extract_thumbnails( + catalog: Catalog, wanted: dict[str, str], console: Console | None +) -> dict[str, str]: + """Extract one representative frame per wanted episode. + + ``wanted`` maps content_hash -> catalog episode_id. Groups by source_uri, + opens each dataset once, and grabs a middle frame's first camera. + """ + from forge.catalog.ingest import _episode_content_hash, _resolve_source + from forge.formats.registry import FormatRegistry + + rows = catalog.sql("SELECT content_hash, source_uri FROM episodes").to_pylist() + by_source: dict[str, set[str]] = defaultdict(set) + for r in rows: + if r["content_hash"] in wanted: + by_source[r["source_uri"]].add(r["content_hash"]) + + thumbs: dict[str, str] = {} + for source_uri, hashes in by_source.items(): + remaining = set(hashes) + try: + path = _resolve_source(source_uri) + reader = FormatRegistry.get_reader(FormatRegistry.detect_format(path)) + for episode in reader.read_episodes(path): + if not remaining: + break + try: + frames = episode.load_frames() + ch = _episode_content_hash(episode, frames) + if ch not in remaining: + continue + remaining.discard(ch) + mid = frames[len(frames) // 2] + if mid.images: + cam = next(iter(mid.images)) + uri = _thumbnail_data_uri(mid.images[cam].load()) + if uri: + thumbs[wanted[ch]] = uri + finally: + episode.clear_cache() + except Exception: + continue + if console is not None: + console.print(f"[dim] thumbnails: {len(thumbs)}/{len(wanted)}[/dim]") + return thumbs + + +def _build_data(catalog: Catalog, threshold: float, max_thumbnails: int): + """Assemble the JSON payload + the episode set to thumbnail.""" + eps = catalog.sql( + """ + SELECT e.episode_id, e.task, e.robot, e.language_instruction AS instr, + e.num_frames, e.fps, e.duration_s, e.source_format, e.content_hash, + q.overall_score AS score, c.label AS label + FROM episodes e + LEFT JOIN v_latest_quality q USING(episode_id) + LEFT JOIN v_curation c USING(episode_id) + ORDER BY q.overall_score DESC NULLS LAST + """ + ).to_pylist() + + has_embeddings = bool(catalog.embedding_model_ids()) + pairs = [] + if has_embeddings and _table_has_rows(catalog, "dedup_edges"): + pairs = catalog.sql( + f""" + SELECT d.episode_a, d.episode_b, d.similarity AS sim, + ea.task AS task, ea.robot AS robot, + qa.overall_score AS qa, qb.overall_score AS qb + FROM dedup_edges d + JOIN episodes ea ON ea.episode_id = d.episode_a + LEFT JOIN v_latest_quality qa ON qa.episode_id = d.episode_a + LEFT JOIN v_latest_quality qb ON qb.episode_id = d.episode_b + WHERE d.similarity >= {float(threshold)} + ORDER BY d.similarity DESC + """ + ).to_pylist() + + stats = catalog.stats() + dup_pairs = _table_count(catalog, "dedup_edges") + labeled = catalog.sql( + "SELECT label, count(*) n FROM v_curation GROUP BY 1" + ).to_pylist() + + # Episodes to thumbnail: all pair episodes first, then top-quality, up to cap. + ordered = [] + for p in pairs: + ordered += [p["episode_a"], p["episode_b"]] + ordered += [e["episode_id"] for e in eps] + seen: set[str] = set() + thumb_ids: list[str] = [] + for eid in ordered: + if eid not in seen: + seen.add(eid) + thumb_ids.append(eid) + if len(thumb_ids) >= max_thumbnails: + break + + hash_by_id = {e["episode_id"]: e["content_hash"] for e in eps} + wanted = {hash_by_id[i]: i for i in thumb_ids if i in hash_by_id} + + data = { + "catalog": catalog.root_uri, + "stats": stats, + "labeled": {r["label"]: r["n"] for r in labeled}, + "episodes": [ + { + "id": e["episode_id"], + "task": e["task"], + "robot": e["robot"], + "instr": e["instr"], + "score": e["score"], + "frames": e["num_frames"], + "dur": _fmt_dur(e["duration_s"]), + "fmt": e["source_format"], + "label": e["label"], + } + for e in eps + ], + "pairs": [ + { + "a": p["episode_a"], "b": p["episode_b"], "sim": p["sim"], + "task": p["task"], "robot": p["robot"], "qa": p["qa"], "qb": p["qb"], + } + for p in pairs + ], + "dedup": {"threshold": threshold, "total_pairs": dup_pairs}, + "has_embeddings": has_embeddings, + } + return data, wanted + + +def _table_has_rows(catalog: Catalog, table: str) -> bool: + return _table_count(catalog, table) > 0 + + +def _table_count(catalog: Catalog, table: str) -> int: + try: + return int(catalog.sql(f"SELECT count(*) n FROM {table}").to_pylist()[0]["n"]) + except Exception: + return 0 + + +def _fmt_dur(seconds) -> str: + s = int(seconds or 0) + return f"{s // 60}:{s % 60:02d}" + + +def generate_studio( + catalog: Catalog, + out_path: str, + *, + threshold: float = DEFAULT_THRESHOLD, + max_thumbnails: int = 120, + console: Console | None = None, +) -> StudioStats: + """Render a Forge Studio HTML file for ``catalog``.""" + data, wanted = _build_data(catalog, threshold, max_thumbnails) + thumbs = _extract_thumbnails(catalog, wanted, console) if wanted else {} + for e in data["episodes"]: + if e["id"] in thumbs: + e["thumb"] = thumbs[e["id"]] + + payload = json.dumps(data) + thumb_json = json.dumps(thumbs) + doc = _HTML.replace("/*__DATA__*/", payload).replace("/*__THUMBS__*/", thumb_json) + doc = doc.replace("__CATALOG__", html.escape(catalog.root_uri)) + + Path(out_path).write_text(doc, encoding="utf-8") + return StudioStats( + out_path=out_path, + episodes=len(data["episodes"]), + pairs=len(data["pairs"]), + thumbnails=len(thumbs), + capped=len(wanted) >= max_thumbnails, + ) + + +# The template lives in a sibling module to keep this file readable. +from forge.catalog._studio_template import HTML as _HTML # noqa: E402 diff --git a/tests/test_catalog_dedup.py b/tests/test_catalog_dedup.py new file mode 100644 index 0000000..d14eddd --- /dev/null +++ b/tests/test_catalog_dedup.py @@ -0,0 +1,210 @@ +"""Tests for catalog dedup + curation + studio (Phase 3). + +Uses hand-built embedding vectors (no torch) so two episodes are near-identical +and reliably form a near-dup pair. +""" + +from __future__ import annotations + +import json +import re +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("pyarrow") +pytest.importorskip("duckdb") + +from forge.catalog import Catalog # noqa: E402 +from forge.catalog.dedup import compute_dedup_edges, curate # noqa: E402 +from forge.catalog.schema import ( # noqa: E402 + CURATION_LABELS_SCHEMA, + DEDUP_EDGES_SCHEMA, + CatalogSchemaError, + validate_rows, +) + +_MODEL = "fake@t" + + +def _now(): + return datetime.now(timezone.utc) + + +def _unit(vec) -> list: + v = np.asarray(vec, dtype=np.float32) + return (v / (np.linalg.norm(v) or 1.0)).tolist() + + +def _ep(i, score, **over) -> dict: + row = { + "episode_id": f"e{i}", "content_hash": f"h{i}", "source_uri": "/ds", + "source_format": "zarr", "robot": "franka", "task": "pick", + "language_instruction": f"task {i}", "operator_id": None, "scene_id": None, + "num_frames": 100 + i, "duration_s": 5.0, "fps": 10.0, "cameras": [], + "state_dim": 7, "action_dim": 7, "ingest_batch_id": "b", "ingested_at": _now(), + "extra": {}, + } + row.update(over) + return {"ep": row, "score": score} + + +def _emb(eid, vec, camera="cam0") -> dict: + return { + "episode_id": eid, "model_id": _MODEL, "level": "episode", "camera": camera, + "pooling": "mean", "dim": len(vec), "vector": list(vec), + "num_frames_sampled": 4, "computed_at": _now(), + } + + +@pytest.fixture +def dup_catalog(tmp_path: Path) -> Catalog: + """Catalog where e0 and e1 are near-identical (dup); e2 is distinct.""" + cat = Catalog.init(str(tmp_path / "cat")) + specs = [_ep(0, 8.0), _ep(1, 6.0), _ep(2, 7.0)] + cat.commit_batch( + episodes=[s["ep"] for s in specs], + quality_scores=[ + {"episode_id": s["ep"]["episode_id"], "scorer_version": "q-v1", + "overall_score": s["score"], "flags": [], "computed_at": _now()} + for s in specs + ], + batch_id="b1", + ) + # e0 ~ e1 (cosine ~1); e2 orthogonal-ish + cat.add_embeddings([ + _emb("e0", _unit([1.0, 0.0, 0.0, 0.02])), + _emb("e1", _unit([1.0, 0.01, 0.0, 0.0])), + _emb("e2", _unit([0.0, 0.0, 1.0, 0.0])), + ], batch_id="b2") + return cat + + +class TestSchemas: + def test_dedup_edge_valid_and_required(self): + row = { + "episode_a": "a", "episode_b": "b", "similarity": 0.98, + "model_id": _MODEL, "method": "ann-cosine", "computed_at": _now(), + } + assert validate_rows([row], DEDUP_EDGES_SCHEMA, table="dedup_edges").num_rows == 1 + del row["similarity"] + with pytest.raises(CatalogSchemaError, match="missing required"): + validate_rows([row], DEDUP_EDGES_SCHEMA, table="dedup_edges") + + def test_curation_label_valid(self): + row = {"episode_id": "a", "label": "approved", "reason": None, + "labeled_by": "user:x", "labeled_at": _now()} + assert validate_rows([row], CURATION_LABELS_SCHEMA, table="curation_labels").num_rows == 1 + + +class TestDedup: + def test_finds_pair_canonical_order(self, dup_catalog): + st = compute_dedup_edges(dup_catalog, threshold=0.97) + assert st.pairs_found == 1 and st.pairs_added == 1 + edge = dup_catalog.sql("SELECT episode_a, episode_b FROM dedup_edges").to_pylist()[0] + assert (edge["episode_a"], edge["episode_b"]) == ("e0", "e1") # a < b + + def test_threshold_excludes_distinct(self, dup_catalog): + # High threshold: only the true dup; e2 never pairs. + compute_dedup_edges(dup_catalog, threshold=0.97) + eps = dup_catalog.sql( + "SELECT DISTINCT episode_a FROM dedup_edges " + "UNION SELECT DISTINCT episode_b FROM dedup_edges" + ).to_pylist() + ids = {r["episode_a"] for r in eps} + assert "e2" not in ids + + def test_idempotent(self, dup_catalog): + compute_dedup_edges(dup_catalog, threshold=0.97) + st2 = compute_dedup_edges(dup_catalog, threshold=0.97) + assert st2.pairs_added == 0 + assert dup_catalog.sql("SELECT count(*) n FROM dedup_edges").to_pylist()[0]["n"] == 1 + + def test_dup_losers_keeps_higher_quality(self, dup_catalog): + compute_dedup_edges(dup_catalog, threshold=0.97) + # e0 (q=8) vs e1 (q=6) -> e1 loses + losers = dup_catalog.sql( + "SELECT episode_id FROM v_dup_losers(0.97, 'keep-higher-quality')" + ).to_pylist() + assert [r["episode_id"] for r in losers] == ["e1"] + + def test_dup_losers_keep_longer(self, dup_catalog): + compute_dedup_edges(dup_catalog, threshold=0.97) + # e1 has more frames (101) than e0 (100) -> e0 loses under keep-longer + losers = {r["episode_id"] for r in dup_catalog.sql( + "SELECT episode_id FROM v_dup_losers(0.97, 'keep-longer')").to_pylist()} + assert losers == {"e0"} + + +class TestCurate: + def test_approves_and_rejects_by_policy(self, dup_catalog): + compute_dedup_edges(dup_catalog, threshold=0.97) + st = curate(dup_catalog, dedup_threshold=0.97, dedup_policy="keep-higher-quality", + label="approved") + assert st.selected == 3 + assert st.rejected == 1 # e1 dropped + assert st.approved == 2 # e0, e2 + labels = {r["episode_id"]: r["label"] for r in dup_catalog.sql( + "SELECT episode_id, label FROM v_curation").to_pylist()} + assert labels == {"e0": "approved", "e1": "rejected", "e2": "approved"} + + def test_where_filter(self, dup_catalog): + st = curate(dup_catalog, where="overall_score >= 7.0", label="approved") + assert st.selected == 2 # e0(8), e2(7) + assert st.approved == 2 + + def test_latest_label_wins(self, dup_catalog): + curate(dup_catalog, where="episode_id = 'e0'", label="held") + curate(dup_catalog, where="episode_id = 'e0'", label="approved") + rows = dup_catalog.sql( + "SELECT label FROM v_curation WHERE episode_id='e0'").to_pylist() + assert len(rows) == 1 and rows[0]["label"] == "approved" + # history preserved + assert dup_catalog.sql( + "SELECT count(*) n FROM curation_labels WHERE episode_id='e0'" + ).to_pylist()[0]["n"] == 2 + + def test_unknown_policy_raises(self, dup_catalog): + with pytest.raises(ValueError, match="unknown dedup policy"): + curate(dup_catalog, dedup_threshold=0.97, dedup_policy="bogus") + + +class TestStudio: + def test_generate_html(self, dup_catalog): + from forge.catalog.studio import generate_studio + + compute_dedup_edges(dup_catalog, threshold=0.97) + out = str(Path(dup_catalog._root) / ".." / "studio.html") + st = generate_studio(dup_catalog, out, threshold=0.97, max_thumbnails=0) + assert st.episodes == 3 and st.pairs == 1 + h = Path(out).read_text() + # embedded data is valid JSON and carries the real rows + data = json.loads(re.search(r"const FORGE_DATA = (\{.*?\});\nconst THUMBS", h, re.S).group(1)) + assert len(data["episodes"]) == 3 + assert len(data["pairs"]) == 1 + assert "Dedup review" in h and "FORGE_DATA" in h + + +class TestCLI: + def test_dedup_curate_studio(self, tmp_path: Path, dup_catalog): + from typer.testing import CliRunner + + from forge.cli import app + + runner = CliRunner() + root = dup_catalog.root_uri + + r = runner.invoke(app, ["catalog", "dedup", "-c", root, "-t", "0.97"]) + assert r.exit_code == 0, r.output + assert "pairs found" in r.output + + r = runner.invoke(app, ["curate", "-c", root, "--dedup", "0.97", "--label", "approved"]) + assert r.exit_code == 0, r.output + assert "rejected" in r.output + + out = str(tmp_path / "s.html") + r = runner.invoke(app, ["studio", "-c", root, "-o", out, "--max-thumbnails", "0"]) + assert r.exit_code == 0, r.output + assert Path(out).exists()