[arpitg1304]

What

Embed the episodes in a catalog and search them by natural language. Uses SigLIP (a shared image–text model), so text queries match episode video, not just metadata.

pip install "forge-robotics[embed]"
forge embed --catalog ./forge-catalog                       # vectors per episode (vision + text)
forge search "picks up the red cup" -c ./forge-catalog --top 10
forge search --like <episode_id>    -c ./forge-catalog      # image → similar episodes

This is Phase 2 of the Forge Data Engine, building on the catalog (#30) and cloud I/O (#29).

New engine forge/embed/

file	role
base.py	EmbeddingModel interface + device selection (cuda → mps → cpu)
siglip.py	SigLIP (siglip-so400m, D=1152, shared image–text space)
registry.py	model registry; model_id = "<name>@<ckpt-hash>" (reproducible across machines)
sampling.py	~1 Hz frame sampling + pooling (mean / first-mid-last)

Catalog integration

• embeddings table — additive; bumps SCHEMA_VERSION 1 → 2 (v1 catalogs open unchanged, gain the table on first embed). One row per (episode_id, model_id, level, camera): episode-level vision per camera + instruction text.
• Catalog.add_embeddings() / search() (text + --like, single-model_id enforcement, DuckDB array_cosine_similarity); forge/catalog/embed.py backfill; CLI forge embed / forge search / forge ingest --embed.

Design highlights

• Vectors as list<float32> (+ dim) so one table holds multiple models; search filters to one model_id, casts to FLOAT[dim]. Vector spaces never mix (enforced).
• Reuses existing internals — readers for frames, forge.io for cloud sources, [video] for decode, the catalog writer/commit machinery. No existing behavior changed.
• Device auto, with Apple MPS first-class (PYTORCH_ENABLE_MPS_FALLBACK, fp32). Same command on a CUDA box.
• Base-install safe — torch/transformers live behind the [embed] extra; forge --help never imports them.

Verified on real data

Ran end-to-end on droid_100 in MinIO: 100 episodes embedded on the Mac's GPU (300 vision + 46 text vectors), and text search ranks the right episodes — "close the drawer" → the episode captioned "Close the drawer", matched by its video. (SigLIP cosines are small in magnitude by design; the ranking is what's semantic.)

The in-repo catalog_example_droid_100/ now ships those vectors (copied and remapped by content_hash, not regenerated), so forge search works from a fresh checkout — --like with no deps, text search with [embed]. The example stays excluded from the built wheel.

Testing

• tests/test_embed.py — a torch-free deterministic fake model drives the full engine / schema / search-ranking / single-model-enforcement / idempotent-re-embed / --like / memory:// cloud / CLI paths. A real-SigLIP smoke test (guarded by importorskip + FORGE_RUN_SIGLIP=1) asserts D=1152 — it caught a transformers-5.x output-shape change during the demo.
• Full existing suite: no regressions.

Docs

README "Semantic search" section, forge/embed/README.md (models, device selection, reproducibility), the updated example README, and a CHANGELOG entry.

🤖 Generated with Claude Code