Spaces:
Sleeping
Sleeping
SudharshanR commited on
Commit ·
a1dd5ba
0
Parent(s):
ElideDB query by example: no text, no model at query time
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .dockerignore +9 -0
- .gitattributes +4 -0
- Dockerfile +34 -0
- README.md +38 -0
- deploy/entrypoint-qbe.sh +16 -0
- deploy/qbe_serve.py +189 -0
- deploy/qbe_ui.html +198 -0
- deploy/requirements-qbe.txt +13 -0
- python/.DS_Store +0 -0
- python/elidedb.egg-info/PKG-INFO +123 -0
- python/elidedb.egg-info/SOURCES.txt +20 -0
- python/elidedb.egg-info/dependency_links.txt +1 -0
- python/elidedb.egg-info/entry_points.txt +2 -0
- python/elidedb.egg-info/requires.txt +18 -0
- python/elidedb.egg-info/top_level.txt +1 -0
- python/elidedb/__init__.py +14 -0
- python/elidedb/action_channel.py +77 -0
- python/elidedb/action_probe.py +190 -0
- python/elidedb/ann.py +154 -0
- python/elidedb/answer.py +338 -0
- python/elidedb/bptree.py +123 -0
- python/elidedb/chanbench.py +221 -0
- python/elidedb/cli.py +293 -0
- python/elidedb/context.py +1112 -0
- python/elidedb/corpus.py +96 -0
- python/elidedb/cracked.py +133 -0
- python/elidedb/ctxprune.py +516 -0
- python/elidedb/ctxtower.py +554 -0
- python/elidedb/data/ssv2_classes.txt +174 -0
- python/elidedb/derive.py +148 -0
- python/elidedb/desk.py +1723 -0
- python/elidedb/desk_ui.html +0 -0
- python/elidedb/device.py +56 -0
- python/elidedb/dinov3.py +91 -0
- python/elidedb/distill.py +176 -0
- python/elidedb/embeddings.py +744 -0
- python/elidedb/encprune.py +307 -0
- python/elidedb/fdnnstudent.py +231 -0
- python/elidedb/fdnnv2.py +249 -0
- python/elidedb/fdnnvideo.py +500 -0
- python/elidedb/fftools.py +40 -0
- python/elidedb/fusion.py +110 -0
- python/elidedb/grounding.py +218 -0
- python/elidedb/identity.py +915 -0
- python/elidedb/ingest.py +158 -0
- python/elidedb/itm.py +175 -0
- python/elidedb/iv2.py +133 -0
- python/elidedb/lexicon.py +61 -0
- python/elidedb/log.py +293 -0
- python/elidedb/migrate.py +163 -0
.dockerignore
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Only what the image needs enters the build context. The lake and
|
| 2 |
+
# datasets are tens of GB; the demo store ships from deploy/demo.
|
| 3 |
+
*
|
| 4 |
+
!python
|
| 5 |
+
!scripts/get_iv2.py
|
| 6 |
+
!deploy
|
| 7 |
+
deploy/demo/lake/bench/tables/*/_cache
|
| 8 |
+
!.dockerignore
|
| 9 |
+
python/**/__pycache__
|
.gitattributes
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.h264 filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
Dockerfile
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ElideDB query-by-example service. Build from the REPO ROOT:
|
| 2 |
+
# docker build -f deploy/Dockerfile.qbe -t elidedb-qbe .
|
| 3 |
+
# docker run -p 7860:7860 -v $PWD/lake:/app/lake elidedb-qbe
|
| 4 |
+
#
|
| 5 |
+
# Deliberately not the text image. That one carries torch and five text
|
| 6 |
+
# towers because a sentence has to be encoded before it can be matched.
|
| 7 |
+
# Query by example matches vectors the store already holds, so this image
|
| 8 |
+
# has no model in it at all - it boots as fast as the store downloads.
|
| 9 |
+
FROM python:3.11-slim
|
| 10 |
+
|
| 11 |
+
# ffmpeg muxes the clips the results play; every frame is decoded from the
|
| 12 |
+
# stored video's own bytes, so there is no image or preview cache anywhere.
|
| 13 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 14 |
+
ffmpeg && rm -rf /var/lib/apt/lists/*
|
| 15 |
+
|
| 16 |
+
WORKDIR /app
|
| 17 |
+
COPY deploy/requirements-qbe.txt /app/deploy/requirements-qbe.txt
|
| 18 |
+
RUN pip install --no-cache-dir -r deploy/requirements-qbe.txt
|
| 19 |
+
|
| 20 |
+
COPY python /app/python
|
| 21 |
+
COPY deploy/entrypoint-qbe.sh /app/deploy/entrypoint-qbe.sh
|
| 22 |
+
COPY deploy/qbe_serve.py /app/deploy/qbe_serve.py
|
| 23 |
+
COPY deploy/qbe_ui.html /app/deploy/qbe_ui.html
|
| 24 |
+
RUN chmod +x /app/deploy/entrypoint-qbe.sh
|
| 25 |
+
|
| 26 |
+
ENV PYTHONPATH=/app/python \
|
| 27 |
+
ELIDEDB_DEVICE=cpu \
|
| 28 |
+
DEMO_STORE_DATASET=SudharshanR/elidedb-demo-store \
|
| 29 |
+
ELIDEDB_LAKE=lake \
|
| 30 |
+
PORT=7860 \
|
| 31 |
+
HF_HOME=/app/hf_cache
|
| 32 |
+
|
| 33 |
+
EXPOSE 7860
|
| 34 |
+
CMD ["/app/deploy/entrypoint-qbe.sh"]
|
README.md
ADDED
|
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: ElideDB Query by Example
|
| 3 |
+
emoji: 🎬
|
| 4 |
+
colorFrom: green
|
| 5 |
+
colorTo: blue
|
| 6 |
+
sdk: docker
|
| 7 |
+
app_port: 7860
|
| 8 |
+
pinned: false
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
# ElideDB — query by example
|
| 12 |
+
|
| 13 |
+
1,122 robot episodes, and not one word of text. Pick a few clips that show
|
| 14 |
+
the same kind of moment; the archive returns its own kind.
|
| 15 |
+
|
| 16 |
+
**Why this exists.** A text encoder scores *"open the drawer"* against
|
| 17 |
+
*"close the drawer"* at cosine 0.977, and a pooled video embedding scores a
|
| 18 |
+
clip against its own reversal at 1.000000. Description destroys the
|
| 19 |
+
distinction before search begins. A clip needs no description: it is an
|
| 20 |
+
instance of what it shows.
|
| 21 |
+
|
| 22 |
+
**What runs when you click.** Nothing loads. Every vector was computed when
|
| 23 |
+
the video was written, so a query is arithmetic over parquet columns —
|
| 24 |
+
about 250 ms over the whole archive. Six channels vote, weighted by how
|
| 25 |
+
tightly each one pulls *your* picks together relative to how it holds the
|
| 26 |
+
archive: a statistic taken from the query itself, with no labels and no
|
| 27 |
+
training. The panel shows those weights, because which model recognised
|
| 28 |
+
your pick is the interesting part of the answer.
|
| 29 |
+
|
| 30 |
+
**How many come back.** The count is a ceiling, not a target. Each pick is
|
| 31 |
+
held out in turn to see how deep its own kind ranks, and the set stops
|
| 32 |
+
where that evidence stops — returned means believed.
|
| 33 |
+
|
| 34 |
+
Thumbnails and clips are decoded from the stored video's own byte ranges on
|
| 35 |
+
demand. There is no preview cache anywhere in the store.
|
| 36 |
+
|
| 37 |
+
The text-query demo of the same engine is at
|
| 38 |
+
[SudharshanR/elidedb-demo](https://huggingface.co/spaces/SudharshanR/elidedb-demo).
|
deploy/entrypoint-qbe.sh
ADDED
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/sh
|
| 2 |
+
# First start fetches the demo store (Space repos cap at 1 GB, so it ships
|
| 3 |
+
# as a public dataset). Nothing else is downloaded: this service loads no
|
| 4 |
+
# model, which is the point of it.
|
| 5 |
+
set -e
|
| 6 |
+
cd /app
|
| 7 |
+
if [ -n "$DEMO_STORE_DATASET" ] && [ ! -d /app/lake/bench ]; then
|
| 8 |
+
python - <<'PY'
|
| 9 |
+
import os
|
| 10 |
+
from huggingface_hub import snapshot_download
|
| 11 |
+
snapshot_download(os.environ["DEMO_STORE_DATASET"],
|
| 12 |
+
repo_type="dataset", local_dir="/app/lake")
|
| 13 |
+
print("demo store downloaded", flush=True)
|
| 14 |
+
PY
|
| 15 |
+
fi
|
| 16 |
+
exec python deploy/qbe_serve.py
|
deploy/qbe_serve.py
ADDED
|
@@ -0,0 +1,189 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The query-by-example Space: point at a moment, get its kind back.
|
| 2 |
+
|
| 3 |
+
WHY A SECOND DEPLOYMENT. The text Space answers "describe what you want"
|
| 4 |
+
and pays for it: five text towers, model weights on first boot, minutes
|
| 5 |
+
before the first query. This one answers the question a text encoder
|
| 6 |
+
provably cannot - direction - and costs no model at all. Everything it
|
| 7 |
+
ranks was computed at ingest; serving is numpy over parquet columns.
|
| 8 |
+
That is the whole argument for the storage layer, and it deserves to be
|
| 9 |
+
demonstrable in seconds rather than after a fifteen minute warm.
|
| 10 |
+
|
| 11 |
+
The read path is elidedb.qbe.search_like verbatim - multi-seed centroid,
|
| 12 |
+
per-channel seed-coherence weighting, z-fusion, and a return cut
|
| 13 |
+
calibrated from the query's own held-out seeds. Nothing here re-tunes it;
|
| 14 |
+
this file is a surface, not a second implementation.
|
| 15 |
+
|
| 16 |
+
Media (thumbnails, clips) reuses elidedb.desk's byte-range decode path,
|
| 17 |
+
which is the same one the text Space serves from - one media path, not a
|
| 18 |
+
copy that can drift.
|
| 19 |
+
|
| 20 |
+
python deploy/qbe_serve.py [--port 7860] [--root lake]
|
| 21 |
+
"""
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import argparse
|
| 25 |
+
import json
|
| 26 |
+
import os
|
| 27 |
+
import sys
|
| 28 |
+
import time
|
| 29 |
+
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
| 30 |
+
from pathlib import Path
|
| 31 |
+
from urllib.parse import parse_qs, urlparse
|
| 32 |
+
|
| 33 |
+
sys.path.insert(0, "python")
|
| 34 |
+
|
| 35 |
+
import numpy as np # noqa: E402
|
| 36 |
+
|
| 37 |
+
from elidedb import desk as D # noqa: E402 (media + store discovery)
|
| 38 |
+
from elidedb import qbe # noqa: E402
|
| 39 |
+
|
| 40 |
+
HERE = Path(__file__).resolve().parent
|
| 41 |
+
UI = HERE / "qbe_ui.html"
|
| 42 |
+
_EPISODES: dict = {}
|
| 43 |
+
_SPACES: dict = {}
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def episodes(key):
|
| 47 |
+
"""[(stream, ts, label)] for one store, ordered. The label is shown
|
| 48 |
+
AFTER retrieval so a visitor can check the answer; it is never read
|
| 49 |
+
by the ranker - search_like sees vectors and nothing else."""
|
| 50 |
+
if key in _EPISODES:
|
| 51 |
+
return _EPISODES[key]
|
| 52 |
+
db = D.STORES[key]
|
| 53 |
+
ep = db.table("episodes").scan()
|
| 54 |
+
cols = ep.column_names
|
| 55 |
+
lab = next((c for c in ("text", "task", "instruction", "label", "name")
|
| 56 |
+
if c in cols), None)
|
| 57 |
+
rows = list(zip(ep.column("stream").to_pylist(),
|
| 58 |
+
ep.column("ts").to_pylist(),
|
| 59 |
+
ep.column(lab).to_pylist() if lab
|
| 60 |
+
else [""] * ep.num_rows))
|
| 61 |
+
rows = [(str(s), int(t), (str(x) if x else "")) for s, t, x in rows]
|
| 62 |
+
rows.sort(key=lambda r: (r[0], r[1]))
|
| 63 |
+
_EPISODES[key] = rows
|
| 64 |
+
return rows
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def warm(key):
|
| 68 |
+
"""Build the channel matrices once, at boot.
|
| 69 |
+
|
| 70 |
+
spaces() reads every per-episode vector table and pools it. On the
|
| 71 |
+
demo store that is a couple of seconds and a few hundred megabytes
|
| 72 |
+
of parquet; paying it per query would make the first click look
|
| 73 |
+
broken for no reason, and paying it at import means the platform
|
| 74 |
+
reports RUNNING only when queries are genuinely warm - the same
|
| 75 |
+
lesson the text Space learned the expensive way.
|
| 76 |
+
"""
|
| 77 |
+
t0 = time.time()
|
| 78 |
+
keys, M = qbe.spaces(D.STORES[key])
|
| 79 |
+
_SPACES[key] = (keys, M)
|
| 80 |
+
print(f"warm {key}: {len(keys)} episodes, "
|
| 81 |
+
f"{len(M)} channels ({', '.join(sorted(M))}), "
|
| 82 |
+
f"{time.time() - t0:.1f}s", flush=True)
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def like(key, seeds, k_max):
|
| 86 |
+
db = D.STORES[key]
|
| 87 |
+
t0 = time.time()
|
| 88 |
+
out = qbe.search_like(db, [(s, int(t)) for s, t in seeds], k_max=k_max)
|
| 89 |
+
lab = {(s, t): x for s, t, x in episodes(key)}
|
| 90 |
+
return dict(
|
| 91 |
+
clips=[dict(stream=s, ts=t, label=lab.get((s, int(t)), ""))
|
| 92 |
+
for s, t in out["clips"]],
|
| 93 |
+
weights={c: round(float(w), 3)
|
| 94 |
+
for c, w in sorted(out.get("weights", {}).items(),
|
| 95 |
+
key=lambda kv: -kv[1])},
|
| 96 |
+
note=out.get("note", ""), ms=int((time.time() - t0) * 1000))
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
class Handler(BaseHTTPRequestHandler):
|
| 100 |
+
def log_message(self, *a):
|
| 101 |
+
pass
|
| 102 |
+
|
| 103 |
+
def _send(self, code, body, ctype, cache=False):
|
| 104 |
+
b = body if isinstance(body, bytes) else json.dumps(body).encode()
|
| 105 |
+
self.send_response(code)
|
| 106 |
+
self.send_header("Content-Type", ctype)
|
| 107 |
+
self.send_header("Content-Length", str(len(b)))
|
| 108 |
+
if cache:
|
| 109 |
+
self.send_header("Cache-Control", "public, max-age=86400")
|
| 110 |
+
self.end_headers()
|
| 111 |
+
self.wfile.write(b)
|
| 112 |
+
|
| 113 |
+
def _json(self, obj, code=200):
|
| 114 |
+
self._send(code, obj, "application/json")
|
| 115 |
+
|
| 116 |
+
def do_GET(self):
|
| 117 |
+
u = urlparse(self.path)
|
| 118 |
+
q = {k: v[0] for k, v in parse_qs(u.query).items()}
|
| 119 |
+
try:
|
| 120 |
+
if u.path in ("/", "/index.html"):
|
| 121 |
+
return self._send(200, UI.read_bytes(),
|
| 122 |
+
"text/html; charset=utf-8")
|
| 123 |
+
if u.path == "/api/archive":
|
| 124 |
+
key = q.get("store") or next(iter(D.STORES))
|
| 125 |
+
rows = episodes(key)
|
| 126 |
+
return self._json(dict(
|
| 127 |
+
store=key, stores=sorted(D.STORES),
|
| 128 |
+
channels=sorted(_SPACES.get(key, ([], {}))[1]),
|
| 129 |
+
episodes=[dict(stream=s, ts=t, label=x)
|
| 130 |
+
for s, t, x in rows]))
|
| 131 |
+
if u.path == "/api/thumb":
|
| 132 |
+
jpg = D.api_thumb(q.get("store") or next(iter(D.STORES)),
|
| 133 |
+
q.get("stream", ""), int(q["t"]),
|
| 134 |
+
int(q.get("w", "260")))
|
| 135 |
+
if jpg is None:
|
| 136 |
+
return self._json({"error": "no frame"}, 404)
|
| 137 |
+
return self._send(200, jpg, "image/jpeg", cache=True)
|
| 138 |
+
if u.path == "/api/clip":
|
| 139 |
+
mp4 = D.api_clip(q.get("store") or next(iter(D.STORES)),
|
| 140 |
+
q.get("stream", ""), int(q["t0"]),
|
| 141 |
+
int(q["t1"]), int(q.get("w", "480")))
|
| 142 |
+
if isinstance(mp4, dict):
|
| 143 |
+
return self._json(mp4, 422)
|
| 144 |
+
if mp4 is None:
|
| 145 |
+
return self._json({"error": "no frames"}, 404)
|
| 146 |
+
return self._send(200, mp4, "video/mp4", cache=True)
|
| 147 |
+
return self._json({"error": "not found"}, 404)
|
| 148 |
+
except Exception as e:
|
| 149 |
+
return self._json({"error": f"{type(e).__name__}: {e}"}, 500)
|
| 150 |
+
|
| 151 |
+
def do_POST(self):
|
| 152 |
+
u = urlparse(self.path)
|
| 153 |
+
try:
|
| 154 |
+
n = int(self.headers.get("Content-Length", 0))
|
| 155 |
+
body = json.loads(self.rfile.read(n) or b"{}")
|
| 156 |
+
if u.path == "/api/like":
|
| 157 |
+
seeds = body.get("seeds", [])
|
| 158 |
+
if not seeds:
|
| 159 |
+
return self._json({"error": "pick at least one clip"},
|
| 160 |
+
400)
|
| 161 |
+
key = body.get("store") or next(iter(D.STORES))
|
| 162 |
+
return self._json(like(key, seeds,
|
| 163 |
+
int(body.get("k", 24))))
|
| 164 |
+
return self._json({"error": "not found"}, 404)
|
| 165 |
+
except Exception as e:
|
| 166 |
+
return self._json({"error": f"{type(e).__name__}: {e}"}, 500)
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
def main():
|
| 170 |
+
ap = argparse.ArgumentParser()
|
| 171 |
+
ap.add_argument("--port", type=int,
|
| 172 |
+
default=int(os.environ.get("PORT", "7860")))
|
| 173 |
+
ap.add_argument("--host", default="0.0.0.0")
|
| 174 |
+
ap.add_argument("--root", default=os.environ.get("ELIDEDB_LAKE", "lake"))
|
| 175 |
+
a = ap.parse_args()
|
| 176 |
+
D.ROOT = Path.cwd()
|
| 177 |
+
D.LAKE = Path(a.root)
|
| 178 |
+
D.discover()
|
| 179 |
+
if not D.STORES:
|
| 180 |
+
raise SystemExit(f"no store under {a.root}/ - set ELIDEDB_LAKE or "
|
| 181 |
+
f"download the demo store first")
|
| 182 |
+
for key in D.STORES:
|
| 183 |
+
warm(key)
|
| 184 |
+
print(f"query by example on http://{a.host}:{a.port}", flush=True)
|
| 185 |
+
ThreadingHTTPServer((a.host, a.port), Handler).serve_forever()
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
if __name__ == "__main__":
|
| 189 |
+
main()
|
deploy/qbe_ui.html
ADDED
|
@@ -0,0 +1,198 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!doctype html>
|
| 2 |
+
<meta charset="utf-8">
|
| 3 |
+
<meta name="viewport" content="width=device-width, initial-scale=1">
|
| 4 |
+
<title>ElideDB · Query by Example</title>
|
| 5 |
+
<style>
|
| 6 |
+
:root{
|
| 7 |
+
--bg:#0b0e10; --panel:#131719; --line:#232a2e; --ink:#e2e8ea;
|
| 8 |
+
--dim:#8a969c; --accent:#6ee7a8; --accent-dim:#1d3a2e; --warn:#e3b341;
|
| 9 |
+
--mono:ui-monospace,SFMono-Regular,"SF Mono",Menlo,monospace;
|
| 10 |
+
}
|
| 11 |
+
*{box-sizing:border-box}
|
| 12 |
+
body{margin:0;background:var(--bg);color:var(--ink);
|
| 13 |
+
font:14px/1.55 var(--mono)}
|
| 14 |
+
header{padding:16px 22px;border-bottom:1px solid var(--line)}
|
| 15 |
+
h1{margin:0;font-size:15px;font-weight:600;letter-spacing:.01em}
|
| 16 |
+
header p{margin:5px 0 0;color:var(--dim);font-size:12px;max-width:76ch}
|
| 17 |
+
header b{color:var(--ink);font-weight:600}
|
| 18 |
+
main{padding:16px 22px 40px}
|
| 19 |
+
.bar{position:sticky;top:0;z-index:5;background:var(--bg);
|
| 20 |
+
border-bottom:1px solid var(--line);padding:11px 0;margin-bottom:14px;
|
| 21 |
+
display:flex;gap:12px;align-items:center;flex-wrap:wrap}
|
| 22 |
+
.bar .count{color:var(--dim);font-size:12px}
|
| 23 |
+
.bar .count b{color:var(--accent)}
|
| 24 |
+
button{background:var(--accent-dim);border:1px solid #2f5c48;
|
| 25 |
+
color:var(--accent);padding:8px 16px;border-radius:6px;
|
| 26 |
+
font:13px var(--mono);cursor:pointer}
|
| 27 |
+
button:disabled{opacity:.35;cursor:default}
|
| 28 |
+
button.ghost{background:transparent;border-color:var(--line);
|
| 29 |
+
color:var(--dim)}
|
| 30 |
+
h2{font-size:11px;letter-spacing:.08em;text-transform:uppercase;
|
| 31 |
+
color:var(--dim);margin:22px 0 9px;font-weight:600}
|
| 32 |
+
.grid{display:grid;grid-template-columns:repeat(auto-fill,minmax(150px,1fr));
|
| 33 |
+
gap:9px}
|
| 34 |
+
.c{position:relative;border:1px solid var(--line);border-radius:6px;
|
| 35 |
+
overflow:hidden;background:var(--panel);cursor:pointer;aspect-ratio:4/3}
|
| 36 |
+
.c img,.c video{width:100%;height:100%;object-fit:cover;display:block}
|
| 37 |
+
.c .ph{width:100%;height:100%;background:#0f1315}
|
| 38 |
+
.c.on{border-color:var(--accent);box-shadow:0 0 0 1px var(--accent)}
|
| 39 |
+
.c .t{position:absolute;left:0;bottom:0;right:0;padding:3px 6px;
|
| 40 |
+
background:linear-gradient(transparent,rgba(0,0,0,.82));
|
| 41 |
+
font-size:10px;color:#c7d0d4}
|
| 42 |
+
.c .n{position:absolute;top:5px;left:6px;font-size:10px;color:var(--accent);
|
| 43 |
+
background:rgba(0,0,0,.6);padding:1px 5px;border-radius:3px}
|
| 44 |
+
#results .c{cursor:default}
|
| 45 |
+
.why{border:1px solid var(--line);border-radius:8px;background:var(--panel);
|
| 46 |
+
padding:12px 14px;margin:10px 0 4px}
|
| 47 |
+
.why .h{font-size:11px;letter-spacing:.06em;text-transform:uppercase;
|
| 48 |
+
color:var(--dim);margin-bottom:9px}
|
| 49 |
+
.w{display:flex;align-items:center;gap:9px;margin:4px 0;font-size:12px}
|
| 50 |
+
.w .nm{width:62px;color:var(--ink)}
|
| 51 |
+
.w .track{flex:1;height:7px;background:#0f1315;border-radius:4px;
|
| 52 |
+
overflow:hidden}
|
| 53 |
+
.w .fill{height:100%;background:var(--accent)}
|
| 54 |
+
.w .v{width:44px;text-align:right;color:var(--dim)}
|
| 55 |
+
.note{color:var(--dim);font-size:11.5px;margin-top:9px;line-height:1.5}
|
| 56 |
+
.note b{color:var(--ink);font-weight:600}
|
| 57 |
+
.err{color:var(--warn)}
|
| 58 |
+
.tray{display:flex;gap:6px;flex-wrap:wrap;margin-bottom:10px}
|
| 59 |
+
.tray .chip{width:62px;height:47px;object-fit:cover;border-radius:4px;
|
| 60 |
+
border:1px solid var(--accent);cursor:pointer}
|
| 61 |
+
.pager{display:flex;gap:10px;align-items:center;margin:6px 0 10px}
|
| 62 |
+
</style>
|
| 63 |
+
<header>
|
| 64 |
+
<h1>ElideDB · query by example</h1>
|
| 65 |
+
<p><b>1,122 robot episodes, and not one word of text.</b> Pick clips that
|
| 66 |
+
show the same kind of moment; the archive returns its own kind. Nothing is
|
| 67 |
+
described, nothing is labelled, and no model runs to answer you — every
|
| 68 |
+
vector was computed when the video was written, so a query is arithmetic
|
| 69 |
+
over columns.</p>
|
| 70 |
+
</header>
|
| 71 |
+
<main>
|
| 72 |
+
<div class="bar">
|
| 73 |
+
<span class="count"><b id="nsel">0</b> picked <span id="hint">·
|
| 74 |
+
8–10 of the same kind works best</span></span>
|
| 75 |
+
<button id="go" disabled>find more like these</button>
|
| 76 |
+
<button id="clear" class="ghost">clear</button>
|
| 77 |
+
<span class="count" id="stat"></span>
|
| 78 |
+
</div>
|
| 79 |
+
<div id="seltray" class="tray"></div>
|
| 80 |
+
<div id="out"></div>
|
| 81 |
+
<h2 id="archhead">the archive</h2>
|
| 82 |
+
<div class="pager" id="pager"></div>
|
| 83 |
+
<div class="grid" id="archive"></div>
|
| 84 |
+
<div class="pager" id="pager2"></div>
|
| 85 |
+
</main>
|
| 86 |
+
<script>
|
| 87 |
+
const $=s=>document.querySelector(s);
|
| 88 |
+
const esc=t=>String(t).replace(/[&<>"']/g,c=>({'&':'&','<':'<','>':'>','"':'"',"'":'''}[c]));
|
| 89 |
+
let EPS=[], STORE='', sel=[];
|
| 90 |
+
|
| 91 |
+
const thumb=(e,w)=>`/api/thumb?store=${encodeURIComponent(STORE)}&stream=${encodeURIComponent(e.stream)}&t=${e.ts}&w=${w||220}`;
|
| 92 |
+
const clip =e=>`/api/clip?store=${encodeURIComponent(STORE)}&stream=${encodeURIComponent(e.stream)}&t0=${e.ts}&t1=${e.ts+6000000000}&w=420`;
|
| 93 |
+
const shortName=e=>e.stream.split('/').pop()+' · '+new Date(e.ts/1e6).toISOString().slice(11,19);
|
| 94 |
+
|
| 95 |
+
// Thumbnails are decoded on demand from the video's own bytes - there is no
|
| 96 |
+
// pre-baked image anywhere in the store. A page at a time is what makes that
|
| 97 |
+
// affordable: 120 decodes at ~50 ms, not 1,122. They load eagerly on purpose.
|
| 98 |
+
// Both deferred approaches - IntersectionObserver and loading="lazy" - left
|
| 99 |
+
// the grid blank in verification, and an archive you cannot see is not an
|
| 100 |
+
// archive you can pick from. A bounded page costs less than a clever one.
|
| 101 |
+
const PAGE=120; let page=0;
|
| 102 |
+
|
| 103 |
+
function card(e,i){
|
| 104 |
+
return `<div class="c${sel.includes(i)?' on':''}" data-i="${i}">
|
| 105 |
+
<img decoding="async" src="${thumb(e)}" alt="">
|
| 106 |
+
<div class="t">${esc(shortName(e))}</div></div>`;
|
| 107 |
+
}
|
| 108 |
+
|
| 109 |
+
function paintArchive(){
|
| 110 |
+
const n=Math.ceil(EPS.length/PAGE);
|
| 111 |
+
page=Math.max(0,Math.min(page,n-1));
|
| 112 |
+
const a=page*PAGE, b=Math.min(a+PAGE,EPS.length);
|
| 113 |
+
$('#archive').innerHTML=EPS.slice(a,b).map((e,k)=>card(e,a+k)).join('');
|
| 114 |
+
$('#pager').innerHTML=`<button class="ghost" id="prev" ${page?'':'disabled'}>←</button>
|
| 115 |
+
<span class="count">${a+1}–${b} of ${EPS.length}</span>
|
| 116 |
+
<button class="ghost" id="next" ${page<n-1?'':'disabled'}>→</button>`;
|
| 117 |
+
const pv=$('#prev'), nx=$('#next');
|
| 118 |
+
if(pv)pv.onclick=()=>{page--;paintArchive();};
|
| 119 |
+
if(nx)nx.onclick=()=>{page++;paintArchive();};
|
| 120 |
+
}
|
| 121 |
+
|
| 122 |
+
function refreshSel(){
|
| 123 |
+
$('#nsel').textContent=sel.length;
|
| 124 |
+
$('#go').disabled=sel.length===0;
|
| 125 |
+
$('#hint').textContent=sel.length>=8?'· good sample size':(sel.length?'· add a few more of the same kind':'· 8–10 of the same kind works best');
|
| 126 |
+
[...$('#archive').children].forEach(d=>{
|
| 127 |
+
d.classList.toggle('on', sel.includes(+d.dataset.i));
|
| 128 |
+
});
|
| 129 |
+
const t=$('#seltray');
|
| 130 |
+
t.innerHTML=sel.length?sel.map(i=>
|
| 131 |
+
`<img class="chip" src="${thumb(EPS[i],120)}" data-i="${i}" title="click to drop">`).join(''):'';
|
| 132 |
+
}
|
| 133 |
+
|
| 134 |
+
$('#archive').addEventListener('click',ev=>{
|
| 135 |
+
const c=ev.target.closest('.c'); if(!c)return;
|
| 136 |
+
const i=+c.dataset.i, k=sel.indexOf(i);
|
| 137 |
+
if(k>=0) sel.splice(k,1); else sel.push(i);
|
| 138 |
+
refreshSel();
|
| 139 |
+
});
|
| 140 |
+
$('#seltray').addEventListener('click',ev=>{
|
| 141 |
+
if(ev.target.dataset.i===undefined)return;
|
| 142 |
+
const k=sel.indexOf(+ev.target.dataset.i);
|
| 143 |
+
if(k>=0){sel.splice(k,1);refreshSel();}
|
| 144 |
+
});
|
| 145 |
+
$('#clear').onclick=()=>{sel=[];refreshSel();$('#out').innerHTML='';$('#stat').textContent='';};
|
| 146 |
+
|
| 147 |
+
$('#go').onclick=async()=>{
|
| 148 |
+
$('#go').disabled=true; $('#stat').textContent='ranking…';
|
| 149 |
+
const body={store:STORE,k:24,seeds:sel.map(i=>[EPS[i].stream,EPS[i].ts])};
|
| 150 |
+
let d;
|
| 151 |
+
try{ d=await (await fetch('/api/like',{method:'POST',body:JSON.stringify(body)})).json(); }
|
| 152 |
+
catch(e){ d={error:String(e)}; }
|
| 153 |
+
$('#go').disabled=false;
|
| 154 |
+
if(d.error){ $('#stat').innerHTML=`<span class="err">${esc(d.error)}</span>`; return; }
|
| 155 |
+
$('#stat').textContent=`${d.ms} ms`;
|
| 156 |
+
const ws=Object.entries(d.weights||{});
|
| 157 |
+
const mx=Math.max(...ws.map(x=>x[1]),1e-9);
|
| 158 |
+
$('#out').innerHTML=`
|
| 159 |
+
<h2>what came back — ${d.clips.length} moments</h2>
|
| 160 |
+
<div class="grid" id="results">${d.clips.map((e,i)=>
|
| 161 |
+
`<div class="c" data-r="${i}">
|
| 162 |
+
<img decoding="async" src="${thumb(e)}" alt="">
|
| 163 |
+
<span class="n">${i+1}</span>
|
| 164 |
+
<div class="t">${esc(shortName(e))}</div></div>`).join('')}
|
| 165 |
+
</div>
|
| 166 |
+
<div class="why">
|
| 167 |
+
<div class="h">which model recognised your pick</div>
|
| 168 |
+
${ws.map(([c,w])=>`<div class="w"><span class="nm">${esc(c)}</span>
|
| 169 |
+
<span class="track"><span class="fill" style="width:${(100*w/mx).toFixed(1)}%"></span></span>
|
| 170 |
+
<span class="v">${w.toFixed(2)}</span></div>`).join('')}
|
| 171 |
+
<div class="note">Every channel votes, but not equally. The weight is
|
| 172 |
+
how tightly that model pulls <b>your</b> picks together relative to how
|
| 173 |
+
it holds the whole archive — computed from the query itself, with no
|
| 174 |
+
labels and no training. ${d.note?`<br><b>${esc(d.note)}</b>`:''}
|
| 175 |
+
<br>The number returned is a <b>ceiling, not a target</b>: each pick is
|
| 176 |
+
held out in turn to see how deep its own kind ranks, and the set stops
|
| 177 |
+
where that evidence stops. Returned means believed.</div>
|
| 178 |
+
</div>`;
|
| 179 |
+
const g=$('#results');
|
| 180 |
+
[...g.children].forEach((div,i)=>{
|
| 181 |
+
const e=d.clips[i];
|
| 182 |
+
div.onclick=()=>{
|
| 183 |
+
if(div.querySelector('video'))return;
|
| 184 |
+
const v=document.createElement('video');
|
| 185 |
+
v.src=clip(e); v.controls=true; v.autoplay=true; v.loop=true; v.muted=true;
|
| 186 |
+
div.querySelector('img').replaceWith(v);
|
| 187 |
+
};
|
| 188 |
+
});
|
| 189 |
+
window.scrollTo({top:0,behavior:'smooth'});
|
| 190 |
+
};
|
| 191 |
+
|
| 192 |
+
(async()=>{
|
| 193 |
+
const d=await (await fetch('/api/archive')).json();
|
| 194 |
+
EPS=d.episodes; STORE=d.store;
|
| 195 |
+
$('#archhead').textContent=`the archive — ${EPS.length} episodes, ${d.channels.length} channels (${d.channels.join(', ')})`;
|
| 196 |
+
paintArchive();
|
| 197 |
+
})();
|
| 198 |
+
</script>
|
deploy/requirements-qbe.txt
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# The query-by-example service. NO MODEL RUNS AT QUERY TIME: every vector
|
| 2 |
+
# was computed at ingest, so serving is numpy over parquet columns and the
|
| 3 |
+
# torch stack that the text service needs is simply absent. That is the
|
| 4 |
+
# difference between a fifteen minute first boot and a fifteen second one,
|
| 5 |
+
# and it is the storage argument made concrete rather than asserted.
|
| 6 |
+
numpy<2.1
|
| 7 |
+
pyarrow>=17
|
| 8 |
+
pillow>=10
|
| 9 |
+
huggingface_hub>=0.26
|
| 10 |
+
# the frame decoder behind every thumbnail and clip: frames are decoded
|
| 11 |
+
# from the stored video's own byte ranges on demand, never pre-baked, so
|
| 12 |
+
# the decoder is not optional even though no model is.
|
| 13 |
+
opencv-python-headless>=4.10
|
python/.DS_Store
ADDED
|
Binary file (6.15 kB). View file
|
|
|
python/elidedb.egg-info/PKG-INFO
ADDED
|
@@ -0,0 +1,123 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Metadata-Version: 2.4
|
| 2 |
+
Name: elidedb
|
| 3 |
+
Version: 2.0.0
|
| 4 |
+
Summary: Parquet-native, timestamp-first multimodal database — the best read is the read elided
|
| 5 |
+
Requires-Python: >=3.10
|
| 6 |
+
Description-Content-Type: text/markdown
|
| 7 |
+
Requires-Dist: pyarrow>=15
|
| 8 |
+
Requires-Dist: numpy
|
| 9 |
+
Requires-Dist: pandas
|
| 10 |
+
Requires-Dist: duckdb
|
| 11 |
+
Requires-Dist: pillow
|
| 12 |
+
Requires-Dist: opencv-python
|
| 13 |
+
Requires-Dist: mlx-embeddings
|
| 14 |
+
Requires-Dist: scikit-learn
|
| 15 |
+
Requires-Dist: hdbscan
|
| 16 |
+
Requires-Dist: umap-learn
|
| 17 |
+
Requires-Dist: matplotlib
|
| 18 |
+
Requires-Dist: jupyter
|
| 19 |
+
Provides-Extra: ml
|
| 20 |
+
Requires-Dist: mlx-embeddings; extra == "ml"
|
| 21 |
+
Requires-Dist: scikit-learn; extra == "ml"
|
| 22 |
+
Requires-Dist: hdbscan; extra == "ml"
|
| 23 |
+
Requires-Dist: umap-learn; extra == "ml"
|
| 24 |
+
|
| 25 |
+
# ElideDB
|
| 26 |
+
|
| 27 |
+
**A Parquet-native, timestamp-first database for multimodal data.** Put any
|
| 28 |
+
timestamped data in — sensor rows, video, GPS, audio, logs — and get back
|
| 29 |
+
time-window reads, SQL, and semantic search, while the engine reads as few
|
| 30 |
+
bytes as physically possible. *The best read is the read elided.*
|
| 31 |
+
|
| 32 |
+
- **Everything is Parquet.** Every table is plain Parquet files under a
|
| 33 |
+
Delta-Lake-style transaction log. No custom formats; DuckDB, Spark, pandas —
|
| 34 |
+
anything that reads Parquet reads your database.
|
| 35 |
+
- **Timestamps are the law.** Every table has a `ts` column (int64
|
| 36 |
+
nanoseconds). That one rule is what makes cross-modal queries, alignment,
|
| 37 |
+
and pruning work.
|
| 38 |
+
- **Media is indexed, never copied.** Video files stay exactly where they
|
| 39 |
+
are; ElideDB stores byte ranges into them and decodes only the frames a
|
| 40 |
+
query touches.
|
| 41 |
+
- **Search by meaning.** Local SigLIP embeddings (no cloud) turn video into
|
| 42 |
+
searchable windows: `"a person crossing the street"` → playable clips.
|
| 43 |
+
|
| 44 |
+
## Install
|
| 45 |
+
|
| 46 |
+
```bash
|
| 47 |
+
pip install -e ".[ml]" # from the repo root; [ml] adds semantic search
|
| 48 |
+
brew install ffmpeg # clip playback + video indexing
|
| 49 |
+
```
|
| 50 |
+
|
| 51 |
+
## 60 seconds to your first database
|
| 52 |
+
|
| 53 |
+
```bash
|
| 54 |
+
elidedb create lake/mydb --name "my project"
|
| 55 |
+
|
| 56 |
+
# any timestamped rows: CSV/Parquet, ISO dates or epoch s/ms/us/ns (auto-detected)
|
| 57 |
+
elidedb add lake/mydb readings sensor_log.csv --ts-col time
|
| 58 |
+
|
| 59 |
+
# any video (the file is indexed in place, not copied)
|
| 60 |
+
elidedb video lake/mydb dashcam.mp4 --stream front
|
| 61 |
+
|
| 62 |
+
elidedb ls lake/mydb # what's inside
|
| 63 |
+
elidedb sql lake/mydb "SELECT count(*) FROM readings" # SQL via DuckDB
|
| 64 |
+
elidedb embed lake/mydb # local ML, one command
|
| 65 |
+
elidedb search lake/mydb "a cyclist passing a bus" # → time windows
|
| 66 |
+
elidedb desk # browse it
|
| 67 |
+
```
|
| 68 |
+
|
| 69 |
+
The same five verbs in Python:
|
| 70 |
+
|
| 71 |
+
```python
|
| 72 |
+
from elidedb import Store
|
| 73 |
+
db = Store.create("lake/mydb", "my project")
|
| 74 |
+
db.ingest_rows("readings", df, ts_column="time") # DataFrame/CSV/Parquet
|
| 75 |
+
db.ingest_video("frames", "dashcam.mp4", stream="front")
|
| 76 |
+
db.embed_windows(); import elidedb; elidedb.cluster(db)
|
| 77 |
+
hits, _ = db.search_text("a cyclist passing a bus") # → (stream, t0, t1)
|
| 78 |
+
window, stats = db.window(hits[0]["t0"], hits[0]["t1"])
|
| 79 |
+
print(stats) # bytes touched vs corpus
|
| 80 |
+
```
|
| 81 |
+
|
| 82 |
+
## Documentation
|
| 83 |
+
|
| 84 |
+
| doc | what it covers |
|
| 85 |
+
|---|---|
|
| 86 |
+
| [docs/GETTING_STARTED.md](docs/GETTING_STARTED.md) | step-by-step: install → create → add your data → query → browse |
|
| 87 |
+
| [docs/API.md](docs/API.md) | every class, method, and CLI verb |
|
| 88 |
+
| [DESIGN.md](DESIGN.md) | architecture + which idea came from which system (Delta, Spark, C-Store, warehouses) |
|
| 89 |
+
| [notebooks/elidedb_demo.ipynb](notebooks/elidedb_demo.ipynb) | every query style, executed on real data with outputs baked in |
|
| 90 |
+
| [BENCHMARKS.md](BENCHMARKS.md) | measured numbers on 28.5 GB of real captures |
|
| 91 |
+
|
| 92 |
+
## ElideDB Desk
|
| 93 |
+
|
| 94 |
+
`elidedb desk` (or double-click `desk/ElideDB Desk.app` on macOS) opens the
|
| 95 |
+
database browser: every store's tables and timelines, a semantic map where
|
| 96 |
+
hovering any point decodes its frame live, click-to-play clips (with the
|
| 97 |
+
sensor's audio track when the store has one), and a query console for
|
| 98 |
+
SQL / semantic / window queries.
|
| 99 |
+
|
| 100 |
+
## The numbers that matter (measured, [BENCHMARKS.md](BENCHMARKS.md))
|
| 101 |
+
|
| 102 |
+
- 2 s window over a **14.16 M-row** audio table: touches **4 MB of 286 MB
|
| 103 |
+
(98.6 % elided)**, 165 ms.
|
| 104 |
+
- Multimodal 2 s window across 15 tables: 10/56 files touched, **98.9 %
|
| 105 |
+
elided, 13.8 ms** (sensor-only).
|
| 106 |
+
- Semantic search: ranking is **microseconds** at thousands of windows;
|
| 107 |
+
end-to-end text query ≈ 2.8 s (model load dominates, then it stays warm).
|
| 108 |
+
- Store is *smaller* than the raw input in both real corpora, while adding
|
| 109 |
+
random access, SQL, and search.
|
| 110 |
+
|
| 111 |
+
## Repository layout
|
| 112 |
+
|
| 113 |
+
```
|
| 114 |
+
python/elidedb/ the database (store, log, video, embeddings, cli, desk)
|
| 115 |
+
notebooks/ executed demo notebook
|
| 116 |
+
desk/ macOS app bundle (thin launcher for elidedb.desk)
|
| 117 |
+
src/, tests/ v1: the original C++20 engine with hand-built formats
|
| 118 |
+
(SDX/SFI) — the mechanisms ElideDB now hosts on Parquet
|
| 119 |
+
scripts/ dataset ETL adapters (REIP, Oxford RobotCar) + tooling
|
| 120 |
+
```
|
| 121 |
+
|
| 122 |
+
Raw data (`data/`), generated databases (`lake/`, `store*/`), and build
|
| 123 |
+
output are git-ignored — the repo carries code and docs only.
|
python/elidedb.egg-info/SOURCES.txt
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
README.md
|
| 2 |
+
pyproject.toml
|
| 3 |
+
python/elidedb/__init__.py
|
| 4 |
+
python/elidedb/ann.py
|
| 5 |
+
python/elidedb/bptree.py
|
| 6 |
+
python/elidedb/cli.py
|
| 7 |
+
python/elidedb/desk.py
|
| 8 |
+
python/elidedb/desk_ui.html
|
| 9 |
+
python/elidedb/embeddings.py
|
| 10 |
+
python/elidedb/fftools.py
|
| 11 |
+
python/elidedb/log.py
|
| 12 |
+
python/elidedb/migrate.py
|
| 13 |
+
python/elidedb/store.py
|
| 14 |
+
python/elidedb/video.py
|
| 15 |
+
python/elidedb.egg-info/PKG-INFO
|
| 16 |
+
python/elidedb.egg-info/SOURCES.txt
|
| 17 |
+
python/elidedb.egg-info/dependency_links.txt
|
| 18 |
+
python/elidedb.egg-info/entry_points.txt
|
| 19 |
+
python/elidedb.egg-info/requires.txt
|
| 20 |
+
python/elidedb.egg-info/top_level.txt
|
python/elidedb.egg-info/dependency_links.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
|
python/elidedb.egg-info/entry_points.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[console_scripts]
|
| 2 |
+
elidedb = elidedb.cli:main
|
python/elidedb.egg-info/requires.txt
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
pyarrow>=15
|
| 2 |
+
numpy
|
| 3 |
+
pandas
|
| 4 |
+
duckdb
|
| 5 |
+
pillow
|
| 6 |
+
opencv-python
|
| 7 |
+
mlx-embeddings
|
| 8 |
+
scikit-learn
|
| 9 |
+
hdbscan
|
| 10 |
+
umap-learn
|
| 11 |
+
matplotlib
|
| 12 |
+
jupyter
|
| 13 |
+
|
| 14 |
+
[ml]
|
| 15 |
+
mlx-embeddings
|
| 16 |
+
scikit-learn
|
| 17 |
+
hdbscan
|
| 18 |
+
umap-learn
|
python/elidedb.egg-info/top_level.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
elidedb
|
python/elidedb/__init__.py
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ElideDB — a Parquet-only, timestamp-first multimodal store.
|
| 2 |
+
|
| 3 |
+
>>> from elidedb import Store
|
| 4 |
+
>>> db = Store.open("lake/oxford")
|
| 5 |
+
>>> db.describe()
|
| 6 |
+
>>> data, stats = db.window(t0, t1)
|
| 7 |
+
>>> db.sql("SELECT count(*) FROM gps")
|
| 8 |
+
>>> db.search_text("pedestrians crossing")
|
| 9 |
+
"""
|
| 10 |
+
from .store import Store, Table, QueryStats
|
| 11 |
+
from .embeddings import cluster, embed_text
|
| 12 |
+
|
| 13 |
+
__version__ = "2.0.0"
|
| 14 |
+
__all__ = ["Store", "Table", "QueryStats", "cluster", "embed_text"]
|
python/elidedb/action_channel.py
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The `act` channel: SSv2 action posteriors as an INDEX-ONLY ranking
|
| 2 |
+
signal. Ingested once per episode (scripts/action_ingest.py: V-JEPA 2
|
| 3 |
+
ViT-L + Meta's released attentive probe, 174 classes); a query costs a
|
| 4 |
+
cached text-vector pass over the 174 class names plus one 174-d dot per
|
| 5 |
+
episode.
|
| 6 |
+
|
| 7 |
+
Adoption measurement (2026-07-24, labeled episodes, zero fitting):
|
| 8 |
+
put-in vs take-out AUC 0.889 with the literal class pair — the exact
|
| 9 |
+
containment direction the green-drawer failure exposed and no other
|
| 10 |
+
channel measures. Directional queries score a CONTRAST of class
|
| 11 |
+
weights, w(text) − w(swap), mirroring the swap-contrast law used
|
| 12 |
+
everywhere else in this system.
|
| 13 |
+
"""
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
|
| 18 |
+
_IDX = {}
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def _index(store):
|
| 22 |
+
ver = store.table("action_probs").state().version
|
| 23 |
+
key = (str(store.dir), ver)
|
| 24 |
+
if key not in _IDX:
|
| 25 |
+
from .embeddings import _vec_table
|
| 26 |
+
tbl, _ = _vec_table(store, "action_probs")
|
| 27 |
+
ss = tbl.column("stream").to_pylist()
|
| 28 |
+
sa = [int(v) for v in tbl.column("ts").to_pylist()]
|
| 29 |
+
sb = [int(v) for v in tbl.column("t1").to_pylist()]
|
| 30 |
+
idx = {}
|
| 31 |
+
for r, (s, a, b) in enumerate(zip(ss, sa, sb)):
|
| 32 |
+
idx.setdefault(str(s), []).append((a, b, r))
|
| 33 |
+
for s in idx:
|
| 34 |
+
idx[s].sort()
|
| 35 |
+
if len(_IDX) > 8:
|
| 36 |
+
_IDX.clear()
|
| 37 |
+
_IDX[key] = idx
|
| 38 |
+
return _IDX[key]
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def act_lookup(store, text, contrast=None):
|
| 42 |
+
"""(lookup(stream, t0, t1) -> weighted posterior | nan,
|
| 43 |
+
candidates top-64). Rows in action_probs are one per episode.
|
| 44 |
+
`contrast`: an explicit 174-d weight vector (canonical_contrast)
|
| 45 |
+
overrides the text-mapped weights."""
|
| 46 |
+
from .action_probe import query_class_weights
|
| 47 |
+
from .embeddings import _vec_table
|
| 48 |
+
from .rerank import directional_swap
|
| 49 |
+
|
| 50 |
+
idx = _index(store)
|
| 51 |
+
_, probs = _vec_table(store, "action_probs")
|
| 52 |
+
|
| 53 |
+
if contrast is not None:
|
| 54 |
+
w = contrast
|
| 55 |
+
else:
|
| 56 |
+
w = query_class_weights(text)
|
| 57 |
+
sq = directional_swap(text, store)
|
| 58 |
+
if sq is not None:
|
| 59 |
+
w = w - query_class_weights(sq)
|
| 60 |
+
sc = np.asarray(probs) @ w
|
| 61 |
+
|
| 62 |
+
def lookup(s, a, b):
|
| 63 |
+
lst = idx.get(str(s))
|
| 64 |
+
if not lst:
|
| 65 |
+
return float("nan")
|
| 66 |
+
starts = [x[0] for x in lst]
|
| 67 |
+
j = int(np.searchsorted(starts, a, side="right")) - 1
|
| 68 |
+
if j >= 0 and b <= lst[j][1] + 1:
|
| 69 |
+
return float(sc[lst[j][2]])
|
| 70 |
+
return float("nan")
|
| 71 |
+
|
| 72 |
+
cands = []
|
| 73 |
+
for s, lst in idx.items():
|
| 74 |
+
for a, b, r in lst:
|
| 75 |
+
cands.append((s, a, b, float(sc[r])))
|
| 76 |
+
cands.sort(key=lambda x: -x[3])
|
| 77 |
+
return lookup, cands[:64]
|
python/elidedb/action_probe.py
ADDED
|
@@ -0,0 +1,190 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ACTION PROBE — video-native verb grounding, zero LLM/VLM anywhere.
|
| 2 |
+
|
| 3 |
+
The calibration that forced this: both VLM judge tiers overclaim scene
|
| 4 |
+
matches (7B est 1.0 on visually ~10%-pure sets) — a language model shown
|
| 5 |
+
frames cannot bind actions. User directive: no LLM/VLM judges, ever.
|
| 6 |
+
|
| 7 |
+
Replacement: Meta's released SSv2 attentive probe on the V-JEPA 2 ViT-L
|
| 8 |
+
encoder ALREADY in the stack (291 ms/16-frame clip on MPS, measured).
|
| 9 |
+
Something-Something-v2's 174 classes are literally the failing query
|
| 10 |
+
taxonomy of this database — "Closing something", "Opening something",
|
| 11 |
+
"Picking something up", "Putting something into something", "Taking
|
| 12 |
+
something out of something", "Covering something with something" (= lid
|
| 13 |
+
on vessel), "Folding something". The probe is 174-way evidence about
|
| 14 |
+
WHAT HAPPENED in the clip, produced by a video model trained on motion,
|
| 15 |
+
not a text model guessing from two frames.
|
| 16 |
+
|
| 17 |
+
Faithful port of vjepa2/src/models/attentive_pooler.py (MIT): 3
|
| 18 |
+
self-attn blocks over ALL encoder tokens -> 1-query cross-attn pool ->
|
| 19 |
+
linear(1024 -> 174). Class list vendored as data/ssv2_classes.txt —
|
| 20 |
+
this is the MODEL's output vocabulary (like a tokenizer), not dataset
|
| 21 |
+
metadata; the no-metadata rule stays intact.
|
| 22 |
+
|
| 23 |
+
Query -> class mapping rides the PE text tower already cached: cosine
|
| 24 |
+
between the query and the 174 class names, softmaxed over the top few.
|
| 25 |
+
"""
|
| 26 |
+
from __future__ import annotations
|
| 27 |
+
|
| 28 |
+
from pathlib import Path
|
| 29 |
+
|
| 30 |
+
import numpy as np
|
| 31 |
+
|
| 32 |
+
_STATE = {}
|
| 33 |
+
|
| 34 |
+
PROBE_CKPT = Path(__file__).resolve().parents[2] / \
|
| 35 |
+
"models/ssv2-vitl-16x2x3.pt"
|
| 36 |
+
ENCODER_ID = "facebook/vjepa2-vitl-fpc64-256"
|
| 37 |
+
N_CLASSES = 174
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def ssv2_classes():
|
| 41 |
+
if "classes" not in _STATE:
|
| 42 |
+
p = Path(__file__).parent / "data/ssv2_classes.txt"
|
| 43 |
+
_STATE["classes"] = [ln.strip() for ln in
|
| 44 |
+
p.read_text().splitlines() if ln.strip()]
|
| 45 |
+
assert len(_STATE["classes"]) == N_CLASSES
|
| 46 |
+
return _STATE["classes"]
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _build_probe():
|
| 50 |
+
import torch
|
| 51 |
+
import torch.nn as nn
|
| 52 |
+
import torch.nn.functional as F
|
| 53 |
+
|
| 54 |
+
dim, heads = 1024, 16
|
| 55 |
+
|
| 56 |
+
class MLP(nn.Module):
|
| 57 |
+
def __init__(self):
|
| 58 |
+
super().__init__()
|
| 59 |
+
self.fc1 = nn.Linear(dim, dim * 4)
|
| 60 |
+
self.act = nn.GELU()
|
| 61 |
+
self.fc2 = nn.Linear(dim * 4, dim)
|
| 62 |
+
|
| 63 |
+
def forward(self, x):
|
| 64 |
+
return self.fc2(self.act(self.fc1(x)))
|
| 65 |
+
|
| 66 |
+
class Attention(nn.Module):
|
| 67 |
+
def __init__(self):
|
| 68 |
+
super().__init__()
|
| 69 |
+
self.qkv = nn.Linear(dim, dim * 3, bias=True)
|
| 70 |
+
self.proj = nn.Linear(dim, dim)
|
| 71 |
+
|
| 72 |
+
def forward(self, x):
|
| 73 |
+
B, N, C = x.shape
|
| 74 |
+
qkv = self.qkv(x).reshape(B, N, 3, heads, C // heads) \
|
| 75 |
+
.permute(2, 0, 3, 1, 4)
|
| 76 |
+
y = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2])
|
| 77 |
+
return self.proj(y.transpose(1, 2).reshape(B, N, C))
|
| 78 |
+
|
| 79 |
+
class Block(nn.Module):
|
| 80 |
+
def __init__(self):
|
| 81 |
+
super().__init__()
|
| 82 |
+
self.norm1 = nn.LayerNorm(dim)
|
| 83 |
+
self.attn = Attention()
|
| 84 |
+
self.norm2 = nn.LayerNorm(dim)
|
| 85 |
+
self.mlp = MLP()
|
| 86 |
+
|
| 87 |
+
def forward(self, x):
|
| 88 |
+
x = x + self.attn(self.norm1(x))
|
| 89 |
+
return x + self.mlp(self.norm2(x))
|
| 90 |
+
|
| 91 |
+
class CrossAttention(nn.Module):
|
| 92 |
+
# NOTE: Meta's probe CrossAttention has NO output projection
|
| 93 |
+
def __init__(self):
|
| 94 |
+
super().__init__()
|
| 95 |
+
self.q = nn.Linear(dim, dim, bias=True)
|
| 96 |
+
self.kv = nn.Linear(dim, dim * 2, bias=True)
|
| 97 |
+
|
| 98 |
+
def forward(self, q, x):
|
| 99 |
+
B, n, C = q.shape
|
| 100 |
+
qh = self.q(q).reshape(B, n, heads, C // heads) \
|
| 101 |
+
.transpose(1, 2)
|
| 102 |
+
N = x.shape[1]
|
| 103 |
+
kv = self.kv(x).reshape(B, N, 2, heads, C // heads) \
|
| 104 |
+
.permute(2, 0, 3, 1, 4)
|
| 105 |
+
y = F.scaled_dot_product_attention(qh, kv[0], kv[1])
|
| 106 |
+
return y.transpose(1, 2).reshape(B, n, C)
|
| 107 |
+
|
| 108 |
+
class CrossAttentionBlock(nn.Module):
|
| 109 |
+
# norm1 normalizes the CONTEXT tokens, not the query (Meta)
|
| 110 |
+
def __init__(self):
|
| 111 |
+
super().__init__()
|
| 112 |
+
self.norm1 = nn.LayerNorm(dim)
|
| 113 |
+
self.xattn = CrossAttention()
|
| 114 |
+
self.norm2 = nn.LayerNorm(dim)
|
| 115 |
+
self.mlp = MLP()
|
| 116 |
+
|
| 117 |
+
def forward(self, q, x):
|
| 118 |
+
q = q + self.xattn(q, self.norm1(x))
|
| 119 |
+
return q + self.mlp(self.norm2(q))
|
| 120 |
+
|
| 121 |
+
class AttentivePooler(nn.Module):
|
| 122 |
+
def __init__(self):
|
| 123 |
+
super().__init__()
|
| 124 |
+
self.query_tokens = nn.Parameter(torch.zeros(1, 1, dim))
|
| 125 |
+
self.cross_attention_block = CrossAttentionBlock()
|
| 126 |
+
self.blocks = nn.ModuleList([Block() for _ in range(3)])
|
| 127 |
+
|
| 128 |
+
def forward(self, x):
|
| 129 |
+
for blk in self.blocks:
|
| 130 |
+
x = blk(x)
|
| 131 |
+
q = self.query_tokens.repeat(len(x), 1, 1)
|
| 132 |
+
return self.cross_attention_block(q, x)
|
| 133 |
+
|
| 134 |
+
class AttentiveClassifier(nn.Module):
|
| 135 |
+
def __init__(self):
|
| 136 |
+
super().__init__()
|
| 137 |
+
self.pooler = AttentivePooler()
|
| 138 |
+
self.linear = nn.Linear(dim, N_CLASSES)
|
| 139 |
+
|
| 140 |
+
def forward(self, x):
|
| 141 |
+
return self.linear(self.pooler(x).squeeze(1))
|
| 142 |
+
|
| 143 |
+
return AttentiveClassifier()
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def _load():
|
| 147 |
+
if "model" in _STATE:
|
| 148 |
+
return _STATE
|
| 149 |
+
import torch
|
| 150 |
+
from transformers import AutoModel, AutoVideoProcessor
|
| 151 |
+
from .device import pick
|
| 152 |
+
dev, dtype = pick()
|
| 153 |
+
enc = AutoModel.from_pretrained(ENCODER_ID, dtype=dtype) \
|
| 154 |
+
.to(dev).eval()
|
| 155 |
+
probe = _build_probe()
|
| 156 |
+
sd = torch.load(PROBE_CKPT, map_location="cpu",
|
| 157 |
+
weights_only=False)["classifiers"][0]
|
| 158 |
+
sd = {k.replace("module.", ""): v for k, v in sd.items()}
|
| 159 |
+
probe.load_state_dict(sd, strict=True)
|
| 160 |
+
probe = probe.to(dev).float().eval()
|
| 161 |
+
_STATE.update(model=enc, probe=probe, dev=dev, dtype=dtype,
|
| 162 |
+
proc=AutoVideoProcessor.from_pretrained(ENCODER_ID))
|
| 163 |
+
return _STATE
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def clip_action_probs(frames_u8):
|
| 167 |
+
"""16 HWC uint8 frames -> softmax over 174 SSv2 classes."""
|
| 168 |
+
import torch
|
| 169 |
+
st = _load()
|
| 170 |
+
px = st["proc"](videos=[list(frames_u8)], return_tensors="pt")[
|
| 171 |
+
"pixel_values_videos"].to(st["dev"], st["dtype"])
|
| 172 |
+
with torch.no_grad():
|
| 173 |
+
feats = st["model"](pixel_values_videos=px).last_hidden_state
|
| 174 |
+
logits = st["probe"](feats.float())
|
| 175 |
+
return torch.softmax(logits[0], -1).cpu().numpy()
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
def query_class_weights(text, top=5):
|
| 179 |
+
"""Query text -> sparse weights over SSv2 classes via the PE text
|
| 180 |
+
tower (cached). Softmax over the top matches; everything else 0."""
|
| 181 |
+
from .pe import _text_vec
|
| 182 |
+
if "clsvec" not in _STATE:
|
| 183 |
+
_STATE["clsvec"] = np.stack(
|
| 184 |
+
[_text_vec(c.lower()) for c in ssv2_classes()])
|
| 185 |
+
sims = _STATE["clsvec"] @ _text_vec(text)
|
| 186 |
+
w = np.zeros(N_CLASSES)
|
| 187 |
+
ix = np.argsort(-sims)[:top]
|
| 188 |
+
e = np.exp((sims[ix] - sims[ix].max()) / 0.05)
|
| 189 |
+
w[ix] = e / e.sum()
|
| 190 |
+
return w
|
python/elidedb/ann.py
ADDED
|
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ANN tiers over the embeddings table. Three mechanisms, one contract:
|
| 2 |
+
|
| 3 |
+
- exact one matmul; correct by definition; right up to ~10^6 vectors
|
| 4 |
+
- IVF learned cells (HDBSCAN centroids, embeddings.py) — prune cells,
|
| 5 |
+
scan survivors exactly; noise always scanned
|
| 6 |
+
- HNSW graph ANN (hnswlib) for sub-ms search past the matmul crossover
|
| 7 |
+
- IVF-PQ product quantization: vectors compressed ~48x into subspace
|
| 8 |
+
codes; asymmetric-distance scan + EXACT rerank of the top pool
|
| 9 |
+
(the SCANN/Faiss recipe: approximate to shortlist, never to
|
| 10 |
+
answer)
|
| 11 |
+
|
| 12 |
+
Artifacts are version-suffixed sidecars under tables/embeddings/_index/,
|
| 13 |
+
rebuilt like any derived state, and recorded in the table log. Search picks
|
| 14 |
+
the best available tier automatically; every path supports HYBRID
|
| 15 |
+
retrieval — time-range and stream predicates pushed into candidate
|
| 16 |
+
selection so vector search composes with the store's core dimension, time.
|
| 17 |
+
"""
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import json
|
| 21 |
+
|
| 22 |
+
import numpy as np
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def _run_dir(store):
|
| 26 |
+
d = store.dir / "tables" / "embeddings" / "_index"
|
| 27 |
+
d.mkdir(parents=True, exist_ok=True)
|
| 28 |
+
return d
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def _emb(store):
|
| 32 |
+
from .embeddings import _vec_table
|
| 33 |
+
return _vec_table(store)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def build_hnsw(store, M: int = 16, ef_construction: int = 200) -> dict:
|
| 37 |
+
"""Hierarchical navigable small world graph over unit vectors
|
| 38 |
+
(inner-product space == cosine). O(log n) expected search hops."""
|
| 39 |
+
import hnswlib
|
| 40 |
+
t, vecs = _emb(store)
|
| 41 |
+
v = store.table("embeddings").state().version
|
| 42 |
+
ix = hnswlib.Index(space="ip", dim=vecs.shape[1])
|
| 43 |
+
ix.init_index(max_elements=len(vecs), M=M,
|
| 44 |
+
ef_construction=ef_construction, random_seed=7)
|
| 45 |
+
ix.add_items(vecs, np.arange(len(vecs)))
|
| 46 |
+
path = _run_dir(store) / f"hnsw.v{v}.bin"
|
| 47 |
+
ix.save_index(str(path))
|
| 48 |
+
(path.with_suffix(".json")).write_text(json.dumps(
|
| 49 |
+
{"dim": int(vecs.shape[1]), "n": len(vecs), "M": M}))
|
| 50 |
+
# An ANN index is a derived sidecar keyed to the embeddings DATA version,
|
| 51 |
+
# not a new data version — building it must NOT advance the log (doing so
|
| 52 |
+
# would invalidate the artifact we just named .v{v} against the bumped
|
| 53 |
+
# current version). The .v{v} filename IS the binding.
|
| 54 |
+
return {"n": len(vecs), "bytes": path.stat().st_size, "version": v}
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def load_hnsw(store):
|
| 58 |
+
# Look for the artifact BEFORE importing. hnswlib is an optional
|
| 59 |
+
# accelerator; when it is absent the planner must fall back to the exact
|
| 60 |
+
# scan, not raise ModuleNotFoundError out of the middle of a query and
|
| 61 |
+
# take down search entirely.
|
| 62 |
+
cands = sorted(_run_dir(store).glob("hnsw.v*.bin"), reverse=True)
|
| 63 |
+
if not cands:
|
| 64 |
+
return None
|
| 65 |
+
try:
|
| 66 |
+
import hnswlib
|
| 67 |
+
except ImportError:
|
| 68 |
+
return None
|
| 69 |
+
v = store.table("embeddings").state().version
|
| 70 |
+
for cand in cands:
|
| 71 |
+
if int(cand.stem.split(".v")[-1]) != v:
|
| 72 |
+
continue # stale: embeddings changed since this was built
|
| 73 |
+
meta = json.loads(cand.with_suffix(".json").read_text())
|
| 74 |
+
ix = hnswlib.Index(space="ip", dim=meta["dim"])
|
| 75 |
+
ix.load_index(str(cand), max_elements=meta["n"])
|
| 76 |
+
return ix
|
| 77 |
+
return None
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def build_ivfpq(store, nlist: int | None = None, m: int = 8,
|
| 81 |
+
nbits: int = 8) -> dict:
|
| 82 |
+
"""Coarse k-means cells + per-subspace codebooks. A d=1152 float32
|
| 83 |
+
vector becomes `m` uint8 codes (m=8 → 576x… realistically 4608B → 8B =
|
| 84 |
+
576x raw, ~48x vs the parquet-compressed vectors). Scans read codes, not
|
| 85 |
+
vectors; the exact rerank reads only the shortlist's true vectors."""
|
| 86 |
+
from sklearn.cluster import KMeans
|
| 87 |
+
t, vecs = _emb(store)
|
| 88 |
+
n, d = vecs.shape
|
| 89 |
+
v = store.table("embeddings").state().version
|
| 90 |
+
nlist = nlist or max(1, int(np.sqrt(n)))
|
| 91 |
+
coarse = KMeans(n_clusters=min(nlist, n), n_init=4,
|
| 92 |
+
random_state=0).fit(vecs)
|
| 93 |
+
resid = vecs - coarse.cluster_centers_[coarse.labels_]
|
| 94 |
+
assert d % m == 0, f"dim {d} not divisible by m={m}"
|
| 95 |
+
sub = d // m
|
| 96 |
+
codebooks = np.zeros((m, 2 ** nbits, sub), np.float32)
|
| 97 |
+
codes = np.zeros((n, m), np.uint8)
|
| 98 |
+
for j in range(m):
|
| 99 |
+
block = resid[:, j * sub:(j + 1) * sub]
|
| 100 |
+
km = KMeans(n_clusters=min(2 ** nbits, n), n_init=2,
|
| 101 |
+
random_state=j).fit(block)
|
| 102 |
+
k = km.cluster_centers_.shape[0]
|
| 103 |
+
codebooks[j, :k] = km.cluster_centers_
|
| 104 |
+
codes[:, j] = km.labels_.astype(np.uint8)
|
| 105 |
+
path = _run_dir(store) / f"ivfpq.v{v}.npz"
|
| 106 |
+
np.savez_compressed(path, centers=coarse.cluster_centers_,
|
| 107 |
+
labels=coarse.labels_.astype(np.int32),
|
| 108 |
+
codebooks=codebooks, codes=codes)
|
| 109 |
+
# derived sidecar keyed to the embeddings data version; no log bump (see
|
| 110 |
+
# build_hnsw) — the .v{v} filename binds it to the current vectors.
|
| 111 |
+
return {"n": n, "nlist": int(nlist), "m": m,
|
| 112 |
+
"bytes": path.stat().st_size,
|
| 113 |
+
"code_bytes_per_vec": m, "version": v}
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def load_ivfpq(store):
|
| 117 |
+
v = store.table("embeddings").state().version
|
| 118 |
+
for cand in sorted(_run_dir(store).glob("ivfpq.v*.npz"), reverse=True):
|
| 119 |
+
if int(cand.stem.split(".v")[-1]) == v:
|
| 120 |
+
return np.load(cand)
|
| 121 |
+
return None
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def search_ivfpq(store, q: np.ndarray, k: int, nprobe: int = 8,
|
| 125 |
+
rerank: int = 4, mask: np.ndarray | None = None):
|
| 126 |
+
"""ADC scan: distance ≈ coarse-center dot + sum of per-subspace code
|
| 127 |
+
dots (table lookups, no vector reads), then exact rerank of the top
|
| 128 |
+
`rerank*k` shortlist. Approximation shortlists; it never answers."""
|
| 129 |
+
art = load_ivfpq(store)
|
| 130 |
+
if art is None:
|
| 131 |
+
return None
|
| 132 |
+
t, vecs = _emb(store)
|
| 133 |
+
centers, labels = art["centers"], art["labels"]
|
| 134 |
+
codebooks, codes = art["codebooks"], art["codes"]
|
| 135 |
+
m, ksub, sub = codebooks.shape
|
| 136 |
+
probe = np.argsort(centers @ q)[::-1][:nprobe]
|
| 137 |
+
cand = np.isin(labels, probe)
|
| 138 |
+
if mask is not None:
|
| 139 |
+
cand &= mask
|
| 140 |
+
idx = np.where(cand)[0]
|
| 141 |
+
if len(idx) == 0:
|
| 142 |
+
return [], {"scanned": 0, "total": len(vecs)}
|
| 143 |
+
# lookup tables: q-subvector · every codeword, per subspace
|
| 144 |
+
lut = np.stack([codebooks[j] @ q[j * sub:(j + 1) * sub]
|
| 145 |
+
for j in range(m)]) # [m, ksub]
|
| 146 |
+
approx = centers[labels[idx]] @ q + \
|
| 147 |
+
lut[np.arange(m)[None, :], codes[idx]].sum(axis=1)
|
| 148 |
+
short = idx[np.argsort(approx)[::-1][:max(k * rerank, k)]]
|
| 149 |
+
exact = vecs[short] @ q
|
| 150 |
+
order = np.argsort(exact)[::-1][:k]
|
| 151 |
+
return ([(int(short[i]), float(exact[i])) for i in order],
|
| 152 |
+
{"scanned": int(len(idx)), "total": len(vecs),
|
| 153 |
+
"code_bytes": int(len(idx) * m),
|
| 154 |
+
"reranked": int(len(short))})
|
python/elidedb/answer.py
ADDED
|
@@ -0,0 +1,338 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ANSWER: the join of scene, agent, participants and events. No text.
|
| 2 |
+
|
| 3 |
+
The user's definition, verbatim: "answer is just a join... everything is
|
| 4 |
+
just logical." So this module contains no text tower, no fitted weight,
|
| 5 |
+
no learned combiner - only the element tables and three logical
|
| 6 |
+
operations:
|
| 7 |
+
|
| 8 |
+
AND min over percentile ranks (scale-free conjunction)
|
| 9 |
+
OR max over a candidate episode's objects
|
| 10 |
+
BIND the AND is evaluated PER CANDIDATE OBJECT - "an object that
|
| 11 |
+
looks like X and moves like X" must be satisfied by one
|
| 12 |
+
object, not by one object that looks right and a different
|
| 13 |
+
one that moves right. This binding is the entire difference
|
| 14 |
+
between a join and a channel soup, and it is what "object1
|
| 15 |
+
(and similar) undergoing action1 (and similar)" means.
|
| 16 |
+
|
| 17 |
+
Percentile ranks, not raw cosines, because the terms live on different
|
| 18 |
+
scales (DINOv3 identity cosine, path-delta cosine, V-JEPA cosine) and a
|
| 19 |
+
min over raw scales lets the tightest-distributed term govern
|
| 20 |
+
everything. Rank-normalising is order statistics, not a weight.
|
| 21 |
+
|
| 22 |
+
Element terms per query participant, each one a table:
|
| 23 |
+
|
| 24 |
+
looks object_vectors DINOv3 track descriptor cosine
|
| 25 |
+
moves trajectories signed, speed-normalised (dx,dy,dz) path
|
| 26 |
+
deltas resampled to N steps - direction
|
| 27 |
+
survives, which is the thing text erases
|
| 28 |
+
acts vjepa_part physics tubelet cosine, where both sides
|
| 29 |
+
have one
|
| 30 |
+
|
| 31 |
+
Episode-level terms:
|
| 32 |
+
|
| 33 |
+
scene scene_vectors pooled window cosine, one more rank in the
|
| 34 |
+
conjunction
|
| 35 |
+
events kind overlap a PARTITION, not a score: episodes sharing
|
| 36 |
+
at least one event kind with the query rank
|
| 37 |
+
ahead of episodes sharing none. Logical
|
| 38 |
+
precedence, no number invented.
|
| 39 |
+
"""
|
| 40 |
+
from __future__ import annotations
|
| 41 |
+
|
| 42 |
+
import numpy as np
|
| 43 |
+
|
| 44 |
+
_C = {}
|
| 45 |
+
|
| 46 |
+
NPATH = 16 # path resample steps; 15 deltas x 3 dims
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _ranks(x):
|
| 50 |
+
"""Percentile rank in [0,1] per entry; NaN stays NaN (missing)."""
|
| 51 |
+
x = np.asarray(x, np.float64)
|
| 52 |
+
out = np.full(len(x), np.nan)
|
| 53 |
+
ok = np.isfinite(x)
|
| 54 |
+
if ok.sum() > 1:
|
| 55 |
+
r = x[ok].argsort().argsort()
|
| 56 |
+
out[ok] = r / (ok.sum() - 1)
|
| 57 |
+
elif ok.sum() == 1:
|
| 58 |
+
out[ok] = 1.0
|
| 59 |
+
return out
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def _path_desc(ts, px, py, pz, diag):
|
| 63 |
+
"""Signed, speed-normalised path deltas. Translation-invariant by
|
| 64 |
+
construction (deltas), scale-normalised by the object's own box
|
| 65 |
+
diagonal, direction preserved (open vs close differ by SIGN, the
|
| 66 |
+
one thing appearance embeddings collapse - cos 0.957 measured)."""
|
| 67 |
+
if len(ts) < 3:
|
| 68 |
+
return None
|
| 69 |
+
o = np.argsort(ts)
|
| 70 |
+
t = np.linspace(0, len(o) - 1, NPATH).round().astype(int)
|
| 71 |
+
x, y, z = (np.asarray(v, np.float64)[o][t] for v in (px, py, pz))
|
| 72 |
+
d = np.stack([np.diff(x) / max(diag, 1.0),
|
| 73 |
+
np.diff(y) / max(diag, 1.0),
|
| 74 |
+
np.diff(z) * 4.0], 1).ravel()
|
| 75 |
+
n = np.linalg.norm(d)
|
| 76 |
+
return (d / n).astype(np.float32) if n > 1e-6 else None
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def _elements(store):
|
| 80 |
+
"""Every element table, loaded once per store version, joined on
|
| 81 |
+
the track key (stream, track_ts, t1, object_id)."""
|
| 82 |
+
ver = store.table("trajectories").state().version
|
| 83 |
+
key = (str(store.dir), ver)
|
| 84 |
+
if key in _C:
|
| 85 |
+
return _C[key]
|
| 86 |
+
E = {}
|
| 87 |
+
tr = store.table("trajectories").scan().to_pydict()
|
| 88 |
+
by = {}
|
| 89 |
+
for i in range(len(tr["ts"])):
|
| 90 |
+
k = (str(tr["stream"][i]), int(tr["track_ts"][i]),
|
| 91 |
+
int(tr["t1"][i]), int(tr["object_id"][i]))
|
| 92 |
+
by.setdefault(k, []).append(i)
|
| 93 |
+
paths, agents = {}, {}
|
| 94 |
+
for k, idx in by.items():
|
| 95 |
+
bx = [(tr["x1"][i] - tr["x0"][i], tr["y1"][i] - tr["y0"][i])
|
| 96 |
+
for i in idx]
|
| 97 |
+
diag = float(np.median([np.hypot(w, h) for w, h in bx]))
|
| 98 |
+
d = _path_desc([tr["ts"][i] for i in idx],
|
| 99 |
+
[tr["px"][i] for i in idx],
|
| 100 |
+
[tr["py"][i] for i in idx],
|
| 101 |
+
[tr["pz"][i] for i in idx], diag)
|
| 102 |
+
if d is not None:
|
| 103 |
+
paths[k] = d
|
| 104 |
+
if any(tr["is_agent"][i] for i in idx):
|
| 105 |
+
agents[k] = True
|
| 106 |
+
E["paths"] = paths
|
| 107 |
+
|
| 108 |
+
# key -> ROW index; double-detection twins share a key, so the
|
| 109 |
+
# dict is smaller than the table and the reshape must use the ROW
|
| 110 |
+
# count - keying by len(dict) sheared the matrix off by 2,086 rows
|
| 111 |
+
ov = store.table("object_vectors").scan().to_pydict()
|
| 112 |
+
E["obj"] = {(str(s), int(a), int(b), int(o)): i for i, (s, a, b, o)
|
| 113 |
+
in enumerate(zip(ov["stream"], ov["ts"], ov["t1"],
|
| 114 |
+
ov["object_id"]))}
|
| 115 |
+
V = np.asarray(ov["vector"], np.float32).reshape(len(ov["ts"]), -1)
|
| 116 |
+
E["objV"] = V / np.maximum(
|
| 117 |
+
np.linalg.norm(V, axis=1, keepdims=True), 1e-8)
|
| 118 |
+
|
| 119 |
+
vp = store.table("vjepa_part_vectors").scan().to_pydict()
|
| 120 |
+
E["phys"] = {(str(s), int(a), int(b), int(o)): i for i, (s, a, b, o)
|
| 121 |
+
in enumerate(zip(vp["stream"], vp["ts"], vp["t1"],
|
| 122 |
+
vp["object_id"]))}
|
| 123 |
+
Vp = np.asarray(vp["vector"], np.float32).reshape(len(vp["ts"]), -1)
|
| 124 |
+
E["physV"] = Vp / np.maximum(
|
| 125 |
+
np.linalg.norm(Vp, axis=1, keepdims=True), 1e-8)
|
| 126 |
+
|
| 127 |
+
sc = store.table("scene_vectors").scan().to_pydict()
|
| 128 |
+
E["scene"] = (sc["stream"], np.asarray(sc["ts"], np.int64),
|
| 129 |
+
np.asarray(sc["vector"], np.float32)
|
| 130 |
+
.reshape(len(sc["ts"]), -1))
|
| 131 |
+
|
| 132 |
+
ev = store.table("events").scan().to_pydict()
|
| 133 |
+
E["events"] = ev
|
| 134 |
+
E["agents"] = agents
|
| 135 |
+
|
| 136 |
+
# EVENT-LEVEL join arrays: each event row -> its motion vector, its
|
| 137 |
+
# bound object's DINOv3 descriptor row, its physics row, its
|
| 138 |
+
# episode. The image-plane path descriptor above measured 0.08
|
| 139 |
+
# yield on the direction queries - 2D geometry does not survive a
|
| 140 |
+
# camera change, and the supports span four cameras. The per-event
|
| 141 |
+
# delta-appearance vector does (it is appearance change, not
|
| 142 |
+
# coordinates), so it is the join's moves-term; geometry remains
|
| 143 |
+
# the fallback where a window has no bound events.
|
| 144 |
+
mv = store.table("motion_vectors").scan().to_pydict()
|
| 145 |
+
mrow = {(str(s), int(a), int(b)): i for i, (s, a, b) in
|
| 146 |
+
enumerate(zip(mv["stream"], mv["ts"], mv["t1"]))}
|
| 147 |
+
MV = np.asarray(mv["vector"], np.float32).reshape(len(mv["ts"]), -1)
|
| 148 |
+
MV /= np.maximum(np.linalg.norm(MV, axis=1, keepdims=True), 1e-8)
|
| 149 |
+
tracks_of = {}
|
| 150 |
+
for k in paths:
|
| 151 |
+
tracks_of.setdefault((k[0], k[3]), []).append(k)
|
| 152 |
+
n_ev = len(ev["ts"])
|
| 153 |
+
e_m = np.full(n_ev, -1)
|
| 154 |
+
e_obj = np.full(n_ev, -1)
|
| 155 |
+
e_phys = np.full(n_ev, -1)
|
| 156 |
+
for i in range(n_ev):
|
| 157 |
+
s, a, b = str(ev["stream"][i]), int(ev["ts"][i]), int(ev["t1"][i])
|
| 158 |
+
e_m[i] = mrow.get((s, a, b), -1)
|
| 159 |
+
oid = int(ev["object_id"][i])
|
| 160 |
+
if oid >= 0:
|
| 161 |
+
for k in tracks_of.get((s, oid), ()):
|
| 162 |
+
if k[1] <= b and a <= k[2]:
|
| 163 |
+
e_obj[i] = E["obj"].get(k, -1)
|
| 164 |
+
e_phys[i] = E["phys"].get(k, -1)
|
| 165 |
+
break
|
| 166 |
+
# THE AGENT HAS NO LOOKS-TERM. Its identity is the ROLE - the
|
| 167 |
+
# self-moving thing - not an appearance. Measured: agent-track
|
| 168 |
+
# descriptors cosine 0.477 median against each other, barely
|
| 169 |
+
# above random non-agent pairs at 0.437, because an articulated
|
| 170 |
+
# arm deforms; only 0.2% of same-arm pairs clear the identity
|
| 171 |
+
# cut, which is where 1,517 ids for ~4 arms came from. No
|
| 172 |
+
# threshold merges what does not separate, so an agent event
|
| 173 |
+
# matches by moves and physics alone, and its appearance rank -
|
| 174 |
+
# noise by measurement - never enters the conjunction.
|
| 175 |
+
if (ev["role"][i] == "agent"
|
| 176 |
+
or (oid >= 0 and oid == int(ev["agent_object_id"][i]))):
|
| 177 |
+
e_obj[i] = -1
|
| 178 |
+
E["e_m"], E["e_obj"], E["e_phys"], E["MV"] = e_m, e_obj, e_phys, MV
|
| 179 |
+
|
| 180 |
+
ep = store.table("episodes").scan().to_pydict()
|
| 181 |
+
E["episodes"] = [(str(s), int(a), int(b)) for s, a, b in
|
| 182 |
+
zip(ep["stream"], ep["ts"], ep["t1"])]
|
| 183 |
+
if len(_C) > 4:
|
| 184 |
+
_C.clear()
|
| 185 |
+
_C[key] = E
|
| 186 |
+
return E
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
def _tracks_in(E, stream, t0, t1):
|
| 190 |
+
return [k for k in E["paths"]
|
| 191 |
+
if k[0] == stream and k[1] <= t1 and t0 <= k[2]]
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
def _scene_vec(E, stream, t0, t1):
|
| 195 |
+
ss, ts, V = E["scene"]
|
| 196 |
+
m = np.array([s == stream and t0 <= t <= t1
|
| 197 |
+
for s, t in zip(ss, ts)])
|
| 198 |
+
if not m.any():
|
| 199 |
+
return None
|
| 200 |
+
v = V[m].mean(0)
|
| 201 |
+
return v / (np.linalg.norm(v) + 1e-8)
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def answer_like(store, stream, t0, t1):
|
| 205 |
+
"""Episodes telling the same story as the query window.
|
| 206 |
+
|
| 207 |
+
Returns (keys, score, shared_kind) - score is the conjunction rank
|
| 208 |
+
(higher = better), shared_kind marks the event-kind partition.
|
| 209 |
+
"""
|
| 210 |
+
E = _elements(store)
|
| 211 |
+
eps = E["episodes"]
|
| 212 |
+
n = len(eps)
|
| 213 |
+
|
| 214 |
+
# ---- the query's elements -----------------------------------
|
| 215 |
+
q_tracks = _tracks_in(E, stream, t0, t1)
|
| 216 |
+
ev = E["events"]
|
| 217 |
+
q_kinds = {ev["kind"][i] for i in range(len(ev["ts"]))
|
| 218 |
+
if str(ev["stream"][i]) == stream
|
| 219 |
+
and int(ev["ts"][i]) >= t0 and int(ev["t1"][i]) <= t1
|
| 220 |
+
and ev["kind"][i]}
|
| 221 |
+
ep_of = {}
|
| 222 |
+
for i, (s, a, b) in enumerate(eps):
|
| 223 |
+
ep_of.setdefault(s, []).append((a, b, i))
|
| 224 |
+
|
| 225 |
+
# events of the query window, and each event's episode index
|
| 226 |
+
q_ev = [i for i in range(len(ev["ts"]))
|
| 227 |
+
if str(ev["stream"][i]) == stream
|
| 228 |
+
and int(ev["ts"][i]) >= t0 and int(ev["t1"][i]) <= t1]
|
| 229 |
+
n_ev = len(ev["ts"])
|
| 230 |
+
ev_ep = np.full(n_ev, -1)
|
| 231 |
+
for i in range(n_ev):
|
| 232 |
+
s, a = str(ev["stream"][i]), int(ev["ts"][i])
|
| 233 |
+
for ea, eb, j in ep_of.get(s, ()):
|
| 234 |
+
if a >= ea and a <= eb:
|
| 235 |
+
ev_ep[i] = j
|
| 236 |
+
break
|
| 237 |
+
|
| 238 |
+
# ---- EVENT-LEVEL join (primary) -----------------------------
|
| 239 |
+
# "the same transition happening to the same kind of thing": per
|
| 240 |
+
# query event, every candidate event is scored on moves (delta-
|
| 241 |
+
# appearance of the transition - the view-tolerant motion term;
|
| 242 |
+
# image-plane geometry measured 0.08 on the direction queries
|
| 243 |
+
# because the supports span four cameras) AND looks (DINOv3 of the
|
| 244 |
+
# object each event BINDS to) AND acts (its physics tubelet). One
|
| 245 |
+
# candidate event must satisfy all of it - that is the join.
|
| 246 |
+
# conjunction granularity is the KIND, not the event instance: a
|
| 247 |
+
# window holding five t3 relocations and one t7 asks for "t3-like
|
| 248 |
+
# AND t7-like", not six independent demands. OR (max) within a
|
| 249 |
+
# kind's query events, AND (min) across kinds - the story's clauses
|
| 250 |
+
# are its distinct transitions.
|
| 251 |
+
q_bound = [i for i in q_ev if E["e_m"][i] >= 0]
|
| 252 |
+
by_kind = {}
|
| 253 |
+
for qi in q_bound:
|
| 254 |
+
by_kind.setdefault(ev["kind"][qi] or "", []).append(qi)
|
| 255 |
+
per_kind = []
|
| 256 |
+
for kind, members in by_kind.items():
|
| 257 |
+
rows = []
|
| 258 |
+
for qi in members:
|
| 259 |
+
terms = []
|
| 260 |
+
terms.append(_ranks(E["MV"] @ E["MV"][E["e_m"][qi]]))
|
| 261 |
+
if E["e_obj"][qi] >= 0:
|
| 262 |
+
qv = E["objV"][E["e_obj"][qi]]
|
| 263 |
+
sims = np.full(n_ev, np.nan)
|
| 264 |
+
has = E["e_obj"] >= 0
|
| 265 |
+
sims[has] = E["objV"][E["e_obj"][has]] @ qv
|
| 266 |
+
terms.append(_ranks(sims))
|
| 267 |
+
if E["e_phys"][qi] >= 0:
|
| 268 |
+
qv = E["physV"][E["e_phys"][qi]]
|
| 269 |
+
sims = np.full(n_ev, np.nan)
|
| 270 |
+
has = E["e_phys"] >= 0
|
| 271 |
+
sims[has] = E["physV"][E["e_phys"][has]] @ qv
|
| 272 |
+
terms.append(_ranks(sims))
|
| 273 |
+
with np.errstate(invalid="ignore"):
|
| 274 |
+
bound = np.nanmin(np.stack(terms), 0)
|
| 275 |
+
row = np.full(n, np.nan)
|
| 276 |
+
for i in range(n_ev):
|
| 277 |
+
e = ev_ep[i]
|
| 278 |
+
if e >= 0 and np.isfinite(bound[i]):
|
| 279 |
+
row[e] = bound[i] if np.isnan(row[e]) \
|
| 280 |
+
else max(row[e], bound[i])
|
| 281 |
+
rows.append(row)
|
| 282 |
+
with np.errstate(invalid="ignore"):
|
| 283 |
+
per_kind.append(np.nanmax(np.stack(rows), 0))
|
| 284 |
+
|
| 285 |
+
if per_kind:
|
| 286 |
+
with np.errstate(invalid="ignore"):
|
| 287 |
+
score = np.nanmin(np.stack(per_kind), 0) # AND across kinds
|
| 288 |
+
score = np.where(np.isnan(score), 0.0, score)
|
| 289 |
+
else:
|
| 290 |
+
# ---- track-path fallback: a window with no bound events ----
|
| 291 |
+
all_keys = list(E["paths"].keys())
|
| 292 |
+
P = np.stack([E["paths"][k] for k in all_keys])
|
| 293 |
+
track_ep = np.full(len(all_keys), -1)
|
| 294 |
+
for i, k in enumerate(all_keys):
|
| 295 |
+
for a, b, j in ep_of.get(k[0], ()):
|
| 296 |
+
if k[1] >= a and k[1] <= b:
|
| 297 |
+
track_ep[i] = j
|
| 298 |
+
break
|
| 299 |
+
q_parts = [k for k in q_tracks if E["agents"].get(k)] or q_tracks
|
| 300 |
+
if not q_parts:
|
| 301 |
+
return eps, np.zeros(n, np.float32), np.zeros(n, bool)
|
| 302 |
+
per = []
|
| 303 |
+
for qk in q_parts:
|
| 304 |
+
bound = _ranks(P @ E["paths"][qk])
|
| 305 |
+
row = np.full(n, np.nan)
|
| 306 |
+
for i, e in enumerate(track_ep):
|
| 307 |
+
if e >= 0 and np.isfinite(bound[i]):
|
| 308 |
+
row[e] = bound[i] if np.isnan(row[e]) \
|
| 309 |
+
else max(row[e], bound[i])
|
| 310 |
+
per.append(row)
|
| 311 |
+
with np.errstate(invalid="ignore"):
|
| 312 |
+
score = np.nanmin(np.stack(per), 0)
|
| 313 |
+
score = np.where(np.isnan(score), 0.0, score)
|
| 314 |
+
|
| 315 |
+
qs = _scene_vec(E, stream, t0, t1)
|
| 316 |
+
if qs is not None:
|
| 317 |
+
ssims = np.full(n, np.nan)
|
| 318 |
+
for i, (s, a, b) in enumerate(eps):
|
| 319 |
+
v = _scene_vec(E, s, a, b)
|
| 320 |
+
if v is not None:
|
| 321 |
+
ssims[i] = float(v @ qs)
|
| 322 |
+
sr = _ranks(ssims)
|
| 323 |
+
with np.errstate(invalid="ignore"):
|
| 324 |
+
score = np.fmin(score, np.where(np.isnan(sr), score, sr))
|
| 325 |
+
|
| 326 |
+
# ---- event-kind PARTITION -----------------------------------
|
| 327 |
+
shared = np.zeros(n, bool)
|
| 328 |
+
if q_kinds:
|
| 329 |
+
for i in range(len(ev["ts"])):
|
| 330 |
+
k = ev["kind"][i]
|
| 331 |
+
if k and k in q_kinds:
|
| 332 |
+
s = str(ev["stream"][i])
|
| 333 |
+
a = int(ev["ts"][i])
|
| 334 |
+
for ea, eb, j in ep_of.get(s, ()):
|
| 335 |
+
if a >= ea and a <= eb:
|
| 336 |
+
shared[j] = True
|
| 337 |
+
break
|
| 338 |
+
return eps, score.astype(np.float32), shared
|
python/elidedb/bptree.py
ADDED
|
@@ -0,0 +1,123 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""BPT1 — immutable, bulk-loaded B+ tree (numpy twin of the C++20 reader in
|
| 2 |
+
src/streetdex/index/bptree.{hpp,cpp}; identical bytes on disk).
|
| 3 |
+
|
| 4 |
+
Why a B+ tree here at all: zone maps prune brilliantly on `ts` because files
|
| 5 |
+
are time-sorted, and prune *nothing* on an unsorted column (every row group
|
| 6 |
+
spans nearly the whole domain). A secondary index re-sorts (key → location)
|
| 7 |
+
once, so point/range predicates on ANY numeric column become a descent plus
|
| 8 |
+
a contiguous leaf scan instead of a full-table read. Immutability makes the
|
| 9 |
+
classic hard parts vanish: bulk load at 100% fill, no splits, no
|
| 10 |
+
rebalancing, rebuilt per table version like every other derived artifact.
|
| 11 |
+
|
| 12 |
+
Values are opaque u64s; the store packs (file_idx << 40) | row_in_file, so a
|
| 13 |
+
leaf hit maps straight to one Parquet row group — the index prunes I/O, not
|
| 14 |
+
just rows.
|
| 15 |
+
"""
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
import struct
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
|
| 21 |
+
import numpy as np
|
| 22 |
+
|
| 23 |
+
MAGIC = b"BPT1"
|
| 24 |
+
DEFAULT_ORDER = 256 # 256 × 8B keys = one 4 KiB page per node touch
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def encode_key(values: np.ndarray) -> np.ndarray:
|
| 28 |
+
"""Order-preserving i64 encoding: ints pass through; IEEE-754 doubles use
|
| 29 |
+
the sign-flip trick (identical to bpt_encode_f64 in C++)."""
|
| 30 |
+
if np.issubdtype(values.dtype, np.integer):
|
| 31 |
+
return values.astype("<i8")
|
| 32 |
+
if np.issubdtype(values.dtype, np.floating):
|
| 33 |
+
u = values.astype("<f8").view("<u8").copy()
|
| 34 |
+
neg = (u & 0x8000000000000000) != 0
|
| 35 |
+
u[neg] = ~u[neg]
|
| 36 |
+
u[~neg] |= 0x8000000000000000
|
| 37 |
+
return (u ^ 0x8000000000000000).view("<i8")
|
| 38 |
+
raise TypeError(f"unindexable dtype {values.dtype}")
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def encode_scalar(v) -> int:
|
| 42 |
+
return int(encode_key(np.array([v]))[0])
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def build(keys: np.ndarray, values: np.ndarray,
|
| 46 |
+
order: int = DEFAULT_ORDER) -> bytes:
|
| 47 |
+
order_idx = np.argsort(keys, kind="stable")
|
| 48 |
+
k = np.ascontiguousarray(keys[order_idx], dtype="<i8")
|
| 49 |
+
v = np.ascontiguousarray(values[order_idx], dtype="<u8")
|
| 50 |
+
n = len(k)
|
| 51 |
+
|
| 52 |
+
levels = [] # bottom-up internal levels: first key of each child node
|
| 53 |
+
below = k
|
| 54 |
+
while len(below) > order:
|
| 55 |
+
firsts = below[::order].copy()
|
| 56 |
+
levels.append(firsts)
|
| 57 |
+
below = firsts
|
| 58 |
+
levels.reverse() # root first, as the C++ builder writes them
|
| 59 |
+
|
| 60 |
+
out = bytearray()
|
| 61 |
+
out += MAGIC
|
| 62 |
+
out += struct.pack("<IQII", order, n, len(levels), 0)
|
| 63 |
+
dir_pos = len(out)
|
| 64 |
+
out += b"\0" * (16 * (len(levels) + 1))
|
| 65 |
+
directory = []
|
| 66 |
+
for lv in levels:
|
| 67 |
+
while len(out) % 8:
|
| 68 |
+
out += b"\0"
|
| 69 |
+
directory.append((len(out), len(lv)))
|
| 70 |
+
out += lv.astype("<i8").tobytes()
|
| 71 |
+
while len(out) % 8:
|
| 72 |
+
out += b"\0"
|
| 73 |
+
directory.append((len(out), n))
|
| 74 |
+
leaf = np.empty(n, dtype=[("k", "<i8"), ("v", "<u8")])
|
| 75 |
+
leaf["k"] = k
|
| 76 |
+
leaf["v"] = v
|
| 77 |
+
out += leaf.tobytes()
|
| 78 |
+
for i, (off, cnt) in enumerate(directory):
|
| 79 |
+
struct.pack_into("<QQ", out, dir_pos + 16 * i, off, cnt)
|
| 80 |
+
return bytes(out)
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
class Reader:
|
| 84 |
+
def __init__(self, data: bytes | np.memmap):
|
| 85 |
+
buf = np.frombuffer(data, dtype=np.uint8) if isinstance(data, bytes) \
|
| 86 |
+
else data
|
| 87 |
+
if bytes(buf[:4]) != MAGIC:
|
| 88 |
+
raise ValueError("not a BPT1 index")
|
| 89 |
+
self.order, self.n, nlevels, _ = struct.unpack_from(
|
| 90 |
+
"<IQII", buf.tobytes()[:24] if isinstance(buf, np.memmap)
|
| 91 |
+
else data, 4)
|
| 92 |
+
raw = buf.tobytes() if isinstance(buf, np.memmap) else data
|
| 93 |
+
dirs = [struct.unpack_from("<QQ", raw, 24 + 16 * i)
|
| 94 |
+
for i in range(nlevels + 1)]
|
| 95 |
+
self.levels = [np.frombuffer(raw, "<i8", cnt, off)
|
| 96 |
+
for off, cnt in dirs[:-1]]
|
| 97 |
+
loff, lcnt = dirs[-1]
|
| 98 |
+
leaf = np.frombuffer(raw, dtype=[("k", "<i8"), ("v", "<u8")],
|
| 99 |
+
count=lcnt, offset=loff)
|
| 100 |
+
self.keys = leaf["k"]
|
| 101 |
+
self.vals = leaf["v"]
|
| 102 |
+
|
| 103 |
+
@classmethod
|
| 104 |
+
def open(cls, path: str | Path):
|
| 105 |
+
return cls(np.memmap(path, dtype=np.uint8, mode="r"))
|
| 106 |
+
|
| 107 |
+
def lower_bound(self, k: int) -> int:
|
| 108 |
+
node = 0
|
| 109 |
+
for lv in self.levels:
|
| 110 |
+
begin = node * self.order
|
| 111 |
+
end = min(begin + self.order, len(lv))
|
| 112 |
+
# lower_bound + step-back (see C++ twin): duplicates of k may
|
| 113 |
+
# start in the child before the first child whose first-key == k
|
| 114 |
+
pos = int(np.searchsorted(lv[begin:end], k, side="left")) + begin
|
| 115 |
+
node = pos - 1 if pos > begin else begin
|
| 116 |
+
begin = node * self.order
|
| 117 |
+
end = min(begin + self.order, self.n)
|
| 118 |
+
return int(np.searchsorted(self.keys[begin:end], k, side="left")) + begin
|
| 119 |
+
|
| 120 |
+
def range(self, lo: int, hi: int) -> np.ndarray:
|
| 121 |
+
a = self.lower_bound(lo)
|
| 122 |
+
b = int(np.searchsorted(self.keys, hi, side="right"))
|
| 123 |
+
return self.vals[a:b]
|
python/elidedb/chanbench.py
ADDED
|
@@ -0,0 +1,221 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""How to judge a channel — teacher or student, same yardstick.
|
| 2 |
+
|
| 3 |
+
Params and ms/clip say nothing about whether a representation works.
|
| 4 |
+
Neither does agreement with the teacher: the V-JEPA2 run made that
|
| 5 |
+
concrete, where the teacher's own clips sat at 0.977 mean cosine to
|
| 6 |
+
their centroid, so a student agreeing with it perfectly would inherit a
|
| 7 |
+
representation that cannot rank anything. Fidelity to a broken teacher
|
| 8 |
+
is a broken student with extra steps.
|
| 9 |
+
|
| 10 |
+
Three layers, and a channel has to pass them IN ORDER. Each one is a
|
| 11 |
+
different question, and a failure at any level makes the levels below
|
| 12 |
+
it unreadable.
|
| 13 |
+
|
| 14 |
+
1 SEPARABILITY does the representation distinguish anything at all?
|
| 15 |
+
mean pairwise cosine and effective rank. Catches the
|
| 16 |
+
collapse that both FDNN-V (0.978) and mean-pooled
|
| 17 |
+
V-JEPA2 (0.977) were hiding behind respectable
|
| 18 |
+
cosine numbers. Costs nothing, run it FIRST, on the
|
| 19 |
+
teacher, before training any student against it.
|
| 20 |
+
|
| 21 |
+
2 TASK AUC is it the RIGHT signal? A linear probe over labels
|
| 22 |
+
the WRITE PATH already produced - event kinds from
|
| 23 |
+
geometry, object ids from the identity store - so
|
| 24 |
+
this measures usefulness without ever touching the
|
| 25 |
+
truthset. That matters: the truthset is eval-only,
|
| 26 |
+
and a per-iteration metric that reads it would make
|
| 27 |
+
every later evaluation meaningless.
|
| 28 |
+
|
| 29 |
+
3 FIDELITY does the student rank like the teacher? Only worth
|
| 30 |
+
asking once 1 and 2 say the teacher is worth
|
| 31 |
+
matching.
|
| 32 |
+
|
| 33 |
+
THE SHIPPING GATE is layer 2, not layer 3: a student ships when its task
|
| 34 |
+
AUC is within tolerance of its teacher's, at a fraction of the cost. A
|
| 35 |
+
student can legitimately disagree with its teacher clip-by-clip and
|
| 36 |
+
still be as useful, and that is a pass, not a failure.
|
| 37 |
+
|
| 38 |
+
The product metric - yield and precision on the truthset at
|
| 39 |
+
k = ceil(1.5 x support) - stays the integration test, run once channels
|
| 40 |
+
are assembled. It is too coarse and too slow to steer a per-channel
|
| 41 |
+
loop, and it is the thing these layers exist to protect.
|
| 42 |
+
"""
|
| 43 |
+
from __future__ import annotations
|
| 44 |
+
|
| 45 |
+
import numpy as np
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _l2(x):
|
| 49 |
+
return x / (np.linalg.norm(x, axis=-1, keepdims=True) + 1e-8)
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def separability(V, sample=2000):
|
| 53 |
+
"""Layer 1. Does the space have room to rank anything?
|
| 54 |
+
|
| 55 |
+
effective_rank is the participation ratio of the covariance
|
| 56 |
+
eigenvalues - how many dimensions the data actually uses. A 1024-d
|
| 57 |
+
embedding with an effective rank of 3 is a 3-d embedding that costs
|
| 58 |
+
1024 floats to store and cannot support a nearest-neighbour query.
|
| 59 |
+
"""
|
| 60 |
+
V = _l2(np.asarray(V, np.float32))
|
| 61 |
+
n = min(len(V), sample)
|
| 62 |
+
S = V[:n] @ V[:n].T
|
| 63 |
+
iu = np.triu_indices(n, 1)
|
| 64 |
+
mu = _l2(V.mean(0))
|
| 65 |
+
ev = np.linalg.eigvalsh(np.cov(V[:n].T))[::-1]
|
| 66 |
+
ev = np.clip(ev, 0, None)
|
| 67 |
+
p = ev / (ev.sum() + 1e-12)
|
| 68 |
+
eff = float(np.exp(-(p * np.log(p + 1e-12)).sum()))
|
| 69 |
+
return {"mean_pairwise_cos": round(float(S[iu].mean()), 4),
|
| 70 |
+
"p95_pairwise_cos": round(float(np.percentile(S[iu], 95)), 4),
|
| 71 |
+
"mean_baseline": round(float(np.mean(V @ mu)), 4),
|
| 72 |
+
"effective_rank": round(eff, 1),
|
| 73 |
+
"dims": int(V.shape[1])}
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def probe_auc(V, y, folds=5, seed=0):
|
| 77 |
+
"""Layer 2. Linear-probe AUC for one binary label.
|
| 78 |
+
|
| 79 |
+
A LINEAR probe on purpose: it asks whether the information is
|
| 80 |
+
present and readable, not whether a big enough head can dig it out.
|
| 81 |
+
That is the property a retrieval channel needs, because the thing
|
| 82 |
+
consuming it downstream is a dot product.
|
| 83 |
+
"""
|
| 84 |
+
from sklearn.linear_model import LogisticRegression
|
| 85 |
+
from sklearn.metrics import roc_auc_score
|
| 86 |
+
V, y = _l2(np.asarray(V, np.float32)), np.asarray(y).astype(int)
|
| 87 |
+
if y.sum() < folds or (1 - y).sum() < folds:
|
| 88 |
+
return None # too few of one class to score
|
| 89 |
+
rng = np.random.default_rng(seed)
|
| 90 |
+
idx = rng.permutation(len(V))
|
| 91 |
+
V, y = V[idx], y[idx]
|
| 92 |
+
out = []
|
| 93 |
+
for f in range(folds):
|
| 94 |
+
te = np.zeros(len(V), bool)
|
| 95 |
+
te[f::folds] = True
|
| 96 |
+
if y[~te].sum() == 0 or y[~te].sum() == (~te).sum():
|
| 97 |
+
continue
|
| 98 |
+
if len(set(y[te])) < 2:
|
| 99 |
+
continue
|
| 100 |
+
m = LogisticRegression(max_iter=2000, C=1.0)
|
| 101 |
+
m.fit(V[~te], y[~te])
|
| 102 |
+
out.append(roc_auc_score(y[te], m.decision_function(V[te])))
|
| 103 |
+
return round(float(np.mean(out)), 4) if out else None
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def retrieval_map(V, groups):
|
| 107 |
+
"""Layer 2, the ranking form. Mean average precision when the query
|
| 108 |
+
is a clip and the relevant set is everything sharing its group.
|
| 109 |
+
|
| 110 |
+
AUC says the information is linearly readable; mAP says it survives
|
| 111 |
+
being turned into a ranking, which is what the channel is for.
|
| 112 |
+
"""
|
| 113 |
+
V = _l2(np.asarray(V, np.float32))
|
| 114 |
+
g = np.asarray(groups)
|
| 115 |
+
S = V @ V.T
|
| 116 |
+
np.fill_diagonal(S, -9)
|
| 117 |
+
aps = []
|
| 118 |
+
for i in range(len(V)):
|
| 119 |
+
rel = (g == g[i])
|
| 120 |
+
rel[i] = False
|
| 121 |
+
if not rel.any():
|
| 122 |
+
continue
|
| 123 |
+
order = np.argsort(-S[i])
|
| 124 |
+
hit = rel[order]
|
| 125 |
+
cum = np.cumsum(hit)
|
| 126 |
+
prec = cum / (np.arange(len(hit)) + 1)
|
| 127 |
+
aps.append(float((prec * hit).sum() / hit.sum()))
|
| 128 |
+
return round(float(np.mean(aps)), 4) if aps else None
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def fidelity(P, Y, k=10):
|
| 132 |
+
"""Layer 3. Does the student rank like the teacher?"""
|
| 133 |
+
P, Y = _l2(np.asarray(P, np.float32)), _l2(np.asarray(Y, np.float32))
|
| 134 |
+
n = min(len(P), 400)
|
| 135 |
+
Sp, St = P[:n] @ P[:n].T, Y[:n] @ Y[:n].T
|
| 136 |
+
np.fill_diagonal(Sp, -9); np.fill_diagonal(St, -9)
|
| 137 |
+
kk = min(k, n - 1)
|
| 138 |
+
rp, rt = np.argsort(-Sp, 1)[:, :kk], np.argsort(-St, 1)[:, :kk]
|
| 139 |
+
return {"cosine_to_teacher": round(float(np.mean(np.sum(P * Y, 1))), 4),
|
| 140 |
+
"nn_top1": round(float(np.mean(Sp.argmax(1) == St.argmax(1))), 4),
|
| 141 |
+
f"nn_recall@{kk}": round(float(np.mean(
|
| 142 |
+
[len(set(a) & set(b)) / kk for a, b in zip(rp, rt)])), 4)}
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def code_agreement(student, teacher, C, min_frac=0.05):
|
| 146 |
+
"""Layer 3b: does the student land in the TEACHER'S codebook cell?
|
| 147 |
+
|
| 148 |
+
This metric exists because of how this store prunes. `code` is a
|
| 149 |
+
clustered column and the planner reads only the row groups a query's
|
| 150 |
+
probe selects, so a student that produces a beautiful vector in the
|
| 151 |
+
WRONG cell puts its row in a row group the planner never opens. The
|
| 152 |
+
row is then unreachable at any k. Cosine cannot see that failure and
|
| 153 |
+
neither can nn_recall.
|
| 154 |
+
|
| 155 |
+
Two numbers, and the second is the one that decides:
|
| 156 |
+
|
| 157 |
+
code@1 student cell == teacher cell. Informative, but a
|
| 158 |
+
miss here is survivable - probe() widens.
|
| 159 |
+
code_recall the teacher's cell is INSIDE the probe set the
|
| 160 |
+
student's own vector selects. A miss here is
|
| 161 |
+
permanent data loss: nothing the reader does at
|
| 162 |
+
query time recovers that row.
|
| 163 |
+
|
| 164 |
+
A student at code@1 0.70 with code_recall 0.99 ships. One at code@1
|
| 165 |
+
0.95 with code_recall 0.95 silently loses 5% of the corpus.
|
| 166 |
+
"""
|
| 167 |
+
from .teacher import assign, probe as _probe
|
| 168 |
+
C = np.asarray(C, np.float32)
|
| 169 |
+
ts_code, _ = assign(np.asarray(teacher, np.float32), C)
|
| 170 |
+
st_code, _ = assign(np.asarray(student, np.float32), C)
|
| 171 |
+
top1 = float(np.mean(st_code == ts_code))
|
| 172 |
+
hit, widths = 0, []
|
| 173 |
+
for i, v in enumerate(_l2(np.asarray(student, np.float32))):
|
| 174 |
+
cells = _probe(v, C, min_frac=min_frac)
|
| 175 |
+
widths.append(len(cells))
|
| 176 |
+
if int(ts_code[i]) in cells:
|
| 177 |
+
hit += 1
|
| 178 |
+
return {"code@1": round(top1, 4),
|
| 179 |
+
"code_recall": round(hit / max(len(student), 1), 4),
|
| 180 |
+
"mean_probe_cells": round(float(np.mean(widths)), 2),
|
| 181 |
+
"cells": int(len(C))}
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def report(V, labels=None, groups=None, teacher=None, codebook=None):
|
| 185 |
+
"""One channel, all three layers. `labels` is {name: bool array}."""
|
| 186 |
+
out = {"separability": separability(V)}
|
| 187 |
+
if labels:
|
| 188 |
+
out["task_auc"] = {k: probe_auc(V, y) for k, y in labels.items()}
|
| 189 |
+
vals = [v for v in out["task_auc"].values() if v is not None]
|
| 190 |
+
out["task_auc_mean"] = round(float(np.mean(vals)), 4) if vals else None
|
| 191 |
+
if groups is not None:
|
| 192 |
+
out["retrieval_map"] = retrieval_map(V, groups)
|
| 193 |
+
if teacher is not None:
|
| 194 |
+
out["fidelity"] = fidelity(V, teacher)
|
| 195 |
+
if codebook is not None:
|
| 196 |
+
out["code"] = code_agreement(V, teacher, codebook)
|
| 197 |
+
return out
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
def verdict(student, teacher, tol=0.03, code_recall_min=0.98):
|
| 201 |
+
"""The shipping gate: task AUC within `tol` of the teacher's.
|
| 202 |
+
|
| 203 |
+
Deliberately NOT fidelity. A student that disagrees with its teacher
|
| 204 |
+
clip by clip while carrying the same usable signal has done its job;
|
| 205 |
+
holding it to agreement would reject it for the wrong reason.
|
| 206 |
+
"""
|
| 207 |
+
s, t = student.get("task_auc_mean"), teacher.get("task_auc_mean")
|
| 208 |
+
if s is None or t is None:
|
| 209 |
+
return {"ship": False, "reason": "no task labels to score on"}
|
| 210 |
+
if t < 0.55:
|
| 211 |
+
return {"ship": False,
|
| 212 |
+
"reason": f"teacher itself is uninformative (AUC {t}) - "
|
| 213 |
+
"fix the teacher target before distilling"}
|
| 214 |
+
cr = (student.get("code") or {}).get("code_recall")
|
| 215 |
+
if cr is not None and cr < code_recall_min:
|
| 216 |
+
return {"ship": False, "student_auc": s, "teacher_auc": t,
|
| 217 |
+
"reason": f"code_recall {cr} < {code_recall_min}: rows "
|
| 218 |
+
"whose teacher cell falls outside the student's "
|
| 219 |
+
"probe are unreachable at any k"}
|
| 220 |
+
return {"ship": bool(s >= t - tol), "student_auc": s, "teacher_auc": t,
|
| 221 |
+
"gap": round(s - t, 4), "tolerance": tol, "code_recall": cr}
|
python/elidedb/cli.py
ADDED
|
@@ -0,0 +1,293 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""elidedb — the command line. One verb per lifecycle step:
|
| 2 |
+
|
| 3 |
+
elidedb create lake/mydb --name "my dataset"
|
| 4 |
+
elidedb add lake/mydb readings data.csv --ts-col time
|
| 5 |
+
elidedb video lake/mydb cam.mp4 --stream cam0
|
| 6 |
+
elidedb ls lake/mydb
|
| 7 |
+
elidedb embed lake/mydb
|
| 8 |
+
elidedb search lake/mydb "a person crossing the street"
|
| 9 |
+
elidedb sql lake/mydb "SELECT count(*) FROM readings"
|
| 10 |
+
elidedb desk
|
| 11 |
+
|
| 12 |
+
Times print as +SECONDS relative to the store's start; pass the same +N form
|
| 13 |
+
back into `window`.
|
| 14 |
+
"""
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import argparse
|
| 18 |
+
import sys
|
| 19 |
+
|
| 20 |
+
from .store import Store
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def _rel(store, ns):
|
| 24 |
+
lo = min((d["min_ts"] for d in store.describe()
|
| 25 |
+
if d["rows"] and d["table"] != "centroids"), default=0)
|
| 26 |
+
return (ns - lo) / 1e9, lo
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _parse_t(store, s: str) -> int:
|
| 30 |
+
if s.startswith("+"):
|
| 31 |
+
lo = min((d["min_ts"] for d in store.describe()
|
| 32 |
+
if d["rows"] and d["table"] != "centroids"), default=0)
|
| 33 |
+
return lo + int(float(s[1:]) * 1e9)
|
| 34 |
+
return int(s)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def cmd_create(a):
|
| 38 |
+
Store.create(a.store, a.name or a.store.rstrip("/").split("/")[-1])
|
| 39 |
+
print(f"created database at {a.store}")
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def cmd_ls(a):
|
| 43 |
+
db = Store.open(a.store)
|
| 44 |
+
print(f"{db.name} ({db.dir})")
|
| 45 |
+
rows = db.describe()
|
| 46 |
+
if not rows:
|
| 47 |
+
print(" (empty — add data with `elidedb add` or `elidedb video`)")
|
| 48 |
+
return
|
| 49 |
+
w = max(len(r["table"]) for r in rows)
|
| 50 |
+
for r in rows:
|
| 51 |
+
span = ""
|
| 52 |
+
if r["rows"] and r["table"] != "centroids":
|
| 53 |
+
lo, base = _rel(db, r["min_ts"])
|
| 54 |
+
hi, _ = _rel(db, r["max_ts"])
|
| 55 |
+
span = f" [+{lo:.1f} .. +{hi:.1f}]s"
|
| 56 |
+
print(f" {r['table']:<{w}} {r['kind']:<12} {r['rows']:>12,} rows "
|
| 57 |
+
f"{r['bytes'] / 1e6:9.1f} MB v{r['version']}{span}")
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def cmd_add(a):
|
| 61 |
+
db = Store.open(a.store)
|
| 62 |
+
v = db.ingest_rows(a.table, a.file, ts_column=a.ts_col, ts_unit=a.ts_unit)
|
| 63 |
+
st = db.table(a.table).state()
|
| 64 |
+
print(f"{a.table}: +{st.files[-1].rows:,} rows -> version {v}")
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def cmd_video(a):
|
| 68 |
+
db = Store.open(a.store)
|
| 69 |
+
ts = None
|
| 70 |
+
if a.timestamps:
|
| 71 |
+
ts = [int(line) for line in open(a.timestamps)]
|
| 72 |
+
v = db.ingest_video(a.table, a.file, timestamps_ns=ts, stream=a.stream,
|
| 73 |
+
transcode=a.transcode, gop_s=a.gop_s, crf=a.crf)
|
| 74 |
+
print(f"{a.table}: indexed {a.file} as stream "
|
| 75 |
+
f"'{a.stream or a.file}' -> version {v}")
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def cmd_adopt(a):
|
| 79 |
+
db = Store.open(a.store)
|
| 80 |
+
for d in db.describe():
|
| 81 |
+
if d["kind"] == "frame_index":
|
| 82 |
+
r = db.adopt_media(d["table"])
|
| 83 |
+
print(f"{d['table']}: adopted {r['adopted']} media files "
|
| 84 |
+
f"({r['bytes'] / 1e6:.1f} MB) into {a.store}/media/")
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def cmd_vacuum(a):
|
| 88 |
+
db = Store.open(a.store)
|
| 89 |
+
r = db.vacuum(retain_versions=a.retain, dry_run=a.dry_run)
|
| 90 |
+
verb = "would free" if a.dry_run else "freed"
|
| 91 |
+
print(f"{r['files_removed']} files, {verb} {r['bytes_freed']/1e6:.1f} MB")
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def cmd_index(a):
|
| 95 |
+
db = Store.open(a.store)
|
| 96 |
+
if a.ann:
|
| 97 |
+
from . import ann
|
| 98 |
+
r = (ann.build_hnsw(db) if a.ann == "hnsw"
|
| 99 |
+
else ann.build_ivfpq(db))
|
| 100 |
+
print(f"{a.ann} index: {r['n']:,} vectors, "
|
| 101 |
+
f"{r['bytes']/1e6:.1f} MB (embeddings v{r['version']})")
|
| 102 |
+
else:
|
| 103 |
+
r = db.table(a.table).create_index(a.column)
|
| 104 |
+
print(f"B+ index on {a.table}.{a.column}: {r['keys']:,} keys, "
|
| 105 |
+
f"{r['bytes']/1e3:.0f} KB (v{r['version']})")
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def cmd_optimize(a):
|
| 109 |
+
db = Store.open(a.store)
|
| 110 |
+
for d in db.describe():
|
| 111 |
+
if a.table and d["table"] != a.table:
|
| 112 |
+
continue
|
| 113 |
+
r = db.table(d["table"]).compact()
|
| 114 |
+
if r.get("bytes_before"):
|
| 115 |
+
print(f" {d['table']}: {r['files_before']}->{r['files_after']} files, "
|
| 116 |
+
f"{r['bytes_before']/1e6:.1f}->{r['bytes_after']/1e6:.1f} MB "
|
| 117 |
+
f"({r['ratio']}x)")
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def cmd_embed(a):
|
| 121 |
+
db = Store.open(a.store)
|
| 122 |
+
print(db.embed_windows(window_s=a.window_s,
|
| 123 |
+
frames_per_window=a.frames_per_window))
|
| 124 |
+
if not a.no_cluster:
|
| 125 |
+
from .embeddings import cluster
|
| 126 |
+
print(cluster(db, min_cluster_size=a.min_cluster_size))
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def cmd_search(a):
|
| 130 |
+
db = Store.open(a.store)
|
| 131 |
+
hits, stats = db.search_text(a.text, k=a.k)
|
| 132 |
+
print(f"scanned {stats['scanned']}/{stats['total']} vectors, "
|
| 133 |
+
f"probed {stats['clusters_probed']}/{stats['clusters_total']} clusters")
|
| 134 |
+
for h in hits:
|
| 135 |
+
lo, _ = _rel(db, h["t0"])
|
| 136 |
+
hi, _ = _rel(db, h["t1"])
|
| 137 |
+
print(f" {h['score']:.4f} {h['stream']:<24} [+{lo:.2f} .. +{hi:.2f}]s"
|
| 138 |
+
f" (elidedb window {a.store} +{lo:.2f} +{hi:.2f})")
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def cmd_sql(a):
|
| 142 |
+
db = Store.open(a.store)
|
| 143 |
+
df = db.sql(a.query)
|
| 144 |
+
print(df.to_string(index=False, max_rows=50))
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def cmd_window(a):
|
| 148 |
+
db = Store.open(a.store)
|
| 149 |
+
t0, t1 = _parse_t(db, a.t0), _parse_t(db, a.t1)
|
| 150 |
+
w, stats = db.window(t0, t1)
|
| 151 |
+
print(stats)
|
| 152 |
+
from .video import FrameSet
|
| 153 |
+
for name, v in w.items():
|
| 154 |
+
if isinstance(v, FrameSet):
|
| 155 |
+
print(f" {name}: {len(v)} frames, streams {v.streams()}")
|
| 156 |
+
if a.dump:
|
| 157 |
+
from PIL import Image
|
| 158 |
+
import pathlib
|
| 159 |
+
out = pathlib.Path(a.dump)
|
| 160 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 161 |
+
n = 0
|
| 162 |
+
for s in v.streams():
|
| 163 |
+
for ts, img in v.decode(stream=s, width=a.width):
|
| 164 |
+
lo, _ = _rel(db, ts)
|
| 165 |
+
safe = s.replace("/", "_").replace(" ", "_")
|
| 166 |
+
Image.fromarray(img).save(out / f"{safe}_{lo:.3f}s.jpg")
|
| 167 |
+
n += 1
|
| 168 |
+
print(f" wrote {n} frames to {a.dump}")
|
| 169 |
+
elif len(v):
|
| 170 |
+
print(f" {name}: {len(v):,} rows x {len(v.column_names)} cols")
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def cmd_desk(a):
|
| 174 |
+
from . import desk
|
| 175 |
+
sys.argv = ["elidedb-desk", "--root", a.root, "--port", str(a.port)] + \
|
| 176 |
+
(["--open"] if not a.no_open else [])
|
| 177 |
+
desk.main()
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
def main():
|
| 181 |
+
ap = argparse.ArgumentParser(
|
| 182 |
+
prog="elidedb",
|
| 183 |
+
description="ElideDB — Parquet-native, timestamp-first multimodal "
|
| 184 |
+
"database. The best read is the read elided.")
|
| 185 |
+
sub = ap.add_subparsers(dest="cmd", required=True)
|
| 186 |
+
|
| 187 |
+
p = sub.add_parser("create", help="create an empty database")
|
| 188 |
+
p.add_argument("store")
|
| 189 |
+
p.add_argument("--name")
|
| 190 |
+
p.set_defaults(f=cmd_create)
|
| 191 |
+
|
| 192 |
+
p = sub.add_parser("ls", help="list a database's tables")
|
| 193 |
+
p.add_argument("store")
|
| 194 |
+
p.set_defaults(f=cmd_ls)
|
| 195 |
+
|
| 196 |
+
p = sub.add_parser("add", help="add timestamped rows (CSV/Parquet)")
|
| 197 |
+
p.add_argument("store")
|
| 198 |
+
p.add_argument("table")
|
| 199 |
+
p.add_argument("file")
|
| 200 |
+
p.add_argument("--ts-col", default="ts",
|
| 201 |
+
help="timestamp column name (default: ts)")
|
| 202 |
+
p.add_argument("--ts-unit", default="auto",
|
| 203 |
+
choices=["auto", "s", "ms", "us", "ns"],
|
| 204 |
+
help="timestamp unit (default: auto-detect)")
|
| 205 |
+
p.set_defaults(f=cmd_add)
|
| 206 |
+
|
| 207 |
+
p = sub.add_parser("video", help="index a video file (media stays put)")
|
| 208 |
+
p.add_argument("store")
|
| 209 |
+
p.add_argument("file")
|
| 210 |
+
p.add_argument("--table", default="frames")
|
| 211 |
+
p.add_argument("--stream", help="stream name (default: file stem)")
|
| 212 |
+
p.add_argument("--timestamps",
|
| 213 |
+
help="file with one ns timestamp per frame; omitted = "
|
| 214 |
+
"container timestamps")
|
| 215 |
+
p.add_argument("--transcode", choices=["hevc", "h264"],
|
| 216 |
+
help="re-encode the managed copy (~10-25x smaller; "
|
| 217 |
+
"random access becomes GOP-granular)")
|
| 218 |
+
p.add_argument("--gop-s", type=float, default=1.0,
|
| 219 |
+
help="keyframe interval in seconds (the seekability-vs-"
|
| 220 |
+
"compression dial)")
|
| 221 |
+
p.add_argument("--crf", type=int, default=26)
|
| 222 |
+
p.set_defaults(f=cmd_video)
|
| 223 |
+
|
| 224 |
+
p = sub.add_parser("adopt", help="copy referenced media into the store "
|
| 225 |
+
"(makes it standalone)")
|
| 226 |
+
p.add_argument("store")
|
| 227 |
+
p.set_defaults(f=cmd_adopt)
|
| 228 |
+
|
| 229 |
+
p = sub.add_parser("optimize", help="compact tables (fewer, larger, "
|
| 230 |
+
"delta-encoded files)")
|
| 231 |
+
p.add_argument("store")
|
| 232 |
+
p.add_argument("--table")
|
| 233 |
+
p.set_defaults(f=cmd_optimize)
|
| 234 |
+
|
| 235 |
+
p = sub.add_parser("vacuum", help="garbage-collect unreachable files "
|
| 236 |
+
"(keeps last N versions readable)")
|
| 237 |
+
p.add_argument("store")
|
| 238 |
+
p.add_argument("--retain", type=int, default=3)
|
| 239 |
+
p.add_argument("--dry-run", action="store_true")
|
| 240 |
+
p.set_defaults(f=cmd_vacuum)
|
| 241 |
+
|
| 242 |
+
p = sub.add_parser("index", help="build a B+ tree on a column, or an "
|
| 243 |
+
"ANN index (hnsw/ivfpq) on embeddings")
|
| 244 |
+
p.add_argument("store")
|
| 245 |
+
p.add_argument("--table", default="frames")
|
| 246 |
+
p.add_argument("--column", help="numeric column for a B+ secondary index")
|
| 247 |
+
p.add_argument("--ann", choices=["hnsw", "ivfpq"],
|
| 248 |
+
help="build a vector index over the embeddings table")
|
| 249 |
+
p.set_defaults(f=cmd_index)
|
| 250 |
+
|
| 251 |
+
p = sub.add_parser("embed", help="embed video windows + cluster (local ML)")
|
| 252 |
+
p.add_argument("store")
|
| 253 |
+
p.add_argument("--window-s", type=float, default=2.0)
|
| 254 |
+
p.add_argument("--frames-per-window", type=int, default=2)
|
| 255 |
+
p.add_argument("--min-cluster-size", type=int, default=8)
|
| 256 |
+
p.add_argument("--no-cluster", action="store_true")
|
| 257 |
+
p.set_defaults(f=cmd_embed)
|
| 258 |
+
|
| 259 |
+
p = sub.add_parser("search", help="semantic text search")
|
| 260 |
+
p.add_argument("store")
|
| 261 |
+
p.add_argument("text")
|
| 262 |
+
p.add_argument("-k", type=int, default=8)
|
| 263 |
+
p.set_defaults(f=cmd_search)
|
| 264 |
+
|
| 265 |
+
p = sub.add_parser("sql", help="run SQL (DuckDB) over the database")
|
| 266 |
+
p.add_argument("store")
|
| 267 |
+
p.add_argument("query")
|
| 268 |
+
p.set_defaults(f=cmd_sql)
|
| 269 |
+
|
| 270 |
+
p = sub.add_parser("window", help="read a time window (+SECS or raw ns)")
|
| 271 |
+
p.add_argument("store")
|
| 272 |
+
p.add_argument("t0")
|
| 273 |
+
p.add_argument("t1")
|
| 274 |
+
p.add_argument("--dump", help="write frames as JPEGs to this directory")
|
| 275 |
+
p.add_argument("--width", type=int, default=640)
|
| 276 |
+
p.set_defaults(f=cmd_window)
|
| 277 |
+
|
| 278 |
+
p = sub.add_parser("desk", help="open the browser UI")
|
| 279 |
+
p.add_argument("--root", default="lake")
|
| 280 |
+
p.add_argument("--port", type=int, default=8787)
|
| 281 |
+
p.add_argument("--no-open", action="store_true")
|
| 282 |
+
p.set_defaults(f=cmd_desk)
|
| 283 |
+
|
| 284 |
+
a = ap.parse_args()
|
| 285 |
+
try:
|
| 286 |
+
a.f(a)
|
| 287 |
+
except Exception as e:
|
| 288 |
+
print(f"error: {e}", file=sys.stderr)
|
| 289 |
+
sys.exit(1)
|
| 290 |
+
|
| 291 |
+
|
| 292 |
+
if __name__ == "__main__":
|
| 293 |
+
main()
|
python/elidedb/context.py
ADDED
|
@@ -0,0 +1,1112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Context retrieval: the ingest-side half.
|
| 2 |
+
|
| 3 |
+
WHY THIS EXISTS
|
| 4 |
+
---------------
|
| 5 |
+
`embeddings` holds one mean-pooled SigLIP vector per window. Mean pooling is
|
| 6 |
+
order-blind: reverse the frames and the vector is identical. So the index can
|
| 7 |
+
answer "is there a car and a person here" but never "is the person walking
|
| 8 |
+
*toward* the car" — appearance, not context. Reranking with a VLM fixes the
|
| 9 |
+
ranking but costs seconds per query, which is the wrong place to spend time in
|
| 10 |
+
a database.
|
| 11 |
+
|
| 12 |
+
The fix is the oldest trick a database has: precompute the expensive operator
|
| 13 |
+
into an index and make the query a lookup. Concretely,
|
| 14 |
+
|
| 15 |
+
ingest (once, offline) query (every time, hot)
|
| 16 |
+
------------------------ -----------------------
|
| 17 |
+
VLM reads 3 frames of the text -> SigLIP text vector
|
| 18 |
+
window and writes a
|
| 19 |
+
relational caption ONE matmul against ctx vectors
|
| 20 |
+
|
|
| 21 |
+
SigLIP text tower (no VLM in the loop, ever)
|
| 22 |
+
|
|
| 23 |
+
target vector -----------> train a small temporal tower to
|
| 24 |
+
predict it from cheap frame vectors
|
| 25 |
+
|
| 26 |
+
The VLM only ever labels a subset. The tower generalises the label to every
|
| 27 |
+
window, including windows ingested later, so a new day of footage costs a
|
| 28 |
+
forward pass over frame vectors instead of a day of VLM time. This is
|
| 29 |
+
pseudo-labelling in the sense of "Distilling Vision-Language Models on
|
| 30 |
+
Millions of Videos" (arXiv 2401.06129), applied at ingest instead of at
|
| 31 |
+
pretraining scale.
|
| 32 |
+
|
| 33 |
+
Three tables come out of this module, all ordinary Parquet:
|
| 34 |
+
|
| 35 |
+
frame_vectors ts, stream, vector[d] one row per FRAME
|
| 36 |
+
context_captions ts, t1, stream, caption, vector[d] teacher labels
|
| 37 |
+
context ts, t1, stream, vector[d] student output
|
| 38 |
+
"""
|
| 39 |
+
from __future__ import annotations
|
| 40 |
+
|
| 41 |
+
import re
|
| 42 |
+
import time
|
| 43 |
+
|
| 44 |
+
import numpy as np
|
| 45 |
+
import pyarrow as pa
|
| 46 |
+
import pyarrow.compute as pc
|
| 47 |
+
|
| 48 |
+
from .embeddings import (DEFAULT_MODEL, MODELS, _embed_images,
|
| 49 |
+
_load_model, resolve_model)
|
| 50 |
+
|
| 51 |
+
# The teacher is asked for RELATIONS and CHANGE, not for a list of objects.
|
| 52 |
+
# An object list is exactly what SigLIP already encodes, so a caption that
|
| 53 |
+
# reads "a car, a road, a building" teaches the student nothing it does not
|
| 54 |
+
# already know. Everything interesting is in the verbs.
|
| 55 |
+
CAPTION_PROMPT = (
|
| 56 |
+
"These frames are in time order from one short video clip. "
|
| 57 |
+
"Reply with ONE sentence of at most 25 words describing what is "
|
| 58 |
+
"happening: who or what is present, what each is doing, how they are "
|
| 59 |
+
"positioned relative to each other, and how the scene moves or changes. "
|
| 60 |
+
"Be concrete and literal. Do not say 'frame', 'image', or 'video'."
|
| 61 |
+
)
|
| 62 |
+
|
| 63 |
+
# The caption is the index. If it does not use the words a user would use, the
|
| 64 |
+
# lexical ranker can never fire and the caption-LSA space is built around the
|
| 65 |
+
# wrong distinctions. Measured on BridgeData2: the generic prompt above
|
| 66 |
+
# produced "a robot arm interacts with a wooden box", while the human label
|
| 67 |
+
# for the same clip was "put red object in the drawer" — different noun,
|
| 68 |
+
# different granularity, no overlap for retrieval to work with.
|
| 69 |
+
#
|
| 70 |
+
# So the prompt is a per-domain parameter. This one asks for the ACTION and
|
| 71 |
+
# the OBJECT MOVED, which is how manipulation data is described. It stays
|
| 72 |
+
# deliberately generic: it never names objects that appear in the labels
|
| 73 |
+
# (that would be leaking the eval set into the index), only the SHAPE of the
|
| 74 |
+
# description — what moved, and where it ended up.
|
| 75 |
+
# The caption IS the index: whatever verbs the prompt teaches are the only
|
| 76 |
+
# verbs lexical recall can ever match. The first version of this prompt said
|
| 77 |
+
# "which object the arm picks up or moves, and where it puts it" — and the
|
| 78 |
+
# resulting 2,348 captions contained picks x2371, puts x1673 and ZERO
|
| 79 |
+
# instances of close/open/wipe/push, so "close the drawer" was unfindable by
|
| 80 |
+
# construction. The prompt must be VERB-OPEN: describe the action in its own
|
| 81 |
+
# words, and always report state changes.
|
| 82 |
+
# S0 VIOLATION, FIXED. The previous MANIPULATION_PROMPT enumerated the
|
| 83 |
+
# task vocabulary - "picking up, putting, opening, closing, pushing,
|
| 84 |
+
# pouring, wiping, pressing" - and this file's own comment stated the
|
| 85 |
+
# consequence: "The caption IS the index: whatever verbs the prompt
|
| 86 |
+
# teaches are the only verbs lexical recall can ever match." A verb
|
| 87 |
+
# outside that list was unfindable by construction, and a non-
|
| 88 |
+
# manipulation corpus was unfindable entirely.
|
| 89 |
+
#
|
| 90 |
+
# It was also TUNED AGAINST THE EVAL LABELS. The removed comment
|
| 91 |
+
# recorded the generic prompt producing "a robot arm interacts with a
|
| 92 |
+
# wooden box" where "the human label for the same clip was 'put red
|
| 93 |
+
# object in the drawer'", and the prompt being rewritten to close that
|
| 94 |
+
# gap. bridge_ingest.py deliberately keeps the task strings out of the
|
| 95 |
+
# store; this prompt let them back in through the side door.
|
| 96 |
+
#
|
| 97 |
+
# The replacement names no verb, no object class and no domain. It asks
|
| 98 |
+
# for what MOVED and what CHANGED, which is answerable for a kitchen, a
|
| 99 |
+
# road, a warehouse or a surgical table, and lets the corpus supply its
|
| 100 |
+
# own words.
|
| 101 |
+
MANIPULATION_PROMPT = (
|
| 102 |
+
"These frames are in time order from one short clip. "
|
| 103 |
+
"Reply with ONE short sentence in plain English: what moved, what it "
|
| 104 |
+
"did, and what was different at the end. Use whatever words fit; do "
|
| 105 |
+
"not choose from a list. If anything changed state, say so. "
|
| 106 |
+
"Do not say 'frame', 'image', or 'video'."
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
PROMPTS = {"scene": CAPTION_PROMPT, "manipulation": MANIPULATION_PROMPT}
|
| 110 |
+
|
| 111 |
+
# SigLIP's text tower truncates at 64 tokens, so a rambling caption is
|
| 112 |
+
# silently cut mid-clause and the tail is lost anyway. Trim deliberately
|
| 113 |
+
# instead: drop the VLM's framing preamble, keep whole sentences.
|
| 114 |
+
#
|
| 115 |
+
# The preamble pattern is deliberately narrow — a leading PREPOSITIONAL
|
| 116 |
+
# phrase only ("In the first frame,", "Across this sequence,"). An earlier
|
| 117 |
+
# looser version also matched "The video shows a street with cars," and
|
| 118 |
+
# amputated the subject, leaving captions that began "and various buildings".
|
| 119 |
+
# Requiring a leading preposition and at most two filler words makes that
|
| 120 |
+
# impossible.
|
| 121 |
+
_PREAMBLE = re.compile(
|
| 122 |
+
r"^(?:in|across|throughout|over|during)\s+(?:the|this|these)\s+"
|
| 123 |
+
r"(?:[\w-]+\s+){0,2}?(?:frames?|images?|pictures?|sequence|clip|video)\s*,\s*",
|
| 124 |
+
re.I)
|
| 125 |
+
|
| 126 |
+
# Prompt-echo: asking for "the action verb in plain English" made the VLM
|
| 127 |
+
# write 'The action verb is "picking up" as the arm picks up...' — meta-
|
| 128 |
+
# language that pollutes the lexical index. Strip the frame, keep the deed.
|
| 129 |
+
_VERB_ECHO = re.compile(
|
| 130 |
+
r"^the action(?:\s+verb)?\s+is\s+['\"]?[\w-]+(?:\s+[\w-]+){0,2}?"
|
| 131 |
+
r"['\"]?[,.]?\s+(?:as|because|where|since|:)\s+", re.I)
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
def tidy_caption(text: str, max_words: int = 32) -> str:
|
| 135 |
+
t = " ".join(text.strip().split())
|
| 136 |
+
t = _PREAMBLE.sub("", t)
|
| 137 |
+
t = _VERB_ECHO.sub("", t)
|
| 138 |
+
parts = re.split(r"(?<=[.!?])\s+", t)
|
| 139 |
+
# A generation cut off at max_tokens ends mid-clause. Keep only sentences
|
| 140 |
+
# that actually terminate, unless that would leave nothing at all.
|
| 141 |
+
whole = [p for p in parts if p.rstrip().endswith((".", "!", "?"))]
|
| 142 |
+
parts = whole or parts[:1]
|
| 143 |
+
out = []
|
| 144 |
+
for p in parts:
|
| 145 |
+
if out and len(" ".join(out + [p]).split()) > max_words:
|
| 146 |
+
break
|
| 147 |
+
out.append(p)
|
| 148 |
+
t = " ".join(out).strip()
|
| 149 |
+
return (t[:1].upper() + t[1:]) if t else " ".join(text.split())[:200]
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
# ---------------------------------------------------------------------------
|
| 153 |
+
# 1. Per-frame vectors — the sequence a temporal model needs
|
| 154 |
+
# ---------------------------------------------------------------------------
|
| 155 |
+
def embed_frames(store, frame_table="frames", model=None, width=512,
|
| 156 |
+
batch=32, incremental=True, verbose=True, stride=1,
|
| 157 |
+
streams=None, engine="siglip"):
|
| 158 |
+
"""One SigLIP vector per frame → `frame_vectors`.
|
| 159 |
+
|
| 160 |
+
Deliberately NOT pooled. Pooling is the student's job, and pooling here
|
| 161 |
+
would throw away the only signal that distinguishes context from
|
| 162 |
+
appearance. Decode goes through the same byte-range path queries use, so
|
| 163 |
+
this costs the frames it reads and nothing else.
|
| 164 |
+
"""
|
| 165 |
+
from PIL import Image
|
| 166 |
+
|
| 167 |
+
from .video import FrameSet
|
| 168 |
+
# engine="fdnnv": the distilled streaming encoder — every frame, no
|
| 169 |
+
# stride, state carried per stream. Same output space as the teacher, so
|
| 170 |
+
# everything downstream (pooling, search, context) is unchanged.
|
| 171 |
+
if engine == "fdnnv":
|
| 172 |
+
return _embed_frames_fdnnv(store, frame_table, incremental, verbose,
|
| 173 |
+
streams)
|
| 174 |
+
model = resolve_model(model)
|
| 175 |
+
frames = store.table(frame_table).scan()
|
| 176 |
+
allst = sorted(set(frames.column("stream").to_pylist()))
|
| 177 |
+
streams = [s for s in allst if s in streams] if streams else allst
|
| 178 |
+
|
| 179 |
+
done = {}
|
| 180 |
+
if incremental:
|
| 181 |
+
try:
|
| 182 |
+
prev = store.table("frame_vectors").scan()
|
| 183 |
+
for s_, t_ in zip(prev.column("stream").to_pylist(),
|
| 184 |
+
prev.column("ts").to_pylist()):
|
| 185 |
+
done[s_] = max(done.get(s_, -1), t_)
|
| 186 |
+
except Exception:
|
| 187 |
+
pass
|
| 188 |
+
|
| 189 |
+
t_start = time.time()
|
| 190 |
+
rows_ts, rows_stream, rows_vec = [], [], []
|
| 191 |
+
for s in streams:
|
| 192 |
+
sel = frames.filter(pc.equal(frames.column("stream"), s))
|
| 193 |
+
if s in done:
|
| 194 |
+
sel = sel.filter(pc.greater(sel.column("ts"), done[s]))
|
| 195 |
+
if len(sel) == 0:
|
| 196 |
+
continue
|
| 197 |
+
# `stride` subsamples the frame index before decoding. A 5 Hz robot
|
| 198 |
+
# camera does not need every frame embedded for a 4 s window to be
|
| 199 |
+
# well described, and the cost here is linear in frames decoded, so
|
| 200 |
+
# this is the dial between ingest time and temporal resolution.
|
| 201 |
+
decoded = FrameSet(store, frame_table, sel).decode(width=width,
|
| 202 |
+
stride=stride)
|
| 203 |
+
if verbose:
|
| 204 |
+
print(f" {s}: {len(decoded)} frames decoded", flush=True)
|
| 205 |
+
for i in range(0, len(decoded), batch):
|
| 206 |
+
chunk = decoded[i:i + batch]
|
| 207 |
+
vecs = _embed_images([Image.fromarray(a) for _, a in chunk], model)
|
| 208 |
+
for (ts, _), v in zip(chunk, vecs):
|
| 209 |
+
rows_ts.append(int(ts))
|
| 210 |
+
rows_stream.append(s)
|
| 211 |
+
rows_vec.append(v)
|
| 212 |
+
if not rows_ts:
|
| 213 |
+
return {"frames": 0, "note": "nothing new (incremental)"}
|
| 214 |
+
|
| 215 |
+
dim = len(rows_vec[0])
|
| 216 |
+
tbl = pa.table({
|
| 217 |
+
"ts": pa.array(rows_ts, pa.int64()),
|
| 218 |
+
"stream": pa.array(rows_stream),
|
| 219 |
+
"vector": pa.array([v.tolist() for v in rows_vec],
|
| 220 |
+
pa.list_(pa.float32(), dim)),
|
| 221 |
+
})
|
| 222 |
+
version = store.table("frame_vectors").append(
|
| 223 |
+
tbl, kind="embeddings",
|
| 224 |
+
meta={"model": model, "dim": dim, "decode_width": width,
|
| 225 |
+
"source_table": frame_table})
|
| 226 |
+
return {"frames": len(tbl), "dim": dim, "version": version,
|
| 227 |
+
"seconds": round(time.time() - t_start, 1)}
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
# ---------------------------------------------------------------------------
|
| 231 |
+
# 2. Window plan — shared by teacher and student so labels line up exactly
|
| 232 |
+
# ---------------------------------------------------------------------------
|
| 233 |
+
def plan_windows(store, window_s=2.0, stride_s=0.5, table="frame_vectors",
|
| 234 |
+
min_frames=4):
|
| 235 |
+
"""Sliding windows over each stream's timeline.
|
| 236 |
+
|
| 237 |
+
Stride < window on purpose: overlapping windows are how a *sliding* index
|
| 238 |
+
avoids the boundary problem where an event straddles two tumbling windows
|
| 239 |
+
and lands strongly in neither. Merging overlaps back into one answer is
|
| 240 |
+
already handled downstream by the segment merger.
|
| 241 |
+
"""
|
| 242 |
+
fv = store.table(table).scan()
|
| 243 |
+
win = int(window_s * 1e9)
|
| 244 |
+
stride = int(stride_s * 1e9)
|
| 245 |
+
out = []
|
| 246 |
+
for s in sorted(set(fv.column("stream").to_pylist())):
|
| 247 |
+
rows = fv.filter(pc.equal(fv.column("stream"), s))
|
| 248 |
+
ts = np.sort(rows.column("ts").to_numpy())
|
| 249 |
+
if len(ts) == 0:
|
| 250 |
+
continue
|
| 251 |
+
t = int(ts[0])
|
| 252 |
+
end = int(ts[-1])
|
| 253 |
+
while t <= end - win // 2:
|
| 254 |
+
lo, hi = np.searchsorted(ts, [t, t + win])
|
| 255 |
+
if hi - lo >= min_frames:
|
| 256 |
+
out.append((s, t, t + win - 1))
|
| 257 |
+
t += stride
|
| 258 |
+
return out
|
| 259 |
+
|
| 260 |
+
|
| 261 |
+
def window_sequences(store, windows, table="frame_vectors", max_len=32):
|
| 262 |
+
"""(stream, t0, t1) → (T, d) float32 stack of that window's frame vectors.
|
| 263 |
+
|
| 264 |
+
Subsampled to `max_len` evenly. A 2 s window at 16 Hz is 32 frames; the
|
| 265 |
+
cap keeps the tower's cost independent of frame rate, which is what makes
|
| 266 |
+
the same model valid across a 10 Hz LiDAR-synced camera and a 60 Hz one.
|
| 267 |
+
"""
|
| 268 |
+
fv = store.table(table).scan()
|
| 269 |
+
by_stream = {}
|
| 270 |
+
for s in sorted(set(fv.column("stream").to_pylist())):
|
| 271 |
+
rows = fv.filter(pc.equal(fv.column("stream"), s))
|
| 272 |
+
ts = rows.column("ts").to_numpy()
|
| 273 |
+
order = np.argsort(ts)
|
| 274 |
+
vecs = np.asarray(rows.column("vector").to_pylist(), dtype=np.float32)
|
| 275 |
+
by_stream[s] = (ts[order], vecs[order])
|
| 276 |
+
seqs = []
|
| 277 |
+
for (s, t0, t1) in windows:
|
| 278 |
+
ts, vecs = by_stream[s]
|
| 279 |
+
lo, hi = np.searchsorted(ts, [t0, t1 + 1])
|
| 280 |
+
idx = np.arange(lo, hi)
|
| 281 |
+
if len(idx) > max_len:
|
| 282 |
+
idx = idx[np.linspace(0, len(idx) - 1, max_len).round().astype(int)]
|
| 283 |
+
seqs.append(vecs[idx])
|
| 284 |
+
return seqs
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
# ---------------------------------------------------------------------------
|
| 288 |
+
# 3. The teacher — a VLM that actually reads the pixels, run ONCE per window
|
| 289 |
+
# ---------------------------------------------------------------------------
|
| 290 |
+
def caption_windows(store, windows, frames_per_window=3, model_id=None,
|
| 291 |
+
max_tokens=64, width=448, verbose=True, limit=None,
|
| 292 |
+
every=1, prompt="scene"):
|
| 293 |
+
"""VLM captions for `windows` → `context_captions` table.
|
| 294 |
+
|
| 295 |
+
The VLM is shown several frames of the SAME window in order, so the
|
| 296 |
+
caption can describe motion. A single-frame caption would be another
|
| 297 |
+
appearance label and the student would learn nothing a mean-pool cannot
|
| 298 |
+
already produce.
|
| 299 |
+
"""
|
| 300 |
+
import tempfile
|
| 301 |
+
from pathlib import Path
|
| 302 |
+
|
| 303 |
+
from PIL import Image
|
| 304 |
+
|
| 305 |
+
from .rerank import DEFAULT_VLM, _load
|
| 306 |
+
from .video import FrameSet
|
| 307 |
+
model_id = model_id or DEFAULT_VLM
|
| 308 |
+
from mlx_vlm import generate
|
| 309 |
+
from mlx_vlm.prompt_utils import apply_chat_template
|
| 310 |
+
|
| 311 |
+
vlm, processor, cfg, _, _ = _load(model_id)
|
| 312 |
+
rot = store.meta.get("display", {}).get("rotate", 0)
|
| 313 |
+
frames = store.table("frames").scan()
|
| 314 |
+
tmpdir = Path(tempfile.mkdtemp(prefix="elidedb_ctx_"))
|
| 315 |
+
|
| 316 |
+
# `every` samples the window list uniformly instead of taking a prefix.
|
| 317 |
+
# On a corpus too large to caption in full this matters: a prefix would
|
| 318 |
+
# confine the caption vocabulary to whatever happens early in the
|
| 319 |
+
# timeline, so query terms for anything later would be out of vocabulary
|
| 320 |
+
# and the lexical ranker would abstain on the entire tail.
|
| 321 |
+
if every > 1:
|
| 322 |
+
windows = windows[::every]
|
| 323 |
+
if limit:
|
| 324 |
+
windows = windows[:limit]
|
| 325 |
+
text = PROMPTS.get(prompt, prompt) # preset name or raw prompt
|
| 326 |
+
prompt = apply_chat_template(processor, cfg, text,
|
| 327 |
+
num_images=frames_per_window)
|
| 328 |
+
|
| 329 |
+
recs = {"ts": [], "t1": [], "stream": [], "caption": []}
|
| 330 |
+
t_start = time.time()
|
| 331 |
+
for n, (s, t0, t1) in enumerate(windows):
|
| 332 |
+
sel = frames.filter(pc.and_(
|
| 333 |
+
pc.equal(frames.column("stream"), s),
|
| 334 |
+
pc.and_(pc.greater_equal(frames.column("ts"), t0),
|
| 335 |
+
pc.less_equal(frames.column("ts"), t1))))
|
| 336 |
+
fs = FrameSet(store, "frames", sel)
|
| 337 |
+
dec = fs.decode(width=width)
|
| 338 |
+
if len(dec) < 1:
|
| 339 |
+
continue
|
| 340 |
+
picks = np.linspace(0, len(dec) - 1,
|
| 341 |
+
min(frames_per_window, len(dec))).round().astype(int)
|
| 342 |
+
paths = []
|
| 343 |
+
for j, p in enumerate(picks):
|
| 344 |
+
im = Image.fromarray(dec[p][1])
|
| 345 |
+
if rot:
|
| 346 |
+
im = im.rotate(rot, expand=True)
|
| 347 |
+
fp = tmpdir / f"w{n}_{j}.jpg"
|
| 348 |
+
im.save(fp, "JPEG", quality=85)
|
| 349 |
+
paths.append(str(fp))
|
| 350 |
+
while len(paths) < frames_per_window: # pad short windows
|
| 351 |
+
paths.append(paths[-1])
|
| 352 |
+
r = generate(vlm, processor, prompt, image=paths,
|
| 353 |
+
max_tokens=max_tokens, verbose=False)
|
| 354 |
+
cap = tidy_caption(r.text if hasattr(r, "text") else str(r))
|
| 355 |
+
recs["ts"].append(t0)
|
| 356 |
+
recs["t1"].append(t1)
|
| 357 |
+
recs["stream"].append(s)
|
| 358 |
+
recs["caption"].append(cap)
|
| 359 |
+
if verbose and n % 10 == 0:
|
| 360 |
+
el = time.time() - t_start
|
| 361 |
+
print(f" [{n + 1}/{len(windows)}] {el:.0f}s {s} +"
|
| 362 |
+
f"{(t0 - int(frames.column('ts')[0].as_py())) / 1e9:.1f}s :: "
|
| 363 |
+
f"{cap[:90]}", flush=True)
|
| 364 |
+
if not recs["ts"]:
|
| 365 |
+
return {"captions": 0}
|
| 366 |
+
|
| 367 |
+
# Caption → frozen SigLIP TEXT tower. This is the whole point of using
|
| 368 |
+
# SigLIP as the teacher's codec: the target lands in the SAME space a
|
| 369 |
+
# user's query lands in, so the student is learning to be the image side
|
| 370 |
+
# of a dual encoder — not to regress an arbitrary embedding.
|
| 371 |
+
vecs = embed_texts(recs["caption"])
|
| 372 |
+
dim = vecs.shape[1]
|
| 373 |
+
tbl = pa.table({
|
| 374 |
+
"ts": pa.array(recs["ts"], pa.int64()),
|
| 375 |
+
"t1": pa.array(recs["t1"], pa.int64()),
|
| 376 |
+
"stream": pa.array(recs["stream"]),
|
| 377 |
+
"caption": pa.array(recs["caption"]),
|
| 378 |
+
"vector": pa.array([v.tolist() for v in vecs],
|
| 379 |
+
pa.list_(pa.float32(), dim)),
|
| 380 |
+
})
|
| 381 |
+
version = store.table("context_captions").append(
|
| 382 |
+
tbl, kind="embeddings",
|
| 383 |
+
meta={"teacher": model_id, "text_model": DEFAULT_MODEL, "dim": dim,
|
| 384 |
+
"frames_per_window": frames_per_window, "prompt": text[:200],
|
| 385 |
+
"seconds": round(time.time() - t_start, 1)})
|
| 386 |
+
return {"captions": len(tbl), "version": version,
|
| 387 |
+
"seconds": round(time.time() - t_start, 1)}
|
| 388 |
+
|
| 389 |
+
|
| 390 |
+
def embed_texts(texts, model_id=DEFAULT_MODEL, batch=32):
|
| 391 |
+
"""Batched SigLIP text tower. Same normalisation as image vectors so the
|
| 392 |
+
two are directly comparable by dot product. Delegates to the
|
| 393 |
+
portable single-text path when mlx is not on this machine."""
|
| 394 |
+
from .embeddings import _backend, embed_text
|
| 395 |
+
if _backend() != "mlx":
|
| 396 |
+
return np.stack([embed_text(t if t.strip() else "a scene",
|
| 397 |
+
model_id) for t in texts])
|
| 398 |
+
import mlx.core as mx
|
| 399 |
+
model, processor = _load_model(model_id)
|
| 400 |
+
out = []
|
| 401 |
+
for i in range(0, len(texts), batch):
|
| 402 |
+
chunk = [t if t.strip() else "a scene" for t in texts[i:i + batch]]
|
| 403 |
+
ti = processor(text=chunk, padding="max_length", max_length=64,
|
| 404 |
+
truncation=True, return_tensors="np")
|
| 405 |
+
v = np.array(model.get_text_features(mx.array(ti["input_ids"])),
|
| 406 |
+
dtype=np.float32)
|
| 407 |
+
out.append(v / np.linalg.norm(v, axis=1, keepdims=True))
|
| 408 |
+
return np.concatenate(out, axis=0)
|
| 409 |
+
|
| 410 |
+
|
| 411 |
+
# ---------------------------------------------------------------------------
|
| 412 |
+
# 4. CaptionSpace — the output space, chosen by measurement
|
| 413 |
+
# ---------------------------------------------------------------------------
|
| 414 |
+
# Three candidate spaces were benchmarked against a VLM judge on windows the
|
| 415 |
+
# tower never trained on (mean yes/no logprob margin over each method's top-5,
|
| 416 |
+
# six relational queries, higher is better):
|
| 417 |
+
#
|
| 418 |
+
# appearance only (SigLIP image-text) +0.276
|
| 419 |
+
# caption EMBEDDING (SigLIP text tower, oracle) +0.218 <- worse
|
| 420 |
+
# VLM rerank at query time (2.1 s/query) +0.314
|
| 421 |
+
# caption TEXT, LSA-48 +0.339 <- winner
|
| 422 |
+
# caption TEXT, raw TF-IDF, fused with appearance +0.345
|
| 423 |
+
#
|
| 424 |
+
# The embedding route loses because SigLIP's text tower is trained to sit
|
| 425 |
+
# near IMAGES, not near other text; comparing a query embedding to a caption
|
| 426 |
+
# embedding uses a geometry the model was never optimised for. Matching the
|
| 427 |
+
# caption as TEXT sidesteps that entirely.
|
| 428 |
+
#
|
| 429 |
+
# LSA-48 is chosen over raw TF-IDF despite scoring 0.006 lower: it is dense
|
| 430 |
+
# and fixed-width, so (a) it is a target a small tower can actually regress,
|
| 431 |
+
# which is what lets unlabelled windows get a context vector at all, and
|
| 432 |
+
# (b) it is one more fixed_size_list column, so every index already in the
|
| 433 |
+
# store consumes it unchanged.
|
| 434 |
+
|
| 435 |
+
|
| 436 |
+
class CaptionSpace:
|
| 437 |
+
"""TF-IDF + LSA over the caption corpus. Queries and captions share it.
|
| 438 |
+
|
| 439 |
+
Not persisted as a pickle. The captions themselves are the durable
|
| 440 |
+
artifact — they live in `context_captions` as ordinary Parquet — and the
|
| 441 |
+
lexical index is refit from them on load (a few ms for this corpus) and
|
| 442 |
+
cached. That keeps the index unconditionally consistent with the data and
|
| 443 |
+
free of any sklearn version pinning. At corpus sizes where refitting
|
| 444 |
+
stops being free, persist the vocabulary and the SVD basis; the interface
|
| 445 |
+
does not change.
|
| 446 |
+
"""
|
| 447 |
+
|
| 448 |
+
def __init__(self, vec, svd):
|
| 449 |
+
self.vec, self.svd = vec, svd
|
| 450 |
+
self.dim = svd.n_components
|
| 451 |
+
|
| 452 |
+
@staticmethod
|
| 453 |
+
def fit(texts, dim=48):
|
| 454 |
+
from sklearn.decomposition import TruncatedSVD
|
| 455 |
+
from sklearn.feature_extraction.text import TfidfVectorizer
|
| 456 |
+
vec = TfidfVectorizer(stop_words="english", ngram_range=(1, 2),
|
| 457 |
+
sublinear_tf=True).fit(texts)
|
| 458 |
+
X = vec.transform(texts)
|
| 459 |
+
dim = int(min(dim, X.shape[1] - 1, len(texts) - 1))
|
| 460 |
+
svd = TruncatedSVD(n_components=dim, random_state=0).fit(X)
|
| 461 |
+
return CaptionSpace(vec, svd)
|
| 462 |
+
|
| 463 |
+
def transform(self, texts):
|
| 464 |
+
v = self.svd.transform(self.vec.transform(list(texts)))
|
| 465 |
+
v = np.asarray(v, dtype=np.float32)
|
| 466 |
+
return v / (np.linalg.norm(v, axis=1, keepdims=True) + 1e-8)
|
| 467 |
+
|
| 468 |
+
|
| 469 |
+
_SPACE_CACHE: dict = {}
|
| 470 |
+
|
| 471 |
+
|
| 472 |
+
def caption_space(store, dim=48):
|
| 473 |
+
t = store.table("context_captions")
|
| 474 |
+
key = (str(store.dir), t.state().version, dim)
|
| 475 |
+
if key not in _SPACE_CACHE:
|
| 476 |
+
_SPACE_CACHE.clear()
|
| 477 |
+
_SPACE_CACHE[key] = CaptionSpace.fit(
|
| 478 |
+
t.scan().column("caption").to_pylist(), dim=dim)
|
| 479 |
+
return _SPACE_CACHE[key]
|
| 480 |
+
|
| 481 |
+
|
| 482 |
+
# ---------------------------------------------------------------------------
|
| 483 |
+
# 5. The student — train the tower, then materialise context vectors
|
| 484 |
+
# ---------------------------------------------------------------------------
|
| 485 |
+
def _labelled(store, windows):
|
| 486 |
+
"""Join the caption table onto the window plan by (stream, ts)."""
|
| 487 |
+
caps = store.table("context_captions").scan()
|
| 488 |
+
key = {(s, t): i for i, (s, t) in enumerate(
|
| 489 |
+
zip(caps.column("stream").to_pylist(), caps.column("ts").to_pylist()))}
|
| 490 |
+
vecs = np.asarray(caps.column("vector").to_pylist(), dtype=np.float32)
|
| 491 |
+
texts = caps.column("caption").to_pylist()
|
| 492 |
+
keep, tv, tt = [], [], []
|
| 493 |
+
for w in windows:
|
| 494 |
+
i = key.get((w[0], w[1]))
|
| 495 |
+
if i is not None:
|
| 496 |
+
keep.append(w)
|
| 497 |
+
tv.append(vecs[i])
|
| 498 |
+
tt.append(texts[i])
|
| 499 |
+
return keep, np.asarray(tv, dtype=np.float32), tt
|
| 500 |
+
|
| 501 |
+
|
| 502 |
+
def train_context(store, window_s=2.0, stride_s=0.5, val_frac=0.3,
|
| 503 |
+
d_in=48, d_out=48, epochs=400, cfg=None, verbose=True,
|
| 504 |
+
seed=0):
|
| 505 |
+
"""Fit the input PCA + tower to predict caption-LSA coordinates.
|
| 506 |
+
|
| 507 |
+
The validation split is by TIME, never at random: windows slide with 75%
|
| 508 |
+
overlap, so a random split puts near-duplicate windows on both sides and
|
| 509 |
+
reports a score that measures nothing.
|
| 510 |
+
"""
|
| 511 |
+
from .ctxtower import ContextCodec, retrieval_r1, save_tower, train_tower
|
| 512 |
+
|
| 513 |
+
windows = plan_windows(store, window_s, stride_s)
|
| 514 |
+
windows, _, cap_text = _labelled(store, windows)
|
| 515 |
+
if len(windows) < 16:
|
| 516 |
+
raise RuntimeError(f"only {len(windows)} captioned windows — run "
|
| 517 |
+
"caption_windows() first")
|
| 518 |
+
|
| 519 |
+
fv = store.table("frame_vectors").scan()
|
| 520 |
+
frame_mat = np.asarray(fv.column("vector").to_pylist(), dtype=np.float32)
|
| 521 |
+
codec = ContextCodec.fit(frame_mat, d_in=d_in)
|
| 522 |
+
space = caption_space(store, dim=d_out)
|
| 523 |
+
|
| 524 |
+
seqs = [codec.encode(s) for s in window_sequences(store, windows)]
|
| 525 |
+
targets = space.transform(cap_text)
|
| 526 |
+
|
| 527 |
+
t0s = np.array([w[1] for w in windows])
|
| 528 |
+
cut = np.quantile(t0s, 1.0 - val_frac)
|
| 529 |
+
val_idx = np.where(t0s >= cut)[0]
|
| 530 |
+
if verbose:
|
| 531 |
+
print(f" {len(windows)} labelled windows | train "
|
| 532 |
+
f"{len(windows) - len(val_idx)} | val {len(val_idx)} "
|
| 533 |
+
f"(time split at +{(cut - t0s.min()) / 1e9:.1f}s) | "
|
| 534 |
+
f"target dim {space.dim}", flush=True)
|
| 535 |
+
|
| 536 |
+
cfg = {**dict(d_in=codec.P_in.shape[0], d_out=space.dim), **(cfg or {})}
|
| 537 |
+
model, info = train_tower(seqs, targets, windows, val_idx, cfg=cfg,
|
| 538 |
+
epochs=epochs, verbose=verbose, seed=seed)
|
| 539 |
+
|
| 540 |
+
import mlx.core as mx
|
| 541 |
+
val = np.zeros(len(seqs), bool)
|
| 542 |
+
val[val_idx] = True
|
| 543 |
+
Xva = _pad_stack([s for s, m in zip(seqs, val) if m])
|
| 544 |
+
pred = np.array(model(mx.array(Xva)))
|
| 545 |
+
pred /= np.linalg.norm(pred, axis=1, keepdims=True) + 1e-8
|
| 546 |
+
tgt_va = targets[val]
|
| 547 |
+
|
| 548 |
+
# Baseline: predict the TRAIN-set mean target for every window. That is
|
| 549 |
+
# the best a model can do while knowing nothing about the specific window,
|
| 550 |
+
# so any lift over it is genuine per-window information and not the corpus
|
| 551 |
+
# prior leaking through.
|
| 552 |
+
prior = targets[~val].mean(axis=0)
|
| 553 |
+
prior /= np.linalg.norm(prior) + 1e-8
|
| 554 |
+
prior = np.repeat(prior[None, :], len(tgt_va), axis=0)
|
| 555 |
+
|
| 556 |
+
metrics = {
|
| 557 |
+
"labelled_windows": len(windows),
|
| 558 |
+
"train": int((~val).sum()), "val": int(val.sum()),
|
| 559 |
+
"target_dim": space.dim,
|
| 560 |
+
"val_R@1_tower": retrieval_r1(pred, tgt_va),
|
| 561 |
+
"val_cos_tower": float((pred * tgt_va).sum(1).mean()),
|
| 562 |
+
"val_cos_prior": float((prior * tgt_va).sum(1).mean()),
|
| 563 |
+
"params": int(sum(v.size for v in _flat(model).values())),
|
| 564 |
+
**{k: v for k, v in info.items() if k != "history"},
|
| 565 |
+
}
|
| 566 |
+
save_tower(model, codec,
|
| 567 |
+
{"window_s": window_s, "stride_s": stride_s, "d_out": space.dim,
|
| 568 |
+
"metrics": metrics, "history": info["history"]},
|
| 569 |
+
store.dir / "models" / "context")
|
| 570 |
+
return model, codec, metrics
|
| 571 |
+
|
| 572 |
+
|
| 573 |
+
def _flat(model):
|
| 574 |
+
from mlx.utils import tree_flatten
|
| 575 |
+
return {k: np.array(v) for k, v in tree_flatten(model.parameters())}
|
| 576 |
+
|
| 577 |
+
|
| 578 |
+
def _pad_stack(seqs):
|
| 579 |
+
T = max(s.shape[0] for s in seqs)
|
| 580 |
+
X = np.zeros((len(seqs), T, seqs[0].shape[1]), dtype=np.float32)
|
| 581 |
+
for i, s in enumerate(seqs):
|
| 582 |
+
X[i, :len(s)] = s
|
| 583 |
+
if len(s) < T:
|
| 584 |
+
X[i, len(s):] = s[-1]
|
| 585 |
+
return X
|
| 586 |
+
|
| 587 |
+
|
| 588 |
+
def build_context(store, window_s=None, stride_s=None, batch=64, verbose=True):
|
| 589 |
+
"""Materialise the `context` table: one context vector per window.
|
| 590 |
+
|
| 591 |
+
Where a window has a teacher caption, its EXACT caption-LSA vector is
|
| 592 |
+
stored. Where it does not, the tower's prediction is stored and the row is
|
| 593 |
+
flagged `estimated`. This is the ordinary database distinction between a
|
| 594 |
+
materialised value and an estimated one, and it is the point of having a
|
| 595 |
+
student at all: the VLM labels what you can afford, the tower covers the
|
| 596 |
+
rest, and the query does not care which it got.
|
| 597 |
+
"""
|
| 598 |
+
import time
|
| 599 |
+
|
| 600 |
+
import mlx.core as mx
|
| 601 |
+
|
| 602 |
+
from .ctxtower import load_tower
|
| 603 |
+
model, codec, meta = load_tower(store.dir / "models" / "context")
|
| 604 |
+
window_s = window_s or meta.get("window_s", 2.0)
|
| 605 |
+
stride_s = stride_s or meta.get("stride_s", 0.5)
|
| 606 |
+
space = caption_space(store, dim=meta.get("d_out", 48))
|
| 607 |
+
|
| 608 |
+
windows = plan_windows(store, window_s, stride_s)
|
| 609 |
+
raw = window_sequences(store, windows)
|
| 610 |
+
seqs = [codec.encode(s) for s in raw]
|
| 611 |
+
# The mean-pooled appearance vector for the SAME window, stored alongside.
|
| 612 |
+
# Fusing appearance with context otherwise needs a join between two tables
|
| 613 |
+
# built on different window plans; keeping both columns in one row makes
|
| 614 |
+
# the fused query two matmuls over one Parquet scan and no join at all.
|
| 615 |
+
appear = np.stack([s.mean(axis=0) for s in raw])
|
| 616 |
+
appear /= np.linalg.norm(appear, axis=1, keepdims=True) + 1e-8
|
| 617 |
+
|
| 618 |
+
t_start = time.time()
|
| 619 |
+
out = []
|
| 620 |
+
for i in range(0, len(seqs), batch):
|
| 621 |
+
out.append(np.array(model(mx.array(_pad_stack(seqs[i:i + batch])))))
|
| 622 |
+
z = np.concatenate(out, axis=0)
|
| 623 |
+
z /= np.linalg.norm(z, axis=1, keepdims=True) + 1e-8
|
| 624 |
+
infer_s = time.time() - t_start
|
| 625 |
+
|
| 626 |
+
caps = store.table("context_captions").scan()
|
| 627 |
+
known = {(a, b): c for a, b, c in zip(caps.column("stream").to_pylist(),
|
| 628 |
+
caps.column("ts").to_pylist(),
|
| 629 |
+
caps.column("caption").to_pylist())}
|
| 630 |
+
have = [(i, known[(w[0], w[1])]) for i, w in enumerate(windows)
|
| 631 |
+
if (w[0], w[1]) in known]
|
| 632 |
+
estimated = np.ones(len(windows), bool)
|
| 633 |
+
if have:
|
| 634 |
+
exact = space.transform([c for _, c in have])
|
| 635 |
+
for (i, _), v in zip(have, exact):
|
| 636 |
+
z[i] = v
|
| 637 |
+
estimated[i] = False
|
| 638 |
+
|
| 639 |
+
dim, adim = z.shape[1], appear.shape[1]
|
| 640 |
+
tbl = pa.table({
|
| 641 |
+
"ts": pa.array([w[1] for w in windows], pa.int64()),
|
| 642 |
+
"t1": pa.array([w[2] for w in windows], pa.int64()),
|
| 643 |
+
"stream": pa.array([w[0] for w in windows]),
|
| 644 |
+
"vector": pa.array([v.tolist() for v in z],
|
| 645 |
+
pa.list_(pa.float32(), dim)),
|
| 646 |
+
"appearance": pa.array([v.tolist() for v in appear],
|
| 647 |
+
pa.list_(pa.float32(), adim)),
|
| 648 |
+
"estimated": pa.array(estimated.tolist(), pa.bool_()),
|
| 649 |
+
})
|
| 650 |
+
st = store.table("context").state()
|
| 651 |
+
meta_out = {"dim": dim, "window_s": window_s, "stride_s": stride_s,
|
| 652 |
+
"estimated_rows": int(estimated.sum()),
|
| 653 |
+
"exact_rows": int((~estimated).sum())}
|
| 654 |
+
if st.files: # replace: one atomic commit
|
| 655 |
+
import uuid as _uuid
|
| 656 |
+
|
| 657 |
+
from .log import FileEntry
|
| 658 |
+
from .store import write_parquet
|
| 659 |
+
fname = f"part-{_uuid.uuid4().hex[:12]}.parquet"
|
| 660 |
+
p = store.dir / "tables" / "context" / fname
|
| 661 |
+
tbl = tbl.take(pc.sort_indices(tbl.column("ts")))
|
| 662 |
+
write_parquet(tbl, p)
|
| 663 |
+
tsv = tbl.column("ts").to_numpy()
|
| 664 |
+
version = store.table("context").log.commit(
|
| 665 |
+
op="replace", kind="embeddings", schema=str(tbl.schema),
|
| 666 |
+
add=[FileEntry(fname, len(tbl), p.stat().st_size,
|
| 667 |
+
int(tsv.min()), int(tsv.max()))],
|
| 668 |
+
remove=[f.path for f in st.files], meta=meta_out)
|
| 669 |
+
else:
|
| 670 |
+
version = store.table("context").append(tbl, kind="embeddings",
|
| 671 |
+
meta=meta_out)
|
| 672 |
+
if verbose:
|
| 673 |
+
print(f" {len(tbl)} context vectors ({int((~estimated).sum())} exact, "
|
| 674 |
+
f"{int(estimated.sum())} estimated) | tower inference "
|
| 675 |
+
f"{infer_s * 1000:.0f} ms "
|
| 676 |
+
f"({infer_s / len(tbl) * 1e6:.0f} us/window)")
|
| 677 |
+
return {"windows": len(tbl), "dim": dim, "version": version,
|
| 678 |
+
"inference_s": round(infer_s, 3),
|
| 679 |
+
"us_per_window": round(infer_s / len(tbl) * 1e6, 1), **meta_out}
|
| 680 |
+
|
| 681 |
+
|
| 682 |
+
# ---------------------------------------------------------------------------
|
| 683 |
+
# 6. Post-training cellular turnover — capacity the data cannot support is
|
| 684 |
+
# both latency and overfitting, so apoptosis pays twice.
|
| 685 |
+
# ---------------------------------------------------------------------------
|
| 686 |
+
def _prepare(store, window_s, stride_s, codec, val_frac, d_out):
|
| 687 |
+
windows, _, texts = _labelled(store, plan_windows(store, window_s,
|
| 688 |
+
stride_s))
|
| 689 |
+
seqs = [codec.encode(s) for s in window_sequences(store, windows)]
|
| 690 |
+
targets = caption_space(store, dim=d_out).transform(texts)
|
| 691 |
+
t0s = np.array([w[1] for w in windows])
|
| 692 |
+
val = t0s >= np.quantile(t0s, 1.0 - val_frac)
|
| 693 |
+
return windows, seqs, targets, val
|
| 694 |
+
|
| 695 |
+
|
| 696 |
+
def prune_context(store, val_frac=0.3, ppo_iters=40, sparsity_coef=0.15,
|
| 697 |
+
finetune_epochs=300, rebirth_fraction=0.5,
|
| 698 |
+
select_tolerance=0.03, verbose=True, seed=0):
|
| 699 |
+
"""apoptosis → re-settle → neurogenesis → re-settle, then COMPACT so the
|
| 700 |
+
surviving channels are the only ones that cost anything."""
|
| 701 |
+
import time
|
| 702 |
+
|
| 703 |
+
import mlx.core as mx
|
| 704 |
+
import mlx.nn as nn
|
| 705 |
+
|
| 706 |
+
from .ctxprune import compact, get_channel_mask, run_pruning_cycle
|
| 707 |
+
from .ctxtower import (load_tower, overlap_mask, retrieval_r1,
|
| 708 |
+
save_tower, siglip_loss)
|
| 709 |
+
|
| 710 |
+
model, codec, meta = load_tower(store.dir / "models" / "context")
|
| 711 |
+
windows, seqs, targets, val = _prepare(
|
| 712 |
+
store, meta["window_s"], meta["stride_s"], codec, val_frac,
|
| 713 |
+
meta.get("d_out", 48))
|
| 714 |
+
|
| 715 |
+
Xtr = mx.array(_pad_stack([s for s, m in zip(seqs, ~val) if m]))
|
| 716 |
+
Xva = mx.array(_pad_stack([s for s, m in zip(seqs, val) if m]))
|
| 717 |
+
Ytr, Yva = mx.array(targets[~val]), mx.array(targets[val])
|
| 718 |
+
ig_tr = mx.array(overlap_mask([w for w, m in zip(windows, ~val) if m]))
|
| 719 |
+
ig_va = mx.array(overlap_mask([w for w, m in zip(windows, val) if m]))
|
| 720 |
+
log_t = mx.array(np.float32(meta["metrics"]["log_t"]))
|
| 721 |
+
bias = mx.array(np.float32(meta["metrics"]["bias"]))
|
| 722 |
+
|
| 723 |
+
def _norm(a):
|
| 724 |
+
return a * mx.rsqrt(mx.sum(a * a, axis=-1, keepdims=True) + 1e-8)
|
| 725 |
+
|
| 726 |
+
def _loss(m, x, y, ig):
|
| 727 |
+
v = _norm(m(x))
|
| 728 |
+
u = _norm(y)
|
| 729 |
+
return siglip_loss(v, u, ignore=ig, log_t=log_t, bias=bias) \
|
| 730 |
+
+ 0.3 * mx.mean(1.0 - mx.sum(v * u, axis=-1))
|
| 731 |
+
|
| 732 |
+
def loss_of(m):
|
| 733 |
+
"""The number PPO is scored against: VALIDATION retrieval loss. Using
|
| 734 |
+
train loss here would reward a policy for keeping memorisers."""
|
| 735 |
+
m.set_training(False)
|
| 736 |
+
return float(_loss(m, Xva, Yva, ig_va).item())
|
| 737 |
+
|
| 738 |
+
lg = nn.value_and_grad(model, lambda m: _loss(m, Xtr, Ytr, ig_tr))
|
| 739 |
+
|
| 740 |
+
def r1(m):
|
| 741 |
+
m.set_training(False)
|
| 742 |
+
return retrieval_r1(np.array(_norm(m(Xva))), np.array(_norm(Yva)))
|
| 743 |
+
|
| 744 |
+
def latency(m, reps=20):
|
| 745 |
+
m.set_training(False)
|
| 746 |
+
mx.eval(m(Xva))
|
| 747 |
+
t = time.perf_counter()
|
| 748 |
+
for _ in range(reps):
|
| 749 |
+
mx.eval(m(Xva))
|
| 750 |
+
return (time.perf_counter() - t) / reps / Xva.shape[0] * 1e6
|
| 751 |
+
|
| 752 |
+
before = {"channels": int(get_channel_mask(model).sum()),
|
| 753 |
+
"val_loss": loss_of(model), "val_R@1": r1(model),
|
| 754 |
+
"us_per_window": latency(model),
|
| 755 |
+
"params": int(sum(v.size for v in _flat(model).values()))}
|
| 756 |
+
|
| 757 |
+
rec = run_pruning_cycle(model, np.array(Xtr), loss_of, lg,
|
| 758 |
+
ppo_iters=ppo_iters, sparsity_coef=sparsity_coef,
|
| 759 |
+
finetune_epochs=finetune_epochs,
|
| 760 |
+
rebirth_fraction=rebirth_fraction,
|
| 761 |
+
select_tolerance=select_tolerance, seed=seed,
|
| 762 |
+
verbose=verbose)
|
| 763 |
+
|
| 764 |
+
# Compaction must be a no-op numerically — it only deletes channels the
|
| 765 |
+
# mask already zeroed. Measure both sides and say so if they disagree,
|
| 766 |
+
# rather than quietly shipping a tower that differs from the one the
|
| 767 |
+
# search selected.
|
| 768 |
+
restored = loss_of(model)
|
| 769 |
+
model, kept = compact(model)
|
| 770 |
+
after = {"channels": int(len(kept)), "val_loss": loss_of(model),
|
| 771 |
+
"loss_before_compaction": restored, "val_R@1": r1(model),
|
| 772 |
+
"us_per_window": latency(model),
|
| 773 |
+
"params": int(sum(v.size for v in _flat(model).values()))}
|
| 774 |
+
rec["before"], rec["after"] = before, after
|
| 775 |
+
drift = abs(after["val_loss"] - restored)
|
| 776 |
+
if drift > 1e-3:
|
| 777 |
+
print(f" WARNING: compaction changed val loss by {drift:.4f} "
|
| 778 |
+
f"({restored:.4f} -> {after['val_loss']:.4f}) — expected ~0")
|
| 779 |
+
if verbose:
|
| 780 |
+
print(f"\n channels {before['channels']} -> {after['channels']} | "
|
| 781 |
+
f"params {before['params']:,} -> {after['params']:,} | "
|
| 782 |
+
f"{before['us_per_window']:.1f} -> {after['us_per_window']:.1f} "
|
| 783 |
+
f"us/window | val loss {before['val_loss']:.4f} -> "
|
| 784 |
+
f"{after['val_loss']:.4f}")
|
| 785 |
+
save_tower(model, codec,
|
| 786 |
+
{**{k: v for k, v in meta.items() if k != "history"},
|
| 787 |
+
"pruned": {"before": before, "after": after,
|
| 788 |
+
"ppo_history": rec.get("ppo_history"),
|
| 789 |
+
"keep_probs": rec.get("keep_probs"),
|
| 790 |
+
"reverse_attention": rec.get("reverse_attention"),
|
| 791 |
+
"selected": rec.get("selected"),
|
| 792 |
+
"stages": [{k: v for k, v in s.items()
|
| 793 |
+
if k != "mask"} for s in rec["stages"]]}},
|
| 794 |
+
store.dir / "models" / "context")
|
| 795 |
+
return model, rec
|
| 796 |
+
|
| 797 |
+
|
| 798 |
+
# ---------------------------------------------------------------------------
|
| 799 |
+
# 7. The query path — two matmuls, no VLM, no decode
|
| 800 |
+
# ---------------------------------------------------------------------------
|
| 801 |
+
_CTX_CACHE: dict = {}
|
| 802 |
+
|
| 803 |
+
|
| 804 |
+
def _ctx_matrix(store):
|
| 805 |
+
t = store.table("context")
|
| 806 |
+
key = (str(store.dir), t.state().version)
|
| 807 |
+
if key not in _CTX_CACHE:
|
| 808 |
+
_CTX_CACHE[key] = np.stack([
|
| 809 |
+
np.asarray(v, dtype=np.float32)
|
| 810 |
+
for v in t.scan().column("vector").to_pylist()])
|
| 811 |
+
return _CTX_CACHE[key]
|
| 812 |
+
|
| 813 |
+
|
| 814 |
+
def _lexical(store):
|
| 815 |
+
"""(vectorizer, matrix, has_caption) aligned to the `context` table rows.
|
| 816 |
+
|
| 817 |
+
Windows the VLM never captioned have no text, so they are marked as
|
| 818 |
+
abstentions rather than as empty documents — an empty document would score
|
| 819 |
+
0 on every query and be ranked last by the lexical ranker, which is a veto
|
| 820 |
+
dressed up as evidence.
|
| 821 |
+
"""
|
| 822 |
+
t = store.table("context")
|
| 823 |
+
key = (str(store.dir), t.state().version, "lex")
|
| 824 |
+
if key in _CTX_CACHE:
|
| 825 |
+
return _CTX_CACHE[key]
|
| 826 |
+
ctx = t.scan()
|
| 827 |
+
rows = list(zip(ctx.column("stream").to_pylist(),
|
| 828 |
+
ctx.column("ts").to_pylist()))
|
| 829 |
+
caps = store.table("context_captions").scan()
|
| 830 |
+
known = dict(zip(zip(caps.column("stream").to_pylist(),
|
| 831 |
+
caps.column("ts").to_pylist()),
|
| 832 |
+
caps.column("caption").to_pylist()))
|
| 833 |
+
texts = [known.get(r, "") for r in rows]
|
| 834 |
+
has = np.array([bool(x) for x in texts])
|
| 835 |
+
space = caption_space(store)
|
| 836 |
+
X = space.vec.transform(texts)
|
| 837 |
+
nrm = np.sqrt(np.asarray(X.multiply(X).sum(1))).ravel() + 1e-8
|
| 838 |
+
_CTX_CACHE[key] = (space, X, nrm, has)
|
| 839 |
+
return _CTX_CACHE[key]
|
| 840 |
+
|
| 841 |
+
|
| 842 |
+
DEFAULT_WEIGHTS = {"appearance": 1.0, "context": 1.0, "lexical": 1.0}
|
| 843 |
+
|
| 844 |
+
|
| 845 |
+
def search(store, text, k=10, merge=True, t0=None, t1=None, streams=None,
|
| 846 |
+
weights=None, rrf_k=60.0, neg_weight=0.5, min_score=None,
|
| 847 |
+
percentile=None, rerank=False, rerank_top=12, rerank_alpha=0.7,
|
| 848 |
+
explain_top=0):
|
| 849 |
+
"""Hybrid contextual search: three rankers fused by reciprocal rank.
|
| 850 |
+
|
| 851 |
+
appearance SigLIP image-text cosine over the window's frames.
|
| 852 |
+
Knows what OBJECTS are present. Order-blind.
|
| 853 |
+
context caption-LSA cosine. Knows what is HAPPENING, because the
|
| 854 |
+
caption was written by a VLM that watched three frames.
|
| 855 |
+
lexical TF-IDF over the caption text. Exact term evidence — the
|
| 856 |
+
ranker that actually knows what "red" means.
|
| 857 |
+
|
| 858 |
+
Fusing by RRF rather than by a weighted score sum is the fix for
|
| 859 |
+
"crossing red car" returning any clip of someone crossing: RRF rewards
|
| 860 |
+
agreement across rankers, so a candidate that satisfies one strong signal
|
| 861 |
+
alone can no longer win. See elidedb.fusion.
|
| 862 |
+
|
| 863 |
+
`rerank=True` adds a final VLM pass over the top `rerank_top` — the
|
| 864 |
+
expensive operator, last, on an already-pruned set.
|
| 865 |
+
"""
|
| 866 |
+
from .embeddings import _parse_query, _score_windows, embed_text
|
| 867 |
+
from .fusion import explain_fusion, rrf
|
| 868 |
+
# Fail with a sentence, not a KeyError from three frames down. An empty
|
| 869 |
+
# table has no schema, so the first column access explodes with
|
| 870 |
+
# 'Field "t1" does not exist' — true, useless, and it names the wrong
|
| 871 |
+
# problem.
|
| 872 |
+
if not store.table("context").state().files:
|
| 873 |
+
raise RuntimeError(
|
| 874 |
+
f"store '{store.name}' has no context index. Build it with "
|
| 875 |
+
"store.index_context() (frames -> per-frame vectors -> VLM "
|
| 876 |
+
"captions -> context table), or use store.search_text() for "
|
| 877 |
+
"appearance-only search.")
|
| 878 |
+
pos, neg = _parse_query(text)
|
| 879 |
+
ctx_tbl = store.table("context").scan()
|
| 880 |
+
all_t0 = ctx_tbl.column("ts").to_numpy()
|
| 881 |
+
all_t1 = ctx_tbl.column("t1").to_numpy()
|
| 882 |
+
all_s = ctx_tbl.column("stream").to_numpy(zero_copy_only=False)
|
| 883 |
+
|
| 884 |
+
# hybrid retrieval: time and stream predicates are pushed INTO candidate
|
| 885 |
+
# selection, not applied to a global top-k afterwards
|
| 886 |
+
pred = np.ones(len(all_t0), bool)
|
| 887 |
+
if t0 is not None:
|
| 888 |
+
pred &= all_t1 >= t0
|
| 889 |
+
if t1 is not None:
|
| 890 |
+
pred &= all_t0 <= t1
|
| 891 |
+
if streams:
|
| 892 |
+
pred &= np.isin(all_s, list(streams))
|
| 893 |
+
idx = np.where(pred)[0]
|
| 894 |
+
if len(idx) == 0:
|
| 895 |
+
return [], {"index": "context", "total": len(all_t0), "scanned": 0,
|
| 896 |
+
"segments": 0}
|
| 897 |
+
|
| 898 |
+
APP = _appearance_matrix(store)
|
| 899 |
+
CTX = _ctx_matrix(store)
|
| 900 |
+
space, X, nrm, has_cap = _lexical(store)
|
| 901 |
+
|
| 902 |
+
pos_app = np.stack([embed_text(p) for p in pos])
|
| 903 |
+
neg_app = np.stack([embed_text(n) for n in neg]) if neg else None
|
| 904 |
+
app = _score_windows(APP, idx, pos_app, neg_app, neg_weight)
|
| 905 |
+
|
| 906 |
+
q_join = " ".join(pos)
|
| 907 |
+
ctx_sc = CTX[idx] @ space.transform([q_join])[0]
|
| 908 |
+
|
| 909 |
+
qv = space.vec.transform([q_join])
|
| 910 |
+
lex = np.asarray((X[idx] @ qv.T).todense()).ravel() / nrm[idx]
|
| 911 |
+
lex = np.where(has_cap[idx], lex, np.nan) # abstain, do not veto
|
| 912 |
+
|
| 913 |
+
rankings = {"appearance": app, "context": ctx_sc, "lexical": lex}
|
| 914 |
+
w = {**DEFAULT_WEIGHTS, **(weights or {})}
|
| 915 |
+
scores = rrf(rankings, w, rrf_k)
|
| 916 |
+
|
| 917 |
+
stats = {"index": "context", "total": len(all_t0), "scanned": len(idx),
|
| 918 |
+
"method": "rrf", "rrf_k": rrf_k, "weights": w,
|
| 919 |
+
"predicate_candidates": int(pred.sum()),
|
| 920 |
+
"captioned_candidates": int(has_cap[idx].sum()),
|
| 921 |
+
"positive_terms": pos, "negative_terms": neg}
|
| 922 |
+
if explain_top:
|
| 923 |
+
stats["why"] = explain_fusion(rankings, w, rrf_k, top=explain_top)
|
| 924 |
+
|
| 925 |
+
keep = np.ones(len(idx), bool)
|
| 926 |
+
if percentile is not None:
|
| 927 |
+
keep &= scores >= np.percentile(scores, percentile)
|
| 928 |
+
if min_score is not None:
|
| 929 |
+
keep &= scores >= min_score
|
| 930 |
+
idx, scores = idx[keep], scores[keep]
|
| 931 |
+
stats["after_floor"] = int(len(idx))
|
| 932 |
+
|
| 933 |
+
hits = _segments(idx, scores, all_s, all_t0, all_t1, k, merge, stats)
|
| 934 |
+
if rerank and hits:
|
| 935 |
+
from .rerank import rerank_hits
|
| 936 |
+
hits, info = rerank_hits(store, hits, text, top_n=rerank_top,
|
| 937 |
+
alpha=rerank_alpha)
|
| 938 |
+
stats["rerank"] = info
|
| 939 |
+
return hits, stats
|
| 940 |
+
|
| 941 |
+
|
| 942 |
+
def _segments(idx, scores, all_s, all_t0, all_t1, k, merge, stats):
|
| 943 |
+
"""Merge qualifying windows into maximal runs per stream.
|
| 944 |
+
|
| 945 |
+
Fixed windows are an INDEXING granularity, not an answer granularity: a
|
| 946 |
+
20 s event should come back as one 20 s hit, and a query matching only 2 s
|
| 947 |
+
of it should come back as that 2 s.
|
| 948 |
+
"""
|
| 949 |
+
streams_sel, w_t0, w_t1 = all_s[idx], all_t0[idx], all_t1[idx]
|
| 950 |
+
if len(idx) == 0:
|
| 951 |
+
stats["segments"] = 0
|
| 952 |
+
return []
|
| 953 |
+
if not merge:
|
| 954 |
+
order = np.argsort(scores)[::-1][:k]
|
| 955 |
+
return [{"stream": str(streams_sel[i]), "t0": int(w_t0[i]),
|
| 956 |
+
"t1": int(w_t1[i]), "score": float(scores[i]),
|
| 957 |
+
"windows": 1} for i in order]
|
| 958 |
+
med, top = float(np.median(scores)), float(scores.max())
|
| 959 |
+
thr = med + 0.55 * (top - med)
|
| 960 |
+
stats["threshold"] = round(thr, 6)
|
| 961 |
+
qual = np.where(scores >= thr)[0]
|
| 962 |
+
order = np.lexsort((w_t0[qual], streams_sel[qual]))
|
| 963 |
+
qual = qual[order]
|
| 964 |
+
gap = int(np.median(w_t1[qual] - w_t0[qual])) + 1 if len(qual) else 0
|
| 965 |
+
segs = []
|
| 966 |
+
for i in qual:
|
| 967 |
+
s_, a, b, sc = (str(streams_sel[i]), int(w_t0[i]), int(w_t1[i]),
|
| 968 |
+
float(scores[i]))
|
| 969 |
+
last = segs[-1] if segs else None
|
| 970 |
+
if last and last["stream"] == s_ and a - last["t1"] <= gap:
|
| 971 |
+
last["t1"] = max(last["t1"], b)
|
| 972 |
+
last["score"] = max(last["score"], sc)
|
| 973 |
+
last["mean"] = (last["mean"] * last["windows"] + sc) / (last["windows"] + 1)
|
| 974 |
+
last["windows"] += 1
|
| 975 |
+
else:
|
| 976 |
+
segs.append({"stream": s_, "t0": a, "t1": b, "score": sc,
|
| 977 |
+
"mean": sc, "windows": 1})
|
| 978 |
+
segs.sort(key=lambda g: -g["score"])
|
| 979 |
+
stats["qualifying_windows"] = len(qual)
|
| 980 |
+
stats["segments"] = len(segs)
|
| 981 |
+
return segs[:k]
|
| 982 |
+
|
| 983 |
+
|
| 984 |
+
def _appearance_matrix(store):
|
| 985 |
+
t = store.table("context")
|
| 986 |
+
key = (str(store.dir), t.state().version, "app")
|
| 987 |
+
if key not in _CTX_CACHE:
|
| 988 |
+
_CTX_CACHE[key] = np.stack([
|
| 989 |
+
np.asarray(v, dtype=np.float32)
|
| 990 |
+
for v in t.scan().column("appearance").to_pylist()])
|
| 991 |
+
return _CTX_CACHE[key]
|
| 992 |
+
|
| 993 |
+
|
| 994 |
+
def explain(store, t0, t1, stream=None):
|
| 995 |
+
"""What the database believes is happening in a window — the teacher's own
|
| 996 |
+
words. Lets a result be checked rather than trusted."""
|
| 997 |
+
caps = store.table("context_captions").scan()
|
| 998 |
+
out = []
|
| 999 |
+
if len(caps) == 0 or "stream" not in caps.column_names:
|
| 1000 |
+
return out # store has no captions (yet) — nothing to explain
|
| 1001 |
+
for s, a, b, c in zip(caps.column("stream").to_pylist(),
|
| 1002 |
+
caps.column("ts").to_pylist(),
|
| 1003 |
+
caps.column("t1").to_pylist(),
|
| 1004 |
+
caps.column("caption").to_pylist()):
|
| 1005 |
+
if stream and s != stream:
|
| 1006 |
+
continue
|
| 1007 |
+
if b >= t0 and a <= t1:
|
| 1008 |
+
out.append({"stream": s, "t0": a, "t1": b, "caption": c})
|
| 1009 |
+
return out
|
| 1010 |
+
|
| 1011 |
+
|
| 1012 |
+
def _embed_frames_fdnnv(store, frame_table, incremental, verbose, streams):
|
| 1013 |
+
"""Every frame through the FDNN-V streaming encoder -> frame_vectors."""
|
| 1014 |
+
import pyarrow as pa
|
| 1015 |
+
|
| 1016 |
+
from .fdnnvideo import embed_stream, load_encoder
|
| 1017 |
+
mdir = store.dir / "models" / "fdnnv"
|
| 1018 |
+
if not (mdir / "encoder.json").exists():
|
| 1019 |
+
# a NEW store has no encoder yet — adopt one from a sibling store
|
| 1020 |
+
# and COPY it in, so the store stays self-contained and the exact
|
| 1021 |
+
# weights that wrote its vectors are pinned with its data
|
| 1022 |
+
import shutil
|
| 1023 |
+
donors = sorted(store.dir.parent.glob("*/models/fdnnv/encoder.json"),
|
| 1024 |
+
key=lambda p: p.stat().st_mtime, reverse=True)
|
| 1025 |
+
if not donors:
|
| 1026 |
+
raise RuntimeError(
|
| 1027 |
+
"no FDNN-V encoder found in this store or any sibling — "
|
| 1028 |
+
"train one first (scripts/fdnnv_train.py)")
|
| 1029 |
+
shutil.copytree(donors[0].parent, mdir)
|
| 1030 |
+
if verbose:
|
| 1031 |
+
print(f" adopted encoder from {donors[0].parent}", flush=True)
|
| 1032 |
+
model, meta = load_encoder(mdir)
|
| 1033 |
+
# FRESH-DATA FIDELITY GATE: the fast student was distilled on one
|
| 1034 |
+
# style of footage; on an arbitrary upload its fidelity to the teacher
|
| 1035 |
+
# is unknown. 32 of THIS store's frames go through both encoders; if
|
| 1036 |
+
# mean cosine < 0.90 the store gets the TEACHER (slower ingest, right
|
| 1037 |
+
# space) instead of a fast-but-wrong index. Automatic — a fresh
|
| 1038 |
+
# customer never has to know this exists.
|
| 1039 |
+
try:
|
| 1040 |
+
from PIL import Image
|
| 1041 |
+
from .embeddings import DEFAULT_MODEL, _embed_images
|
| 1042 |
+
from .video import FrameSet
|
| 1043 |
+
frames_all = store.table(frame_table).scan()
|
| 1044 |
+
pick = np.linspace(0, len(frames_all) - 1,
|
| 1045 |
+
min(32, len(frames_all))).round().astype(int)
|
| 1046 |
+
rows = frames_all.take(pick)
|
| 1047 |
+
ts_s, sv, _, _ = embed_stream(store, model, rows)
|
| 1048 |
+
dec = FrameSet(store, frame_table, rows).decode(width=448)
|
| 1049 |
+
imgs = [Image.fromarray(d[1]) for d in sorted(dec)]
|
| 1050 |
+
tv = _embed_images(imgs, DEFAULT_MODEL)
|
| 1051 |
+
n = min(len(sv), len(tv))
|
| 1052 |
+
svn = sv[:n] / (np.linalg.norm(sv[:n], axis=1,
|
| 1053 |
+
keepdims=True) + 1e-8)
|
| 1054 |
+
fid = float((svn * tv[:n]).sum(1).mean())
|
| 1055 |
+
if verbose:
|
| 1056 |
+
print(f" student fidelity on this corpus: {fid:.3f}",
|
| 1057 |
+
flush=True)
|
| 1058 |
+
if fid < 0.90:
|
| 1059 |
+
print(f" fidelity {fid:.3f} < 0.90 — falling back to the "
|
| 1060 |
+
f"TEACHER encoder for this store", flush=True)
|
| 1061 |
+
return embed_frames(store, frame_table, model="fast",
|
| 1062 |
+
incremental=incremental, verbose=verbose,
|
| 1063 |
+
streams=streams, engine="siglip")
|
| 1064 |
+
except Exception as e:
|
| 1065 |
+
if verbose:
|
| 1066 |
+
print(f" fidelity gate skipped ({type(e).__name__})",
|
| 1067 |
+
flush=True)
|
| 1068 |
+
frames = store.table(frame_table).scan()
|
| 1069 |
+
allst = sorted(set(frames.column("stream").to_pylist()))
|
| 1070 |
+
use = [s for s in allst if s in streams] if streams else allst
|
| 1071 |
+
done = {}
|
| 1072 |
+
if incremental:
|
| 1073 |
+
try:
|
| 1074 |
+
prev = store.table("frame_vectors").scan()
|
| 1075 |
+
for s_, t_ in zip(prev.column("stream").to_pylist(),
|
| 1076 |
+
prev.column("ts").to_pylist()):
|
| 1077 |
+
done[s_] = max(done.get(s_, -1), t_)
|
| 1078 |
+
except Exception:
|
| 1079 |
+
pass
|
| 1080 |
+
t_start = time.time()
|
| 1081 |
+
rows_ts, rows_stream, rows_vec = [], [], []
|
| 1082 |
+
for s in use:
|
| 1083 |
+
sel = frames.filter(pc.equal(frames.column("stream"), s))
|
| 1084 |
+
if s in done:
|
| 1085 |
+
sel = sel.filter(pc.greater(sel.column("ts"), done[s]))
|
| 1086 |
+
if len(sel) == 0:
|
| 1087 |
+
continue
|
| 1088 |
+
sel = sel.take(pc.sort_indices(sel.column("ts")))
|
| 1089 |
+
ts, vecs, dec_s, emb_s = embed_stream(store, model, sel)
|
| 1090 |
+
if verbose:
|
| 1091 |
+
print(f" {s}: {len(ts):,} frames (decode {dec_s:.1f}s, "
|
| 1092 |
+
f"embed {emb_s:.1f}s)", flush=True)
|
| 1093 |
+
rows_ts.extend(int(t) for t in ts)
|
| 1094 |
+
rows_stream.extend([s] * len(ts))
|
| 1095 |
+
rows_vec.extend(vecs)
|
| 1096 |
+
if not rows_ts:
|
| 1097 |
+
return {"frames": 0, "note": "nothing new (incremental)"}
|
| 1098 |
+
dim = len(rows_vec[0])
|
| 1099 |
+
tbl = pa.table({
|
| 1100 |
+
"ts": pa.array(rows_ts, pa.int64()),
|
| 1101 |
+
"stream": pa.array(rows_stream),
|
| 1102 |
+
"vector": pa.array([v.tolist() for v in rows_vec],
|
| 1103 |
+
pa.list_(pa.float32(), dim)),
|
| 1104 |
+
})
|
| 1105 |
+
version = store.table("frame_vectors").append(
|
| 1106 |
+
tbl, kind="embeddings",
|
| 1107 |
+
meta={"model": "fdnnv", "teacher": meta.get("teacher"),
|
| 1108 |
+
"dim": dim, "every_frame": True,
|
| 1109 |
+
"source_table": frame_table})
|
| 1110 |
+
return {"frames": len(tbl), "dim": dim, "version": version,
|
| 1111 |
+
"engine": "fdnnv",
|
| 1112 |
+
"seconds": round(time.time() - t_start, 1)}
|
python/elidedb/corpus.py
ADDED
|
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""S5: dataset LAYOUT as configuration, not as constants in six scripts.
|
| 2 |
+
|
| 3 |
+
`CAM = "observation.images.image_0"`, `EPOCH_NS`, `FILE_STRIDE_NS` and
|
| 4 |
+
`FPS = 5.0` were copied into bridge_ingest, bridge4h, bridge_full,
|
| 5 |
+
rebuild_bench and full_write. None of them is a task prior - they
|
| 6 |
+
describe where the bytes are, not what the bytes mean - but six copies
|
| 7 |
+
of a dataset's shape is still the dataset hardwired into the engine, and
|
| 8 |
+
pointing ElideDB at a driving log meant editing six files.
|
| 9 |
+
|
| 10 |
+
A corpus descriptor is data. It lives beside the raw files, it is read,
|
| 11 |
+
and the engine has no opinion about its contents.
|
| 12 |
+
|
| 13 |
+
{
|
| 14 |
+
"name": "bridge",
|
| 15 |
+
"root": "data/bridge",
|
| 16 |
+
"streams": [{"id": "observation.images.image_0",
|
| 17 |
+
"path": "videos/{id}/chunk-000/file-{file:03d}.mp4",
|
| 18 |
+
"fps": 5.0}],
|
| 19 |
+
"epoch_ns": 1704067200000000000,
|
| 20 |
+
"file_stride_ns": 20000000000000,
|
| 21 |
+
"episodes": {"table": "meta/episodes/chunk-000/file-000.parquet",
|
| 22 |
+
"index": "episode_index", "length": "length"},
|
| 23 |
+
"exclude": ["tasks"]
|
| 24 |
+
}
|
| 25 |
+
|
| 26 |
+
`exclude` is the important field and the reason this is not merely
|
| 27 |
+
tidiness: it is where a corpus declares which of its own columns are
|
| 28 |
+
ANNOTATION rather than observation. Bridge's `tasks` column says what
|
| 29 |
+
each clip is about, which is precisely what the database is supposed to
|
| 30 |
+
work out from pixels; ingesting it would turn every retrieval number
|
| 31 |
+
into a join against a label. bridge_ingest.py enforced that by hand, in
|
| 32 |
+
a docstring. Here it is a field the ingest path honours for any corpus,
|
| 33 |
+
so the next dataset cannot forget.
|
| 34 |
+
"""
|
| 35 |
+
from __future__ import annotations
|
| 36 |
+
|
| 37 |
+
import json
|
| 38 |
+
from pathlib import Path
|
| 39 |
+
|
| 40 |
+
# Fields every corpus must supply. No defaults for the ones that encode
|
| 41 |
+
# the dataset's shape: guessing an epoch or a frame rate silently
|
| 42 |
+
# misaligns every timestamp in the store, and a loud failure is cheaper.
|
| 43 |
+
REQUIRED = ("name", "root", "streams")
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
class Corpus:
|
| 47 |
+
"""A dataset's shape, read from disk. Never inferred, never guessed."""
|
| 48 |
+
|
| 49 |
+
def __init__(self, d: dict, path=None):
|
| 50 |
+
missing = [k for k in REQUIRED if k not in d]
|
| 51 |
+
if missing:
|
| 52 |
+
raise ValueError(f"corpus descriptor missing {missing}")
|
| 53 |
+
self.d, self.path = d, path
|
| 54 |
+
self.name = d["name"]
|
| 55 |
+
self.root = Path(d["root"])
|
| 56 |
+
self.streams = d["streams"]
|
| 57 |
+
self.epoch_ns = int(d.get("epoch_ns", 0))
|
| 58 |
+
self.file_stride_ns = int(d.get("file_stride_ns", 0))
|
| 59 |
+
# columns the corpus declares as ANNOTATION - never ingested
|
| 60 |
+
self.exclude = set(d.get("exclude", ()))
|
| 61 |
+
|
| 62 |
+
@staticmethod
|
| 63 |
+
def load(path):
|
| 64 |
+
p = Path(path)
|
| 65 |
+
if p.is_dir():
|
| 66 |
+
p = p / "corpus.json"
|
| 67 |
+
return Corpus(json.loads(p.read_text()), p)
|
| 68 |
+
|
| 69 |
+
def stream(self, sid=None):
|
| 70 |
+
if sid is None:
|
| 71 |
+
return self.streams[0]
|
| 72 |
+
for s in self.streams:
|
| 73 |
+
if s["id"] == sid:
|
| 74 |
+
return s
|
| 75 |
+
raise KeyError(f"no stream {sid} in corpus {self.name}")
|
| 76 |
+
|
| 77 |
+
def fps(self, sid=None):
|
| 78 |
+
return float(self.stream(sid).get("fps", 0)) or None
|
| 79 |
+
|
| 80 |
+
def media(self, file_index, sid=None):
|
| 81 |
+
s = self.stream(sid)
|
| 82 |
+
return self.root / s["path"].format(id=s["id"], file=int(file_index))
|
| 83 |
+
|
| 84 |
+
def base_ns(self, file_index):
|
| 85 |
+
"""Timeline origin for a file. 0/0 means the corpus carries real
|
| 86 |
+
timestamps and no synthetic clock is being imposed."""
|
| 87 |
+
return self.epoch_ns + int(file_index) * self.file_stride_ns
|
| 88 |
+
|
| 89 |
+
def is_annotation(self, column: str) -> bool:
|
| 90 |
+
"""True for columns the corpus declares as annotation.
|
| 91 |
+
|
| 92 |
+
The ingest path must refuse these. A column that says what a clip
|
| 93 |
+
is ABOUT is the answer to the question the database exists to
|
| 94 |
+
answer, and ingesting it makes every later number a lookup.
|
| 95 |
+
"""
|
| 96 |
+
return column in self.exclude
|
python/elidedb/cracked.py
ADDED
|
@@ -0,0 +1,133 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Sharp text search: student shortlist, teacher verdict, cracked cache.
|
| 2 |
+
|
| 3 |
+
THE PROBLEM THIS SOLVES
|
| 4 |
+
-----------------------
|
| 5 |
+
Student (FDNN-V) embeddings rank the whole corpus in ~1 ms but their text
|
| 6 |
+
ranking is measurably weak (top-10 agreement 0.5/10 vs the 4.3/10
|
| 7 |
+
teacher-self ceiling). The teacher ranks text pristinely but costs 27.7 ms
|
| 8 |
+
per frame — unaffordable corpus-wide, affordable on a SHORTLIST.
|
| 9 |
+
|
| 10 |
+
So the query plan is the database's oldest shape, applied to models:
|
| 11 |
+
|
| 12 |
+
cheap operator over everything -> expensive operator over survivors
|
| 13 |
+
student ranks all windows (~1ms) teacher re-scores top-N (~30ms each)
|
| 14 |
+
|
| 15 |
+
THE CRACKING PART
|
| 16 |
+
-----------------
|
| 17 |
+
Every teacher vector computed for a query is WRITTEN BACK to the
|
| 18 |
+
`teacher_windows` table. The next query that touches those windows pays
|
| 19 |
+
nothing. Like database cracking (Idreos et al., CIDR 2007), the index
|
| 20 |
+
materialises as a side effect of the workload: hot regions of the corpus
|
| 21 |
+
become pristine after their first visit, cold regions never cost a cent.
|
| 22 |
+
Cost scales with what users ask, not with what they store.
|
| 23 |
+
|
| 24 |
+
The teacher here is SigLIP-224 ("fast") — measured within the teacher-self
|
| 25 |
+
agreement band of the 384 model at a third of the price — and the query text
|
| 26 |
+
is embedded with the SAME checkpoint, so the rerank compares like with like.
|
| 27 |
+
"""
|
| 28 |
+
from __future__ import annotations
|
| 29 |
+
|
| 30 |
+
import time
|
| 31 |
+
|
| 32 |
+
import numpy as np
|
| 33 |
+
import pyarrow as pa
|
| 34 |
+
import pyarrow.compute as pc
|
| 35 |
+
|
| 36 |
+
TEACHER = "fast" # resolved by embeddings.resolve_model
|
| 37 |
+
_CACHE: dict = {} # (store, version) -> {key: vec}
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def _cache_map(store):
|
| 41 |
+
t = store.table("teacher_windows")
|
| 42 |
+
st = t.state()
|
| 43 |
+
key = (str(store.dir), st.version)
|
| 44 |
+
if key not in _CACHE:
|
| 45 |
+
m = {}
|
| 46 |
+
if st.files:
|
| 47 |
+
tb = t.scan()
|
| 48 |
+
for s, a, v in zip(tb.column("stream").to_pylist(),
|
| 49 |
+
tb.column("ts").to_pylist(),
|
| 50 |
+
tb.column("vector").to_pylist()):
|
| 51 |
+
m[(s, int(a))] = np.asarray(v, np.float32)
|
| 52 |
+
_CACHE.clear()
|
| 53 |
+
_CACHE[key] = m
|
| 54 |
+
return _CACHE[key]
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def _teacher_embed_windows(store, wins):
|
| 58 |
+
"""Decode ONE centre frame per window through the byte-range path and
|
| 59 |
+
embed with the teacher. `wins` = [(stream, t0, t1)]."""
|
| 60 |
+
from PIL import Image
|
| 61 |
+
|
| 62 |
+
from .embeddings import _embed_images
|
| 63 |
+
frames = store.table("frames").scan()
|
| 64 |
+
imgs, keys = [], []
|
| 65 |
+
for (s, a, b) in wins:
|
| 66 |
+
mid = (a + b) // 2
|
| 67 |
+
sel = frames.filter(pc.and_(
|
| 68 |
+
pc.equal(frames.column("stream"), s),
|
| 69 |
+
pc.and_(pc.greater_equal(frames.column("ts"),
|
| 70 |
+
mid - 2_000_000_000),
|
| 71 |
+
pc.less_equal(frames.column("ts"),
|
| 72 |
+
mid + 2_000_000_000))))
|
| 73 |
+
from .video import FrameSet
|
| 74 |
+
dec = FrameSet(store, "frames", sel).decode(stream=s, width=448,
|
| 75 |
+
limit=1)
|
| 76 |
+
if dec:
|
| 77 |
+
imgs.append(Image.fromarray(dec[0][1]))
|
| 78 |
+
keys.append((s, a, b))
|
| 79 |
+
if not imgs:
|
| 80 |
+
return {}
|
| 81 |
+
vecs = _embed_images(imgs, TEACHER)
|
| 82 |
+
return {k: v for k, v in zip(keys, vecs)}
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def search_sharp(store, text, k=10, shortlist=48):
|
| 86 |
+
"""Student ranks everything; the teacher re-scores the shortlist; misses
|
| 87 |
+
are cached into the store. Returns (hits, stats)."""
|
| 88 |
+
from .embeddings import _vec_table, embed_text
|
| 89 |
+
t_start = time.perf_counter()
|
| 90 |
+
tbl, vecs = _vec_table(store, "embeddings")
|
| 91 |
+
a_t0 = tbl.column("ts").to_numpy()
|
| 92 |
+
a_t1 = tbl.column("t1").to_numpy()
|
| 93 |
+
a_s = tbl.column("stream").to_numpy(zero_copy_only=False)
|
| 94 |
+
|
| 95 |
+
q_student = embed_text(text) # 384-space (student's)
|
| 96 |
+
order = np.argsort(-(vecs @ q_student))[:shortlist]
|
| 97 |
+
wins = [(str(a_s[i]), int(a_t0[i]), int(a_t1[i])) for i in order]
|
| 98 |
+
|
| 99 |
+
cache = _cache_map(store)
|
| 100 |
+
missing = [w for w in wins if (w[0], w[1]) not in cache]
|
| 101 |
+
t_miss = time.perf_counter()
|
| 102 |
+
if missing:
|
| 103 |
+
fresh = _teacher_embed_windows(store, missing)
|
| 104 |
+
if fresh:
|
| 105 |
+
ft = pa.table({
|
| 106 |
+
"ts": pa.array([a for (_, a, _b) in fresh], pa.int64()),
|
| 107 |
+
"t1": pa.array([b for (_, _a, b) in fresh], pa.int64()),
|
| 108 |
+
"stream": pa.array([s for (s, _a, _b) in fresh]),
|
| 109 |
+
"vector": pa.array([v.tolist() for v in fresh.values()],
|
| 110 |
+
pa.list_(pa.float32(), 1152)),
|
| 111 |
+
})
|
| 112 |
+
store.table("teacher_windows").append(
|
| 113 |
+
ft, kind="embeddings",
|
| 114 |
+
meta={"model": TEACHER, "written_by": "query cracking"})
|
| 115 |
+
for (s, a, _b), v in fresh.items():
|
| 116 |
+
cache[(s, a)] = v
|
| 117 |
+
miss_ms = (time.perf_counter() - t_miss) * 1e3
|
| 118 |
+
|
| 119 |
+
q_teacher = embed_text(text, model_id=TEACHER) # same checkpoint as
|
| 120 |
+
hits = [] # the cached vectors
|
| 121 |
+
for (s, a, b) in wins:
|
| 122 |
+
tv = cache.get((s, a))
|
| 123 |
+
if tv is None:
|
| 124 |
+
continue
|
| 125 |
+
hits.append({"stream": s, "t0": a, "t1": b,
|
| 126 |
+
"score": float(tv @ q_teacher), "teacher": True})
|
| 127 |
+
hits.sort(key=lambda h: -h["score"])
|
| 128 |
+
stats = {"method": "sharp", "shortlist": len(wins),
|
| 129 |
+
"cache_misses": len(missing),
|
| 130 |
+
"teacher_ms": round(miss_ms, 1),
|
| 131 |
+
"cache_size": len(cache),
|
| 132 |
+
"ms": round((time.perf_counter() - t_start) * 1e3, 1)}
|
| 133 |
+
return hits[:k], stats
|
python/elidedb/ctxprune.py
ADDED
|
@@ -0,0 +1,516 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Post-training cellular turnover for the context tower — PPO + reverse
|
| 2 |
+
attention, ported from FDNN_BrainModel/models/pruning.py.
|
| 3 |
+
|
| 4 |
+
WHY PRUNE AT ALL
|
| 5 |
+
----------------
|
| 6 |
+
Because the only thing this index is for is answering fast. Every channel the
|
| 7 |
+
tower carries is multiplied over T timesteps, for every window, at every
|
| 8 |
+
rebuild. Capacity that is not earning its keep is latency the user pays for
|
| 9 |
+
forever. And with a few hundred labelled windows, an over-provisioned tower
|
| 10 |
+
is also memorising — so apoptosis buys accuracy and speed at the same time,
|
| 11 |
+
which is exactly the regime cellular turnover was designed for.
|
| 12 |
+
|
| 13 |
+
WHAT CHANGED FROM FDNN, AND WHY
|
| 14 |
+
-------------------------------
|
| 15 |
+
FDNN prunes per (layer, neuron): every layer has its own aliveness mask. That
|
| 16 |
+
is right for a plain stack. This tower is a RESIDUAL stack, and in a residual
|
| 17 |
+
stack channel j is a wire that runs the whole depth — masking it in layer 2
|
| 18 |
+
alone does not remove it, because the skip connection keeps carrying it. So
|
| 19 |
+
the prune unit here is the residual-stream CHANNEL, decided once and applied
|
| 20 |
+
to every block.
|
| 21 |
+
|
| 22 |
+
That is not a weakening of the idea, it is the correct form of it for this
|
| 23 |
+
topology (the same structured channel pruning people apply to ResNets), and
|
| 24 |
+
it has a property per-layer masking does not: a dead channel can be
|
| 25 |
+
physically deleted. `compact()` rebuilds the tower without it, so the mask
|
| 26 |
+
turns into real wall-clock speed instead of a multiply by zero. Pruning that
|
| 27 |
+
does not shrink the matmul is not pruning, it is decoration.
|
| 28 |
+
|
| 29 |
+
THE THREE MEASURES (unchanged in spirit from FDNN)
|
| 30 |
+
--------------------------------------------------
|
| 31 |
+
utilization Δ validation RETRIEVAL loss when the channel is silenced.
|
| 32 |
+
Not weight magnitude — the thing we actually care about.
|
| 33 |
+
reverse attention softmax(-utilization): mass on the channels the network
|
| 34 |
+
attends AWAY from, i.e. the prune candidates.
|
| 35 |
+
PPO a clipped-PPO contextual bandit over per-channel features
|
| 36 |
+
that samples masks, scores them by
|
| 37 |
+
-(loss ratio) - λ·kept_fraction, and learns which channels
|
| 38 |
+
can actually go.
|
| 39 |
+
"""
|
| 40 |
+
from __future__ import annotations
|
| 41 |
+
|
| 42 |
+
import mlx.core as mx
|
| 43 |
+
import mlx.nn as nn
|
| 44 |
+
import mlx.optimizers as optim
|
| 45 |
+
import numpy as np
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
# ===========================================================================
|
| 49 |
+
# 0. Channel mask plumbing
|
| 50 |
+
# ===========================================================================
|
| 51 |
+
def n_channels(model):
|
| 52 |
+
return model.blocks[0].max_out_d
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def get_channel_mask(model):
|
| 56 |
+
return np.array(model.blocks[0].mask).copy()
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def set_channel_mask(model, mask):
|
| 60 |
+
"""One decision, applied to every block — see module docstring."""
|
| 61 |
+
m = np.asarray(mask, dtype=np.float32)
|
| 62 |
+
for b in model.turnover_layers():
|
| 63 |
+
b.set_active_mask(m)
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
# ===========================================================================
|
| 67 |
+
# 1. Utilization ("consumption")
|
| 68 |
+
# ===========================================================================
|
| 69 |
+
def measure_utilization(model, loss_of, verbose=False):
|
| 70 |
+
"""Per-channel ablation importance + the cheap descriptive stats.
|
| 71 |
+
|
| 72 |
+
`loss_of(model)` must return the validation retrieval loss as a float.
|
| 73 |
+
Ablation is the increase in that loss when a single channel is silenced
|
| 74 |
+
everywhere — the most directly interpretable notion of how much the tower
|
| 75 |
+
relies on it.
|
| 76 |
+
"""
|
| 77 |
+
model.set_training(False)
|
| 78 |
+
base = loss_of(model)
|
| 79 |
+
mask = get_channel_mask(model)
|
| 80 |
+
n = len(mask)
|
| 81 |
+
|
| 82 |
+
abl = np.zeros(n, dtype=np.float32)
|
| 83 |
+
for c in np.where(mask == 1.0)[0]:
|
| 84 |
+
probe = mask.copy()
|
| 85 |
+
probe[c] = 0.0
|
| 86 |
+
set_channel_mask(model, probe)
|
| 87 |
+
abl[c] = max(loss_of(model) - base, 0.0)
|
| 88 |
+
set_channel_mask(model, mask)
|
| 89 |
+
return {"mask": mask, "ablation": abl, "base_loss": base}
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def descriptive_stats(model, X):
|
| 93 |
+
"""Activation magnitude per channel (averaged over blocks) and the
|
| 94 |
+
downstream weight norm each channel feeds."""
|
| 95 |
+
model.set_training(False)
|
| 96 |
+
outs = []
|
| 97 |
+
h = mx.array(X)
|
| 98 |
+
for b in model.turnover_layers():
|
| 99 |
+
outs.append(np.array(mx.mean(mx.abs(b.neuron_outputs(h)), axis=(0, 1))))
|
| 100 |
+
h = b(h)
|
| 101 |
+
act = np.mean(np.stack(outs), axis=0)
|
| 102 |
+
|
| 103 |
+
n = n_channels(model)
|
| 104 |
+
wf = np.array(model.gru_f.Wx) # (3*hidden, n_hidden)
|
| 105 |
+
wb = np.array(model.gru_b.Wx)
|
| 106 |
+
hw = np.array(model.head.weight) # (d_out, n_hidden + 2*gru_hidden)
|
| 107 |
+
pw = np.array(model.pool.proj.weight)
|
| 108 |
+
down = (np.linalg.norm(wf, axis=0) + np.linalg.norm(wb, axis=0)
|
| 109 |
+
+ np.linalg.norm(hw[:, :n], axis=0)
|
| 110 |
+
+ np.linalg.norm(pw[:, :n], axis=0))
|
| 111 |
+
return act.astype(np.float32), down.astype(np.float32)
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
# ===========================================================================
|
| 115 |
+
# 2. Reverse attention
|
| 116 |
+
# ===========================================================================
|
| 117 |
+
def reverse_attention(importance, mask, temperature=1.0):
|
| 118 |
+
"""Forward attention concentrates on useful units; reverse attention
|
| 119 |
+
inverts it so the LEAST useful surface as prune candidates. Dead channels
|
| 120 |
+
score 0; scores over the live set sum to 1."""
|
| 121 |
+
importance = np.asarray(importance, dtype=np.float64)
|
| 122 |
+
alive = mask == 1.0
|
| 123 |
+
r = np.zeros_like(importance)
|
| 124 |
+
if alive.sum() == 0:
|
| 125 |
+
return r.astype(np.float32)
|
| 126 |
+
imp = importance[alive]
|
| 127 |
+
imp = (imp - imp.mean()) / (imp.std() + 1e-8)
|
| 128 |
+
logits = -imp / max(temperature, 1e-6)
|
| 129 |
+
logits -= logits.max()
|
| 130 |
+
e = np.exp(logits)
|
| 131 |
+
r[alive] = e / (e.sum() + 1e-12)
|
| 132 |
+
return r.astype(np.float32)
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
# ===========================================================================
|
| 136 |
+
# 3. PPO agent
|
| 137 |
+
# ===========================================================================
|
| 138 |
+
class PrunePolicy(nn.Module):
|
| 139 |
+
def __init__(self, feat_dim, hidden=32):
|
| 140 |
+
super().__init__()
|
| 141 |
+
self.fc1 = nn.Linear(feat_dim, hidden)
|
| 142 |
+
self.fc2 = nn.Linear(hidden, hidden)
|
| 143 |
+
self.pi_head = nn.Linear(hidden, 1)
|
| 144 |
+
self.v_head = nn.Linear(hidden, 1)
|
| 145 |
+
|
| 146 |
+
def __call__(self, feats):
|
| 147 |
+
h = nn.relu(self.fc1(feats))
|
| 148 |
+
h = nn.relu(self.fc2(h))
|
| 149 |
+
return self.pi_head(h)[:, 0], self.v_head(h)[:, 0]
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def _bernoulli_logp(actions, logits):
|
| 153 |
+
return actions * logits - mx.logaddexp(mx.zeros_like(logits), logits)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def _entropy(logits):
|
| 157 |
+
p = mx.sigmoid(logits)
|
| 158 |
+
return nn.softplus(logits) - p * logits
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
class PPOPruner:
|
| 162 |
+
"""Contextual bandit over prune masks of the residual stream.
|
| 163 |
+
|
| 164 |
+
One step per episode: observe a fixed per-channel feature matrix from the
|
| 165 |
+
frozen tower, emit a keep-probability per channel, sample a binary mask,
|
| 166 |
+
receive one scalar reward. Maximising it means "prune as much as possible
|
| 167 |
+
without inflating retrieval loss".
|
| 168 |
+
"""
|
| 169 |
+
|
| 170 |
+
def __init__(self, model, X, loss_of, sparsity_coef=0.35,
|
| 171 |
+
loss_tolerance=1.0, hidden=32, lr=3e-3, clip_eps=0.2,
|
| 172 |
+
entropy_coef=0.01, value_coef=0.5, seed=0):
|
| 173 |
+
self.model, self.loss_of = model, loss_of
|
| 174 |
+
self.sparsity_coef = sparsity_coef
|
| 175 |
+
self.loss_tolerance = loss_tolerance
|
| 176 |
+
self.clip_eps, self.entropy_coef = clip_eps, entropy_coef
|
| 177 |
+
self.value_coef = value_coef
|
| 178 |
+
np.random.seed(seed)
|
| 179 |
+
mx.random.seed(seed)
|
| 180 |
+
|
| 181 |
+
self.base_mask = get_channel_mask(model)
|
| 182 |
+
util = measure_utilization(model, loss_of)
|
| 183 |
+
self.base_loss = util["base_loss"]
|
| 184 |
+
self.ablation = util["ablation"]
|
| 185 |
+
act, down = descriptive_stats(model, X)
|
| 186 |
+
ra = reverse_attention(self.ablation, self.base_mask)
|
| 187 |
+
omega = model.blocks[0].omegas_per_neuron
|
| 188 |
+
|
| 189 |
+
self.candidates = np.where(self.base_mask == 1.0)[0]
|
| 190 |
+
f = np.stack([self.ablation[self.candidates], act[self.candidates],
|
| 191 |
+
down[self.candidates], ra[self.candidates],
|
| 192 |
+
omega[self.candidates]], axis=1).astype(np.float32)
|
| 193 |
+
mu, sd = f.mean(0, keepdims=True), f.std(0, keepdims=True) + 1e-6
|
| 194 |
+
f[:, :4] = (f[:, :4] - mu[:, :4]) / sd[:, :4]
|
| 195 |
+
self.feats = mx.array(f)
|
| 196 |
+
self.reverse_attn = ra
|
| 197 |
+
self.policy = PrunePolicy(f.shape[1], hidden=hidden)
|
| 198 |
+
self.opt = optim.Adam(learning_rate=lr)
|
| 199 |
+
self.history = {"iter": [], "reward": [], "kept_frac": [],
|
| 200 |
+
"loss_ratio": [], "entropy": []}
|
| 201 |
+
|
| 202 |
+
def _apply(self, action):
|
| 203 |
+
m = self.base_mask.copy()
|
| 204 |
+
m[self.candidates] = action
|
| 205 |
+
set_channel_mask(self.model, m)
|
| 206 |
+
|
| 207 |
+
def _reward(self, action):
|
| 208 |
+
self._apply(action)
|
| 209 |
+
loss = self.loss_of(self.model)
|
| 210 |
+
set_channel_mask(self.model, self.base_mask)
|
| 211 |
+
kept = float(action.mean()) if len(action) else 0.0
|
| 212 |
+
ratio = loss / (self.base_loss + 1e-12)
|
| 213 |
+
return (-self.loss_tolerance * ratio - self.sparsity_coef * kept,
|
| 214 |
+
kept, ratio)
|
| 215 |
+
|
| 216 |
+
def train(self, n_iters=40, episodes_per_iter=12, ppo_epochs=4,
|
| 217 |
+
verbose=True):
|
| 218 |
+
for it in range(n_iters):
|
| 219 |
+
logits, _ = self.policy(self.feats)
|
| 220 |
+
probs = np.array(mx.sigmoid(logits))
|
| 221 |
+
acts, rews, kfs, lrs = [], [], [], []
|
| 222 |
+
for _ in range(episodes_per_iter):
|
| 223 |
+
a = (np.random.uniform(size=probs.shape) < probs).astype(np.float32)
|
| 224 |
+
r, kf, lr_ = self._reward(a)
|
| 225 |
+
acts.append(a); rews.append(r); kfs.append(kf); lrs.append(lr_)
|
| 226 |
+
acts = np.array(acts, np.float32)
|
| 227 |
+
rews = np.array(rews, np.float32)
|
| 228 |
+
adv = (rews - rews.mean()) / (rews.std() + 1e-6)
|
| 229 |
+
|
| 230 |
+
acts_mx, adv_mx, rews_mx = mx.array(acts), mx.array(adv), mx.array(rews)
|
| 231 |
+
old_logits, _ = self.policy(self.feats)
|
| 232 |
+
old_logp = mx.array(np.array(
|
| 233 |
+
mx.sum(_bernoulli_logp(acts_mx, old_logits[None, :]), axis=1)))
|
| 234 |
+
for _ in range(ppo_epochs):
|
| 235 |
+
_, grads = self._loss_and_grad(acts_mx, old_logp, adv_mx, rews_mx)
|
| 236 |
+
self.opt.update(self.policy, grads)
|
| 237 |
+
mx.eval(self.policy.parameters(), self.opt.state)
|
| 238 |
+
|
| 239 |
+
ent = float(mx.mean(_entropy(logits)).item())
|
| 240 |
+
self.history["iter"].append(it)
|
| 241 |
+
self.history["reward"].append(float(rews.mean()))
|
| 242 |
+
self.history["kept_frac"].append(float(np.mean(kfs)))
|
| 243 |
+
self.history["loss_ratio"].append(float(np.mean(lrs)))
|
| 244 |
+
self.history["entropy"].append(ent)
|
| 245 |
+
if verbose and (it % 5 == 0 or it == n_iters - 1):
|
| 246 |
+
print(f" PPO {it:3d} | reward {rews.mean():+.3f} | kept "
|
| 247 |
+
f"{np.mean(kfs):.2f} | loss x{np.mean(lrs):.3f} | "
|
| 248 |
+
f"H {ent:.3f}", flush=True)
|
| 249 |
+
return self.history
|
| 250 |
+
|
| 251 |
+
def _loss_and_grad(self, actions, old_logp, adv, rewards):
|
| 252 |
+
def loss_fn(policy):
|
| 253 |
+
logits, values = policy(self.feats)
|
| 254 |
+
v = mx.mean(values)
|
| 255 |
+
logp = mx.sum(_bernoulli_logp(actions, logits[None, :]), axis=1)
|
| 256 |
+
ratio = mx.exp(logp - old_logp)
|
| 257 |
+
pl = -mx.mean(mx.minimum(ratio * adv,
|
| 258 |
+
mx.clip(ratio, 1 - self.clip_eps,
|
| 259 |
+
1 + self.clip_eps) * adv))
|
| 260 |
+
vl = mx.mean(mx.square(v - rewards))
|
| 261 |
+
return pl + self.value_coef * vl \
|
| 262 |
+
- self.entropy_coef * mx.mean(_entropy(logits))
|
| 263 |
+
return nn.value_and_grad(self.policy, loss_fn)(self.policy)
|
| 264 |
+
|
| 265 |
+
def apply_greedy(self, keep_threshold=0.5):
|
| 266 |
+
logits, _ = self.policy(self.feats)
|
| 267 |
+
probs = np.array(mx.sigmoid(logits))
|
| 268 |
+
m = self.base_mask.copy()
|
| 269 |
+
m[self.candidates] = (probs >= keep_threshold).astype(np.float32)
|
| 270 |
+
set_channel_mask(self.model, m)
|
| 271 |
+
return m, probs
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
# ===========================================================================
|
| 275 |
+
# 4. Neurogenesis — revive dead channels with fresh, spectrally diverse ones
|
| 276 |
+
# ===========================================================================
|
| 277 |
+
def neurogenesis(model, n_revive, seed=0):
|
| 278 |
+
"""Rebirth. New channels start near-silent (tiny w2) so they cannot shock
|
| 279 |
+
the forward pass, and are re-initialised with FINER's WIDE bias range —
|
| 280 |
+
without that they spawn as low-frequency clones and the population slowly
|
| 281 |
+
loses the spectral diversity that made the KAN basis worth having."""
|
| 282 |
+
rng = np.random.default_rng(seed)
|
| 283 |
+
mask = get_channel_mask(model)
|
| 284 |
+
dead = np.where(mask == 0.0)[0]
|
| 285 |
+
if len(dead) == 0 or n_revive <= 0:
|
| 286 |
+
return 0
|
| 287 |
+
revive = dead[:min(n_revive, len(dead))]
|
| 288 |
+
|
| 289 |
+
for b in model.turnover_layers():
|
| 290 |
+
w1 = np.array(b.w1) # (kernel, in_d, out*k)
|
| 291 |
+
b1 = np.array(b.b1)
|
| 292 |
+
w2 = np.array(b.w2)
|
| 293 |
+
ph = np.array(b.phases)
|
| 294 |
+
gs = np.array(b.gabor_s)
|
| 295 |
+
la = np.array(b.log_alpha)
|
| 296 |
+
om_full = np.repeat(b.omegas_per_neuron, b.k)
|
| 297 |
+
mean_om = float(b.omegas_per_neuron.mean())
|
| 298 |
+
lim = float(np.sqrt(6.0 / (b.kernel * b.in_d)) / mean_om)
|
| 299 |
+
for c in revive:
|
| 300 |
+
s, e = c * b.k, (c + 1) * b.k
|
| 301 |
+
w1[:, :, s:e] = rng.uniform(-lim, lim, (b.kernel, b.in_d, b.k))
|
| 302 |
+
b1[s:e] = rng.uniform(-b.bias_range, b.bias_range, b.k)
|
| 303 |
+
ph[s:e] = rng.uniform(0, 2 * np.pi, b.k)
|
| 304 |
+
gs[s:e] = rng.uniform(0.3, 1.5, b.k)
|
| 305 |
+
lw = np.log(np.clip(om_full[s:e], 1e-3, None))
|
| 306 |
+
la[s:e] = rng.uniform(np.zeros_like(lw), lw)
|
| 307 |
+
w2[c, :] = rng.standard_normal(b.k) * 1e-6
|
| 308 |
+
b.w1 = mx.array(w1.astype(np.float32))
|
| 309 |
+
b.b1 = mx.array(b1.astype(np.float32))
|
| 310 |
+
b.w2 = mx.array(w2.astype(np.float32))
|
| 311 |
+
b.phases = mx.array(ph.astype(np.float32))
|
| 312 |
+
b.gabor_s = mx.array(gs.astype(np.float32))
|
| 313 |
+
b.log_alpha = mx.array(la.astype(np.float32))
|
| 314 |
+
mask[revive] = 1.0
|
| 315 |
+
set_channel_mask(model, mask)
|
| 316 |
+
return len(revive)
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
# ===========================================================================
|
| 320 |
+
# 5. Compaction — where the mask becomes actual speed
|
| 321 |
+
# ===========================================================================
|
| 322 |
+
def compact(model):
|
| 323 |
+
"""Physically delete dead channels: return a smaller, equivalent tower.
|
| 324 |
+
|
| 325 |
+
Every dead channel is removed from each block's output slots, from every
|
| 326 |
+
consumer's input slots (both GRUs, the attention pool, the head), and from
|
| 327 |
+
the residual stream itself. Output is numerically equivalent to the masked
|
| 328 |
+
model but the matmuls are genuinely smaller.
|
| 329 |
+
"""
|
| 330 |
+
from .ctxtower import ContextTower
|
| 331 |
+
keep = np.where(get_channel_mask(model) == 1.0)[0]
|
| 332 |
+
n_new = len(keep)
|
| 333 |
+
if n_new == 0:
|
| 334 |
+
raise RuntimeError("every channel was pruned")
|
| 335 |
+
if n_new == n_channels(model):
|
| 336 |
+
return model, keep
|
| 337 |
+
|
| 338 |
+
cfg = dict(model.cfg)
|
| 339 |
+
cfg["n_hidden"] = int(n_new)
|
| 340 |
+
new = ContextTower(**cfg)
|
| 341 |
+
k = model.blocks[0].k
|
| 342 |
+
sub = np.concatenate([np.arange(c * k, (c + 1) * k) for c in keep])
|
| 343 |
+
|
| 344 |
+
for ob, nb in zip(model.blocks, new.blocks):
|
| 345 |
+
w1 = np.array(ob.w1)
|
| 346 |
+
if not ob.use_residual: # first block: in_d is d_in
|
| 347 |
+
nb.w1 = mx.array(w1[:, :, sub])
|
| 348 |
+
else: # in_d is the stream: slice both
|
| 349 |
+
nb.w1 = mx.array(w1[:, keep][:, :, sub])
|
| 350 |
+
nb.b1 = mx.array(np.array(ob.b1)[sub])
|
| 351 |
+
nb.phases = mx.array(np.array(ob.phases)[sub])
|
| 352 |
+
nb.gabor_s = mx.array(np.array(ob.gabor_s)[sub])
|
| 353 |
+
nb.log_alpha = mx.array(np.array(ob.log_alpha)[sub])
|
| 354 |
+
nb.w2 = mx.array(np.array(ob.w2)[keep])
|
| 355 |
+
# omegas is TRAINED, so it must be sliced from the live array — not
|
| 356 |
+
# rebuilt from omegas_per_neuron, which is the frozen design band and
|
| 357 |
+
# is stale the moment training starts. Rebuilding it silently reverted
|
| 358 |
+
# every learned frequency and moved val loss by 0.015.
|
| 359 |
+
nb.omegas = mx.array(np.array(ob.omegas)[sub])
|
| 360 |
+
nb.omegas_per_neuron = ob.omegas_per_neuron[keep] # design band only
|
| 361 |
+
nb.basis_types = mx.array(np.array(ob.basis_types)[sub])
|
| 362 |
+
nb.freeze(keys=["basis_types"], recurse=False)
|
| 363 |
+
nb.set_active_mask(np.ones(n_new, np.float32))
|
| 364 |
+
|
| 365 |
+
n_old = n_channels(model)
|
| 366 |
+
g = model.cfg["gru_hidden"]
|
| 367 |
+
for og, ng in ((model.gru_f, new.gru_f), (model.gru_b, new.gru_b)):
|
| 368 |
+
ng.Wx = mx.array(np.array(og.Wx)[:, keep])
|
| 369 |
+
ng.Wh = mx.array(np.array(og.Wh))
|
| 370 |
+
ng.b = mx.array(np.array(og.b))
|
| 371 |
+
ng.bhn = mx.array(np.array(og.bhn))
|
| 372 |
+
# consumers of `cat` = [stream | gru_f | gru_b]: keep the stream slice,
|
| 373 |
+
# carry the two GRU slices through untouched.
|
| 374 |
+
tail = np.arange(n_old, n_old + 2 * g)
|
| 375 |
+
cols = np.concatenate([keep, tail])
|
| 376 |
+
new.pool.proj.weight = mx.array(np.array(model.pool.proj.weight)[:, cols])
|
| 377 |
+
new.pool.proj.bias = mx.array(np.array(model.pool.proj.bias))
|
| 378 |
+
new.pool.score.weight = mx.array(np.array(model.pool.score.weight))
|
| 379 |
+
new.pool.score.bias = mx.array(np.array(model.pool.score.bias))
|
| 380 |
+
new.head.weight = mx.array(np.array(model.head.weight)[:, cols])
|
| 381 |
+
new.head.bias = mx.array(np.array(model.head.bias))
|
| 382 |
+
new.log_gamma = mx.array(np.array(model.log_gamma))
|
| 383 |
+
new.set_training(False)
|
| 384 |
+
return new, keep
|
| 385 |
+
|
| 386 |
+
|
| 387 |
+
# ===========================================================================
|
| 388 |
+
# 6. Re-settle
|
| 389 |
+
# ===========================================================================
|
| 390 |
+
def finetune(model, loss_and_grad_fn, epochs=150, lr=5e-4,
|
| 391 |
+
weight_decay=1e-3, loss_of=None, patience=40):
|
| 392 |
+
"""Short re-settling pass with aliveness frozen (see the mask gotcha).
|
| 393 |
+
|
| 394 |
+
Validation-aware on purpose. Survivors re-settle on the TRAIN set, and
|
| 395 |
+
with a couple of hundred windows an unguarded 400-epoch re-settle simply
|
| 396 |
+
re-memorises: measured, it took a pruned tower from 0.311 val back up to
|
| 397 |
+
0.357, undoing everything apoptosis had just gained. Keeping the best
|
| 398 |
+
validation state means re-settling can only help, never hurt.
|
| 399 |
+
"""
|
| 400 |
+
from .ctxtower import _clone_params, _load_params
|
| 401 |
+
opt = optim.AdamW(learning_rate=lr, weight_decay=weight_decay)
|
| 402 |
+
frozen = get_channel_mask(model)
|
| 403 |
+
best, best_w, since = float("inf"), None, 0
|
| 404 |
+
if loss_of is not None:
|
| 405 |
+
best, best_w = loss_of(model), _clone_params(model)
|
| 406 |
+
for _ in range(epochs):
|
| 407 |
+
model.set_training(True)
|
| 408 |
+
loss, grads = loss_and_grad_fn(model)
|
| 409 |
+
opt.update(model, grads)
|
| 410 |
+
set_channel_mask(model, frozen)
|
| 411 |
+
mx.eval(model.parameters(), opt.state)
|
| 412 |
+
if loss_of is None:
|
| 413 |
+
continue
|
| 414 |
+
v = loss_of(model)
|
| 415 |
+
if v < best - 1e-5:
|
| 416 |
+
best, best_w, since = v, _clone_params(model), 0
|
| 417 |
+
else:
|
| 418 |
+
since += 1
|
| 419 |
+
if since >= patience:
|
| 420 |
+
break
|
| 421 |
+
if best_w is not None:
|
| 422 |
+
_load_params(model, best_w)
|
| 423 |
+
set_channel_mask(model, frozen)
|
| 424 |
+
mx.eval(model.parameters())
|
| 425 |
+
model.set_training(False)
|
| 426 |
+
return best
|
| 427 |
+
|
| 428 |
+
|
| 429 |
+
# ===========================================================================
|
| 430 |
+
# 7. The full cycle
|
| 431 |
+
# ===========================================================================
|
| 432 |
+
def run_pruning_cycle(model, X, loss_of, loss_and_grad_fn, ppo_iters=40,
|
| 433 |
+
episodes_per_iter=12, sparsity_coef=0.35,
|
| 434 |
+
keep_threshold=0.5, finetune_epochs=150,
|
| 435 |
+
rebirth_fraction=0.5, select_tolerance=0.02,
|
| 436 |
+
seed=0, verbose=True):
|
| 437 |
+
"""apoptosis → re-settle → neurogenesis → re-settle, on a frozen tower.
|
| 438 |
+
|
| 439 |
+
The cycle is a SEARCH, so it returns the best point it visited, not the
|
| 440 |
+
last one. Rebirth is a bet that the data can support more capacity; on a
|
| 441 |
+
small corpus that bet often loses, and silently shipping the final state
|
| 442 |
+
would hand back a tower worse than the one we started from.
|
| 443 |
+
"""
|
| 444 |
+
from .ctxtower import _clone_params, _load_params
|
| 445 |
+
rec = {"stages": []}
|
| 446 |
+
saved = [] # (stage, loss, alive, params, mask) for every stage
|
| 447 |
+
|
| 448 |
+
def snap(stage):
|
| 449 |
+
u = measure_utilization(model, loss_of)
|
| 450 |
+
alive = int(u["mask"].sum())
|
| 451 |
+
if verbose:
|
| 452 |
+
print(f"[{stage}] val loss {u['base_loss']:.4f} | "
|
| 453 |
+
f"alive {alive}/{len(u['mask'])}", flush=True)
|
| 454 |
+
rec["stages"].append({"stage": stage, "loss": u["base_loss"],
|
| 455 |
+
"alive": alive,
|
| 456 |
+
"ablation": u["ablation"].tolist(),
|
| 457 |
+
"mask": u["mask"].tolist()})
|
| 458 |
+
saved.append((stage, u["base_loss"], alive,
|
| 459 |
+
_clone_params(model), u["mask"].copy()))
|
| 460 |
+
return u
|
| 461 |
+
|
| 462 |
+
snap("before")
|
| 463 |
+
if verbose:
|
| 464 |
+
print("\n== PPO reverse-attention pruning search ==", flush=True)
|
| 465 |
+
pruner = PPOPruner(model, X, loss_of, sparsity_coef=sparsity_coef, seed=seed)
|
| 466 |
+
rec["ppo_history"] = pruner.train(n_iters=ppo_iters,
|
| 467 |
+
episodes_per_iter=episodes_per_iter,
|
| 468 |
+
verbose=verbose)
|
| 469 |
+
_, probs = pruner.apply_greedy(keep_threshold)
|
| 470 |
+
rec["keep_probs"] = probs.tolist()
|
| 471 |
+
rec["reverse_attention"] = pruner.reverse_attn.tolist()
|
| 472 |
+
snap("after_apoptosis")
|
| 473 |
+
|
| 474 |
+
if finetune_epochs:
|
| 475 |
+
if verbose:
|
| 476 |
+
print("\n== re-settling survivors ==", flush=True)
|
| 477 |
+
finetune(model, loss_and_grad_fn, epochs=finetune_epochs,
|
| 478 |
+
loss_of=loss_of)
|
| 479 |
+
snap("after_finetune")
|
| 480 |
+
|
| 481 |
+
if rebirth_fraction > 0:
|
| 482 |
+
dead = int((get_channel_mask(model) == 0.0).sum())
|
| 483 |
+
born = neurogenesis(model, int(round(rebirth_fraction * dead)), seed=seed)
|
| 484 |
+
if verbose:
|
| 485 |
+
print(f"\n== neurogenesis: {born} channels reborn ==", flush=True)
|
| 486 |
+
if born:
|
| 487 |
+
snap("after_rebirth")
|
| 488 |
+
if finetune_epochs:
|
| 489 |
+
finetune(model, loss_and_grad_fn,
|
| 490 |
+
epochs=finetune_epochs, loss_of=loss_of)
|
| 491 |
+
snap("after_rebirth_finetune")
|
| 492 |
+
|
| 493 |
+
# ---- operating-point selection ----------------------------------------
|
| 494 |
+
# Not "lowest loss" — that rule can never prune, because the unpruned
|
| 495 |
+
# tower is usually the most accurate one and the whole exercise then
|
| 496 |
+
# returns its own input. The question a database actually asks is: what
|
| 497 |
+
# is the CHEAPEST tower I can serve without giving up measurable quality?
|
| 498 |
+
# So: take the best loss seen, allow `select_tolerance` relative slack,
|
| 499 |
+
# and among everything inside that band keep the fewest channels.
|
| 500 |
+
floor = min(s[1] for s in saved)
|
| 501 |
+
budget = floor * (1.0 + select_tolerance)
|
| 502 |
+
eligible = [s for s in saved if s[1] <= budget]
|
| 503 |
+
stage, loss, alive, params, mask = min(eligible, key=lambda s: (s[2], s[1]))
|
| 504 |
+
_load_params(model, params)
|
| 505 |
+
set_channel_mask(model, mask)
|
| 506 |
+
mx.eval(model.parameters())
|
| 507 |
+
if verbose:
|
| 508 |
+
print(f"\n== operating point: {stage} — {alive} channels, val loss "
|
| 509 |
+
f"{loss:.4f} (best seen {floor:.4f}, "
|
| 510 |
+
f"budget +{select_tolerance:.0%}) ==", flush=True)
|
| 511 |
+
rec["selected"] = {"stage": stage, "loss": loss, "channels": alive,
|
| 512 |
+
"best_loss_seen": floor,
|
| 513 |
+
"tolerance": select_tolerance,
|
| 514 |
+
"candidates": [{"stage": s[0], "loss": s[1],
|
| 515 |
+
"channels": s[2]} for s in saved]}
|
| 516 |
+
return rec
|
python/elidedb/ctxtower.py
ADDED
|
@@ -0,0 +1,554 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The context tower — a custom SigLIP video tower over frozen frame features.
|
| 2 |
+
|
| 3 |
+
WHAT IT IS
|
| 4 |
+
----------
|
| 5 |
+
SigLIP's image tower is a frozen spatial encoder. This module supplies the
|
| 6 |
+
*other* half of an R(2+1)D-style factorisation: a learned TEMPORAL tower that
|
| 7 |
+
consumes the sequence of per-frame SigLIP vectors in a window and emits one
|
| 8 |
+
vector aligned with SigLIP TEXT embeddings.
|
| 9 |
+
|
| 10 |
+
frame vectors (T, 1152) [frozen SigLIP image tower — the "2D"]
|
| 11 |
+
| PCA (frozen, no params)
|
| 12 |
+
(T, d)
|
| 13 |
+
| dilated temporal conv, taps learned ← the CNN
|
| 14 |
+
| each output neuron = KAN sum over k
|
| 15 |
+
| heterogeneous bases (FINER/Gabor/poly) ← FDNN thesis 1
|
| 16 |
+
(T, n)
|
| 17 |
+
| bi-GRU ← the RNN
|
| 18 |
+
(T, 2h)
|
| 19 |
+
| attention pooling over time
|
| 20 |
+
(2h,)
|
| 21 |
+
| head -> PCA^T of the caption space
|
| 22 |
+
(1152,) ~ a SigLIP text embedding
|
| 23 |
+
|
| 24 |
+
WHY A CONV *AND* A GRU
|
| 25 |
+
----------------------
|
| 26 |
+
They fail differently, which is the only good reason to have both. The dilated
|
| 27 |
+
conv is order-aware but translation-equivariant: it detects "a thing moved
|
| 28 |
+
left-to-right" wherever in the window it happens, with a fixed receptive
|
| 29 |
+
field. The GRU is order-aware and unbounded: it can carry "the car was
|
| 30 |
+
stationary the whole time, then braked" across the window, which no
|
| 31 |
+
fixed-width kernel expresses. CLIP4Clip (arXiv 2104.08860) measured exactly
|
| 32 |
+
this axis on frozen CLIP features — meanP vs seqLSTM vs seqTransf — and found
|
| 33 |
+
learned temporal aggregation helps once the target actually depends on order.
|
| 34 |
+
Mean pooling, what `embeddings` does today, is the degenerate case of both.
|
| 35 |
+
|
| 36 |
+
TSM (Lin et al., ICCV 2019) was the alternative for the conv slot: shifting
|
| 37 |
+
channels along time is free. It was not used because its shift is a fixed
|
| 38 |
+
±1 tap; here the whole point is that the taps are *frequency-selective*, and
|
| 39 |
+
a fixed shift cannot express that.
|
| 40 |
+
|
| 41 |
+
WHY THE BASES ARE THE FDNN BASES
|
| 42 |
+
--------------------------------
|
| 43 |
+
A window's feature trajectory is a signal, and these bases were built for
|
| 44 |
+
signals. Applied after a temporal convolution:
|
| 45 |
+
- a Gabor sub-function is a temporal wavelet — a burst detector, localised
|
| 46 |
+
in time (a door opening, a brake light);
|
| 47 |
+
- a FINER sub-function is a variable-period oscillator — periodic motion
|
| 48 |
+
(gait, wipers, a turning wheel);
|
| 49 |
+
- a polynomial-phase sub-function is a chirp — monotonic acceleration
|
| 50 |
+
(a vehicle pulling away, a zoom).
|
| 51 |
+
The frequency bands then partition the TEMPORAL spectrum instead of every
|
| 52 |
+
neuron competing for it: slow = scene identity, mid = object motion, fast =
|
| 53 |
+
transitions. Each neuron being a sum over k of these is FDNN thesis 1, and
|
| 54 |
+
here it buys genuine expressiveness rather than decoration.
|
| 55 |
+
|
| 56 |
+
IDENTITY-SAFE INIT
|
| 57 |
+
------------------
|
| 58 |
+
The head is initialised so the tower emits the mean caption vector for every
|
| 59 |
+
window. That is "I know nothing" — a legal, centred prediction — so training
|
| 60 |
+
can only add information, and a half-trained tower can never be worse than
|
| 61 |
+
the corpus prior. Same discipline as FDNN's identity-init residual.
|
| 62 |
+
"""
|
| 63 |
+
from __future__ import annotations
|
| 64 |
+
|
| 65 |
+
import json
|
| 66 |
+
from pathlib import Path
|
| 67 |
+
|
| 68 |
+
import mlx.core as mx
|
| 69 |
+
import mlx.nn as nn
|
| 70 |
+
import numpy as np
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
# ===========================================================================
|
| 74 |
+
# Temporal hybrid layer — FDNN's HybridBiomimeticLayer with a time axis
|
| 75 |
+
# ===========================================================================
|
| 76 |
+
class TemporalHybridLayer(nn.Module):
|
| 77 |
+
"""Dilated temporal conv whose output channels are KAN sub-network neurons.
|
| 78 |
+
|
| 79 |
+
Shapes: (B, T, in_d) -> (B, T, max_out_d).
|
| 80 |
+
|
| 81 |
+
The convolution and the sub-function bank are fused: the conv produces
|
| 82 |
+
`max_out_d * k` pre-activations per timestep, the bases are applied
|
| 83 |
+
elementwise, and the k sub-functions of each neuron are summed (KAN).
|
| 84 |
+
So one neuron is not "a channel" — it is a little ensemble of temporal
|
| 85 |
+
filters that disagree about what shape of motion to look for.
|
| 86 |
+
|
| 87 |
+
`omega_bands` are TEMPORAL frequencies here. They are an order of
|
| 88 |
+
magnitude smaller than FDNN's coordinate-network defaults on purpose: the
|
| 89 |
+
input is a PCA of unit-norm embeddings, so |h| is O(1) rather than O(100),
|
| 90 |
+
and reusing omega=200 would put every neuron in the chaotic regime where
|
| 91 |
+
gradients are noise.
|
| 92 |
+
"""
|
| 93 |
+
|
| 94 |
+
def __init__(self, in_d, max_out_d, initial_active=None, k_width=4,
|
| 95 |
+
kernel=3, dilation=1, omega_bands=(2.0, 6.0, 18.0),
|
| 96 |
+
band_fractions=(0.34, 0.33, 0.33), bias_range=2.0,
|
| 97 |
+
use_residual=True, dropout_p=0.0, seed=0):
|
| 98 |
+
super().__init__()
|
| 99 |
+
rng = np.random.default_rng(seed)
|
| 100 |
+
self.in_d = in_d
|
| 101 |
+
self.max_out_d = max_out_d
|
| 102 |
+
self.k = k_width
|
| 103 |
+
self.kernel = kernel
|
| 104 |
+
self.dilation = dilation
|
| 105 |
+
self.bias_range = bias_range
|
| 106 |
+
self.use_residual = use_residual and (in_d == max_out_d)
|
| 107 |
+
self.dropout_p = dropout_p
|
| 108 |
+
self._training = True
|
| 109 |
+
|
| 110 |
+
# ---- frequency-banded omega per neuron (temporal spectrum split) ----
|
| 111 |
+
omegas = []
|
| 112 |
+
for om, fr in zip(omega_bands, band_fractions):
|
| 113 |
+
omegas.extend([om] * int(round(fr * max_out_d)))
|
| 114 |
+
omegas = (omegas + [omega_bands[-1]] * max_out_d)[:max_out_d]
|
| 115 |
+
self.omegas_per_neuron = np.array(omegas, dtype=np.float32)
|
| 116 |
+
om_exp = np.repeat(self.omegas_per_neuron, k_width)
|
| 117 |
+
self.omegas = mx.array(om_exp)
|
| 118 |
+
|
| 119 |
+
# ---- heterogeneous basis assignment: 50% FINER, 25% Gabor, 25% poly
|
| 120 |
+
half, quarter = max(k_width // 2, 1), max(k_width // 4, 1)
|
| 121 |
+
per = np.array([0] * half + [1] * quarter
|
| 122 |
+
+ [3] * max(k_width - half - quarter, 0),
|
| 123 |
+
dtype=np.int32)[:k_width]
|
| 124 |
+
if per.size < k_width:
|
| 125 |
+
per = np.concatenate([per, np.zeros(k_width - per.size, np.int32)])
|
| 126 |
+
self.basis_types = mx.array(np.tile(per, max_out_d).astype(np.int32))
|
| 127 |
+
|
| 128 |
+
# ---- the temporal convolution: (kernel, in_d, max_out_d*k) ----------
|
| 129 |
+
# Fan-in is kernel*in_d, and the SIREN convention divides by omega so
|
| 130 |
+
# the pre-activation lands in the basis's useful range at init.
|
| 131 |
+
mean_om = float(self.omegas_per_neuron.mean())
|
| 132 |
+
limit = float(np.sqrt(6.0 / (kernel * in_d)) / mean_om)
|
| 133 |
+
self.w1 = mx.array(rng.uniform(
|
| 134 |
+
-limit, limit, (kernel, in_d, max_out_d * k_width)).astype(np.float32))
|
| 135 |
+
self.b1 = mx.array(rng.uniform(
|
| 136 |
+
-bias_range, bias_range, (max_out_d * k_width,)).astype(np.float32))
|
| 137 |
+
self.phases = mx.array(rng.uniform(
|
| 138 |
+
0, 2 * np.pi, (max_out_d * k_width,)).astype(np.float32))
|
| 139 |
+
self.gabor_s = mx.array(rng.uniform(
|
| 140 |
+
0.3, 1.5, (max_out_d * k_width,)).astype(np.float32))
|
| 141 |
+
log_om = np.log(np.clip(om_exp, 1e-3, None))
|
| 142 |
+
self.log_alpha = mx.array(
|
| 143 |
+
rng.uniform(0.0, log_om).astype(np.float32))
|
| 144 |
+
|
| 145 |
+
w2s = float(np.sqrt(6.0 / (max_out_d * k_width)))
|
| 146 |
+
self.w2 = mx.array(rng.uniform(
|
| 147 |
+
-w2s, w2s, (max_out_d, k_width)).astype(np.float32))
|
| 148 |
+
|
| 149 |
+
m = np.zeros((max_out_d,), dtype=np.float32)
|
| 150 |
+
m[:(initial_active or max_out_d)] = 1.0
|
| 151 |
+
self.mask = mx.array(m)
|
| 152 |
+
|
| 153 |
+
# basis_types is a CATEGORICAL SELECTOR, not a weight. Every mx.array
|
| 154 |
+
# attribute joins the parameter tree, so without this the optimizer
|
| 155 |
+
# takes gradient steps on it: int32 gets promoted to float32 and the
|
| 156 |
+
# values drift (measured: 1.0 -> 0.9992). Dispatch is
|
| 157 |
+
# `where(basis_types == 1, gabor, ...)`, so a drifted selector matches
|
| 158 |
+
# nothing and every Gabor neuron silently falls through to the poly
|
| 159 |
+
# branch — the heterogeneous basis quietly stops being heterogeneous.
|
| 160 |
+
# Freezing keeps it out of trainable_parameters() entirely.
|
| 161 |
+
self.freeze(keys=["basis_types"], recurse=False)
|
| 162 |
+
|
| 163 |
+
# -- temporal gather: centred taps so the window is read bidirectionally --
|
| 164 |
+
def _conv(self, x):
|
| 165 |
+
"""x: (B, T, in_d) -> (B, T, max_out_d*k)."""
|
| 166 |
+
T = x.shape[1]
|
| 167 |
+
off0 = (self.kernel - 1) // 2
|
| 168 |
+
acc = None
|
| 169 |
+
for t_i in range(self.kernel):
|
| 170 |
+
shift = (t_i - off0) * self.dilation
|
| 171 |
+
if shift == 0:
|
| 172 |
+
xs = x
|
| 173 |
+
elif shift > 0: # look forward, edge-pad tail
|
| 174 |
+
xs = mx.concatenate(
|
| 175 |
+
[x[:, shift:, :],
|
| 176 |
+
mx.repeat(x[:, -1:, :], shift, axis=1)], axis=1)
|
| 177 |
+
else: # look back, edge-pad head
|
| 178 |
+
s = -shift
|
| 179 |
+
xs = mx.concatenate(
|
| 180 |
+
[mx.repeat(x[:, :1, :], s, axis=1),
|
| 181 |
+
x[:, :T - s, :]], axis=1)
|
| 182 |
+
term = xs @ self.w1[t_i]
|
| 183 |
+
acc = term if acc is None else acc + term
|
| 184 |
+
return acc + self.b1
|
| 185 |
+
|
| 186 |
+
def _bases(self, h):
|
| 187 |
+
omega_h = self.omegas * h
|
| 188 |
+
h_sq = h * h
|
| 189 |
+
alpha = mx.exp(self.log_alpha)
|
| 190 |
+
finer = mx.sin(self.omegas * (mx.abs(h) + 1.0) * h + self.phases)
|
| 191 |
+
gabor = mx.exp(-(self.gabor_s ** 2) * h_sq) * mx.sin(omega_h + self.phases)
|
| 192 |
+
sine = mx.sin(omega_h + self.phases)
|
| 193 |
+
poly = mx.sin(alpha * h_sq + omega_h + self.phases)
|
| 194 |
+
return mx.where(self.basis_types == 0, finer,
|
| 195 |
+
mx.where(self.basis_types == 1, gabor,
|
| 196 |
+
mx.where(self.basis_types == 2, sine, poly)))
|
| 197 |
+
|
| 198 |
+
def neuron_outputs(self, x):
|
| 199 |
+
"""(B, T, max_out_d) per-neuron signal BEFORE mask and residual.
|
| 200 |
+
Deterministic regardless of mode — every pruning metric reads this."""
|
| 201 |
+
acts = self._bases(self._conv(x))
|
| 202 |
+
acts = acts.reshape(acts.shape[0], acts.shape[1], self.max_out_d, self.k)
|
| 203 |
+
return mx.sum(acts * self.w2, axis=-1)
|
| 204 |
+
|
| 205 |
+
def set_active_mask(self, mask_np):
|
| 206 |
+
self.mask = mx.array(np.asarray(mask_np, dtype=np.float32))
|
| 207 |
+
|
| 208 |
+
def __call__(self, x):
|
| 209 |
+
acts = self._bases(self._conv(x))
|
| 210 |
+
acts = acts.reshape(acts.shape[0], acts.shape[1], self.max_out_d, self.k)
|
| 211 |
+
if self._training and self.dropout_p > 0.0:
|
| 212 |
+
keep = 1.0 - self.dropout_p
|
| 213 |
+
drop = (mx.random.uniform(0.0, 1.0, acts.shape) < keep).astype(acts.dtype)
|
| 214 |
+
acts = acts * drop / keep
|
| 215 |
+
out = mx.sum(acts * self.w2, axis=-1) * self.mask
|
| 216 |
+
return out + x if self.use_residual else out
|
| 217 |
+
|
| 218 |
+
|
| 219 |
+
# ===========================================================================
|
| 220 |
+
# Attention pooling over time
|
| 221 |
+
# ===========================================================================
|
| 222 |
+
class AttnPool(nn.Module):
|
| 223 |
+
"""Learned soft-argmax over the window.
|
| 224 |
+
|
| 225 |
+
Mean pooling says every instant matters equally, which is wrong for a
|
| 226 |
+
driving clip where 1.8 s is empty road and 0.2 s is the pedestrian
|
| 227 |
+
stepping off the kerb. Attention lets the tower spend its output on the
|
| 228 |
+
part of the window that carries the event.
|
| 229 |
+
"""
|
| 230 |
+
|
| 231 |
+
def __init__(self, d, hidden=32):
|
| 232 |
+
super().__init__()
|
| 233 |
+
self.proj = nn.Linear(d, hidden)
|
| 234 |
+
self.score = nn.Linear(hidden, 1)
|
| 235 |
+
|
| 236 |
+
def __call__(self, h): # (B, T, d) -> (B, d)
|
| 237 |
+
a = self.score(mx.tanh(self.proj(h))) # (B, T, 1)
|
| 238 |
+
w = mx.softmax(a, axis=1)
|
| 239 |
+
return mx.sum(h * w, axis=1), w[..., 0]
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
# ===========================================================================
|
| 243 |
+
# The tower
|
| 244 |
+
# ===========================================================================
|
| 245 |
+
class ContextTower(nn.Module):
|
| 246 |
+
def __init__(self, d_in=48, n_hidden=48, k_width=4, gru_hidden=32,
|
| 247 |
+
d_out=32, kernel=3, dilations=(1, 2), dropout_p=0.1,
|
| 248 |
+
input_noise_std=0.02, omega_bands=(2.0, 6.0, 18.0), seed=0):
|
| 249 |
+
super().__init__()
|
| 250 |
+
self.cfg = dict(d_in=d_in, n_hidden=n_hidden, k_width=k_width,
|
| 251 |
+
gru_hidden=gru_hidden, d_out=d_out, kernel=kernel,
|
| 252 |
+
dilations=list(dilations), dropout_p=dropout_p,
|
| 253 |
+
input_noise_std=input_noise_std,
|
| 254 |
+
omega_bands=list(omega_bands), seed=seed)
|
| 255 |
+
self.input_noise_std = input_noise_std
|
| 256 |
+
self._training = True
|
| 257 |
+
|
| 258 |
+
self.blocks = []
|
| 259 |
+
for i, dil in enumerate(dilations):
|
| 260 |
+
self.blocks.append(TemporalHybridLayer(
|
| 261 |
+
in_d=d_in if i == 0 else n_hidden,
|
| 262 |
+
max_out_d=n_hidden, k_width=k_width, kernel=kernel,
|
| 263 |
+
dilation=dil, omega_bands=omega_bands,
|
| 264 |
+
use_residual=(i > 0), dropout_p=dropout_p, seed=seed + i))
|
| 265 |
+
self.gru_f = nn.GRU(n_hidden, gru_hidden)
|
| 266 |
+
self.gru_b = nn.GRU(n_hidden, gru_hidden)
|
| 267 |
+
self.pool = AttnPool(n_hidden + 2 * gru_hidden)
|
| 268 |
+
self.head = nn.Linear(n_hidden + 2 * gru_hidden, d_out)
|
| 269 |
+
# Head starts at exactly zero: the tower's first prediction is the
|
| 270 |
+
# corpus prior (see module docstring). Bias is zero too — the prior
|
| 271 |
+
# itself is added outside, in the frozen decode basis.
|
| 272 |
+
self.head.weight = mx.zeros(self.head.weight.shape)
|
| 273 |
+
self.head.bias = mx.zeros((d_out,))
|
| 274 |
+
self.log_gamma = mx.array(np.float32(np.log(0.5)))
|
| 275 |
+
|
| 276 |
+
def set_training(self, mode: bool):
|
| 277 |
+
self._training = mode
|
| 278 |
+
for b in self.blocks:
|
| 279 |
+
b._training = mode
|
| 280 |
+
return self
|
| 281 |
+
|
| 282 |
+
def turnover_layers(self):
|
| 283 |
+
return list(self.blocks)
|
| 284 |
+
|
| 285 |
+
def trunk(self, x):
|
| 286 |
+
"""(B, T, d_in) -> (B, n_hidden + 2*gru_hidden) pooled features."""
|
| 287 |
+
if self._training and self.input_noise_std > 0:
|
| 288 |
+
x = x + mx.random.normal(x.shape) * self.input_noise_std
|
| 289 |
+
h = x
|
| 290 |
+
for b in self.blocks:
|
| 291 |
+
h = b(h)
|
| 292 |
+
f = self.gru_f(h)
|
| 293 |
+
b_rev = self.gru_b(h[:, ::-1, :])[:, ::-1, :]
|
| 294 |
+
cat = mx.concatenate([h, f, b_rev], axis=-1)
|
| 295 |
+
pooled, attn = self.pool(cat)
|
| 296 |
+
# NOTE: deliberately not stashed on `self`. An mx.array assigned to a
|
| 297 |
+
# Module attribute joins the parameter tree, and the optimizer then
|
| 298 |
+
# tries to take a step on it (KeyError: 'last_attn'). Use
|
| 299 |
+
# `attention_weights()` when you want to inspect it.
|
| 300 |
+
return pooled
|
| 301 |
+
|
| 302 |
+
def attention_weights(self, x):
|
| 303 |
+
"""(B, T) softmax weights — which instants the tower spent itself on."""
|
| 304 |
+
h = x
|
| 305 |
+
for b in self.blocks:
|
| 306 |
+
h = b(h)
|
| 307 |
+
f = self.gru_f(h)
|
| 308 |
+
b_rev = self.gru_b(h[:, ::-1, :])[:, ::-1, :]
|
| 309 |
+
return self.pool(mx.concatenate([h, f, b_rev], axis=-1))[1]
|
| 310 |
+
|
| 311 |
+
def __call__(self, x):
|
| 312 |
+
"""-> (B, d_out) latent coordinates in the caption PCA basis."""
|
| 313 |
+
return self.head(self.trunk(x)) * mx.exp(self.log_gamma)
|
| 314 |
+
|
| 315 |
+
|
| 316 |
+
# ===========================================================================
|
| 317 |
+
# Codec: frozen PCA in, frozen PCA out. Zero learned parameters either side.
|
| 318 |
+
# ===========================================================================
|
| 319 |
+
class ContextCodec:
|
| 320 |
+
"""Frozen PCA on the INPUT side of the tower.
|
| 321 |
+
|
| 322 |
+
We have a few hundred labelled windows. A learned 1152->d encoder would be
|
| 323 |
+
~150k parameters fitted from ~150 examples, which is not learning, it is
|
| 324 |
+
memorising. PCA is the optimal linear compressor under reconstruction
|
| 325 |
+
error and costs zero training examples, so all the sample budget goes to
|
| 326 |
+
the temporal dynamics — the only part mean pooling cannot already do.
|
| 327 |
+
|
| 328 |
+
The OUTPUT side is not a PCA of SigLIP text embeddings. That was measured
|
| 329 |
+
and it lost: ranking a query against a predicted caption *embedding*
|
| 330 |
+
scored +0.218 against the VLM judge, below plain appearance search at
|
| 331 |
+
+0.276, because SigLIP is trained for image-text similarity and its text
|
| 332 |
+
tower is not calibrated for text-to-text comparison. The output space is
|
| 333 |
+
now `CaptionSpace` — LSA over the caption text itself — which measured
|
| 334 |
+
+0.339. See docs/CONTEXT.md.
|
| 335 |
+
"""
|
| 336 |
+
|
| 337 |
+
def __init__(self, P_in, mu_in):
|
| 338 |
+
self.P_in, self.mu_in = P_in, mu_in # (d_in, D), (D,)
|
| 339 |
+
|
| 340 |
+
@staticmethod
|
| 341 |
+
def fit(frame_vecs, d_in=48):
|
| 342 |
+
from sklearn.decomposition import PCA
|
| 343 |
+
d_in = min(d_in, *frame_vecs.shape)
|
| 344 |
+
pi = PCA(n_components=d_in, random_state=0).fit(frame_vecs)
|
| 345 |
+
return ContextCodec(pi.components_.astype(np.float32),
|
| 346 |
+
pi.mean_.astype(np.float32))
|
| 347 |
+
|
| 348 |
+
def encode(self, seq): # (T, D) -> (T, d_in)
|
| 349 |
+
return (seq - self.mu_in) @ self.P_in.T
|
| 350 |
+
|
| 351 |
+
def save(self, path):
|
| 352 |
+
np.savez(path, P_in=self.P_in, mu_in=self.mu_in)
|
| 353 |
+
|
| 354 |
+
@staticmethod
|
| 355 |
+
def load(path):
|
| 356 |
+
z = np.load(path)
|
| 357 |
+
return ContextCodec(z["P_in"], z["mu_in"])
|
| 358 |
+
|
| 359 |
+
|
| 360 |
+
# ===========================================================================
|
| 361 |
+
# SigLIP's own loss
|
| 362 |
+
# ===========================================================================
|
| 363 |
+
def siglip_loss(v, u, ignore=None, log_t=None, bias=None):
|
| 364 |
+
"""Pairwise sigmoid loss (Zhai et al., ICCV 2023, arXiv 2303.15343).
|
| 365 |
+
|
| 366 |
+
Chosen over softmax-InfoNCE for one concrete reason: the paper's own
|
| 367 |
+
ablation shows sigmoid wins below ~16k batch, and our batch is the whole
|
| 368 |
+
labelled set — a couple of hundred. Softmax normalises over the batch, so
|
| 369 |
+
at this size the partition function is estimated from almost nothing.
|
| 370 |
+
Sigmoid treats every pair as an independent binary problem and never
|
| 371 |
+
needs that global view.
|
| 372 |
+
|
| 373 |
+
`ignore` masks pairs that must not be counted as negatives. Windows slide
|
| 374 |
+
with 75% overlap, so window i+1 genuinely depicts the same moment as
|
| 375 |
+
window i; calling it a negative would teach the tower to separate
|
| 376 |
+
identical content. Those pairs are dropped, not down-weighted.
|
| 377 |
+
"""
|
| 378 |
+
t = mx.exp(log_t)
|
| 379 |
+
logits = t * (v @ u.T) + bias
|
| 380 |
+
n = v.shape[0]
|
| 381 |
+
y = 2.0 * mx.eye(n) - 1.0 # +1 diagonal, -1 elsewhere
|
| 382 |
+
z = y * logits
|
| 383 |
+
# -log sigmoid(z) == softplus(-z), stable for both signs
|
| 384 |
+
per_pair = nn.softplus(-z)
|
| 385 |
+
if ignore is not None:
|
| 386 |
+
keep = 1.0 - ignore
|
| 387 |
+
return mx.sum(per_pair * keep) / mx.maximum(mx.sum(keep), 1.0)
|
| 388 |
+
return mx.mean(per_pair)
|
| 389 |
+
|
| 390 |
+
|
| 391 |
+
def overlap_mask(windows):
|
| 392 |
+
"""1.0 where two windows are the same stream and overlap in time."""
|
| 393 |
+
n = len(windows)
|
| 394 |
+
m = np.zeros((n, n), dtype=np.float32)
|
| 395 |
+
st = [w[0] for w in windows]
|
| 396 |
+
t0 = np.array([w[1] for w in windows], dtype=np.int64)
|
| 397 |
+
t1 = np.array([w[2] for w in windows], dtype=np.int64)
|
| 398 |
+
for i in range(n):
|
| 399 |
+
same = np.array([s == st[i] for s in st])
|
| 400 |
+
ov = same & (t0 < t1[i]) & (t1 > t0[i])
|
| 401 |
+
m[i] = ov.astype(np.float32)
|
| 402 |
+
np.fill_diagonal(m, 0.0) # the positive stays a positive
|
| 403 |
+
return m
|
| 404 |
+
|
| 405 |
+
|
| 406 |
+
# ===========================================================================
|
| 407 |
+
# Training
|
| 408 |
+
# ===========================================================================
|
| 409 |
+
def train_tower(seqs, targets, windows, val_idx, cfg=None, epochs=400,
|
| 410 |
+
lr=3e-3, weight_decay=1e-3, align_w=0.3, verbose=True,
|
| 411 |
+
seed=0):
|
| 412 |
+
"""Fit the tower. `seqs` are (T, d_in) encoded sequences, `targets` are
|
| 413 |
+
(d_out,) caption coordinates in the output basis, both already codec-mapped.
|
| 414 |
+
|
| 415 |
+
Returns (model, history). Validation is a TIME split supplied by the
|
| 416 |
+
caller — never a random split, because 75%-overlapping windows would put
|
| 417 |
+
near-duplicates on both sides and report a fantasy score.
|
| 418 |
+
"""
|
| 419 |
+
import mlx.optimizers as optim
|
| 420 |
+
cfg = cfg or {}
|
| 421 |
+
mx.random.seed(seed)
|
| 422 |
+
model = ContextTower(**cfg)
|
| 423 |
+
|
| 424 |
+
T = max(s.shape[0] for s in seqs)
|
| 425 |
+
X = np.zeros((len(seqs), T, seqs[0].shape[1]), dtype=np.float32)
|
| 426 |
+
for i, s in enumerate(seqs): # edge-pad short windows
|
| 427 |
+
X[i, :len(s)] = s
|
| 428 |
+
if len(s) < T:
|
| 429 |
+
X[i, len(s):] = s[-1]
|
| 430 |
+
Y = np.asarray(targets, dtype=np.float32)
|
| 431 |
+
|
| 432 |
+
val = np.zeros(len(seqs), dtype=bool)
|
| 433 |
+
val[val_idx] = True
|
| 434 |
+
tr = ~val
|
| 435 |
+
Xtr, Ytr = mx.array(X[tr]), mx.array(Y[tr])
|
| 436 |
+
Xva, Yva = mx.array(X[val]), mx.array(Y[val])
|
| 437 |
+
ig_tr = mx.array(overlap_mask([w for w, m in zip(windows, tr) if m]))
|
| 438 |
+
ig_va = mx.array(overlap_mask([w for w, m in zip(windows, val) if m]))
|
| 439 |
+
|
| 440 |
+
# logit scale/bias are learned, initialised as in the SigLIP paper
|
| 441 |
+
state = {"log_t": mx.array(np.float32(np.log(10.0))),
|
| 442 |
+
"bias": mx.array(np.float32(-10.0))}
|
| 443 |
+
|
| 444 |
+
def _norm(a):
|
| 445 |
+
# eps INSIDE the sqrt, not added to the norm afterwards. The head is
|
| 446 |
+
# zero-initialised on purpose, so the very first forward pass produces
|
| 447 |
+
# an exactly-zero vector — and d||x||/dx = x/||x|| is 0/0 = NaN there.
|
| 448 |
+
# Smoothing the radicand keeps the gradient finite at the origin.
|
| 449 |
+
return a * mx.rsqrt(mx.sum(a * a, axis=-1, keepdims=True) + 1e-8)
|
| 450 |
+
|
| 451 |
+
def loss_fn(m, x, y, ig):
|
| 452 |
+
v = _norm(m(x))
|
| 453 |
+
u = _norm(y)
|
| 454 |
+
# alignment term: pull each prediction onto its own caption. The
|
| 455 |
+
# contrastive term only fixes ORDER; this one fixes absolute position,
|
| 456 |
+
# which is what makes the vectors usable against unseen query text.
|
| 457 |
+
align = mx.mean(1.0 - mx.sum(v * u, axis=-1))
|
| 458 |
+
return siglip_loss(v, u, ignore=ig, log_t=state["log_t"],
|
| 459 |
+
bias=state["bias"]) + align_w * align
|
| 460 |
+
|
| 461 |
+
opt = optim.AdamW(learning_rate=lr, weight_decay=weight_decay)
|
| 462 |
+
lg = nn.value_and_grad(model, loss_fn)
|
| 463 |
+
sopt = optim.Adam(learning_rate=lr)
|
| 464 |
+
|
| 465 |
+
def state_loss(s):
|
| 466 |
+
model.set_training(False)
|
| 467 |
+
v = _norm(model(Xtr))
|
| 468 |
+
return siglip_loss(v, _norm(Ytr), ignore=ig_tr,
|
| 469 |
+
log_t=s["log_t"], bias=s["bias"])
|
| 470 |
+
|
| 471 |
+
# The aliveness masks are mx.array attributes, so they live in
|
| 472 |
+
# model.parameters() and AdamW would happily weight-decay them off their
|
| 473 |
+
# 0/1 values — silently rescaling every neuron and breaking the
|
| 474 |
+
# `mask == 1.0` alive-check the pruner depends on. Snapshot and restore
|
| 475 |
+
# after every step. (FDNN hit exactly this; it is not hypothetical.)
|
| 476 |
+
frozen_masks = [np.array(l.mask).copy() for l in model.turnover_layers()]
|
| 477 |
+
|
| 478 |
+
hist = {"train": [], "val": [], "val_r1": []}
|
| 479 |
+
best, best_w, best_state = 1e9, None, None
|
| 480 |
+
for ep in range(epochs):
|
| 481 |
+
model.set_training(True)
|
| 482 |
+
loss, grads = lg(model, Xtr, Ytr, ig_tr)
|
| 483 |
+
opt.update(model, grads)
|
| 484 |
+
for lyr, fm in zip(model.turnover_layers(), frozen_masks):
|
| 485 |
+
lyr.mask = mx.array(fm)
|
| 486 |
+
sl, sg = mx.value_and_grad(state_loss)(state)
|
| 487 |
+
state = sopt.apply_gradients(sg, state) # plain dict, not a Module
|
| 488 |
+
mx.eval(model.parameters(), opt.state, state)
|
| 489 |
+
|
| 490 |
+
model.set_training(False)
|
| 491 |
+
vl = float(loss_fn(model, Xva, Yva, ig_va).item())
|
| 492 |
+
r1 = retrieval_r1(np.array(_norm(model(Xva))), np.array(_norm(Yva)))
|
| 493 |
+
hist["train"].append(float(loss.item()))
|
| 494 |
+
hist["val"].append(vl)
|
| 495 |
+
hist["val_r1"].append(r1)
|
| 496 |
+
if vl < best:
|
| 497 |
+
best, best_w = vl, _clone_params(model)
|
| 498 |
+
best_state = {k: mx.array(np.array(v)) for k, v in state.items()}
|
| 499 |
+
if verbose and (ep % 50 == 0 or ep == epochs - 1):
|
| 500 |
+
print(f" ep {ep:4d} train {float(loss.item()):.4f} "
|
| 501 |
+
f"val {vl:.4f} val_R@1 {r1:.3f}", flush=True)
|
| 502 |
+
if best_w is not None:
|
| 503 |
+
_load_params(model, best_w)
|
| 504 |
+
state = best_state
|
| 505 |
+
model.set_training(False)
|
| 506 |
+
return model, {"history": hist, "best_val": best,
|
| 507 |
+
"log_t": float(state["log_t"].item()),
|
| 508 |
+
"bias": float(state["bias"].item())}
|
| 509 |
+
|
| 510 |
+
|
| 511 |
+
def retrieval_r1(v, u):
|
| 512 |
+
"""Fraction of windows whose own caption is its nearest caption.
|
| 513 |
+
|
| 514 |
+
This is the metric that matters: not "is the vector close to the target"
|
| 515 |
+
but "does the vector RANK the right target first" — the same question the
|
| 516 |
+
query path asks.
|
| 517 |
+
"""
|
| 518 |
+
if len(v) < 2:
|
| 519 |
+
return float("nan")
|
| 520 |
+
s = v @ u.T
|
| 521 |
+
return float((s.argmax(axis=1) == np.arange(len(v))).mean())
|
| 522 |
+
|
| 523 |
+
|
| 524 |
+
def _clone_params(model):
|
| 525 |
+
from mlx.utils import tree_flatten
|
| 526 |
+
return {k: np.array(v) for k, v in tree_flatten(model.parameters())}
|
| 527 |
+
|
| 528 |
+
|
| 529 |
+
def _load_params(model, flat):
|
| 530 |
+
from mlx.utils import tree_unflatten
|
| 531 |
+
model.update(tree_unflatten([(k, mx.array(v)) for k, v in flat.items()]))
|
| 532 |
+
|
| 533 |
+
|
| 534 |
+
def save_tower(model, codec, meta, path):
|
| 535 |
+
path = Path(path)
|
| 536 |
+
path.mkdir(parents=True, exist_ok=True)
|
| 537 |
+
np.savez(path / "tower.npz", **_clone_params(model))
|
| 538 |
+
codec.save(path / "codec.npz")
|
| 539 |
+
# model.cfg goes LAST and wins. `meta` is often a carried-forward copy of
|
| 540 |
+
# a previous tower.json and still holds that tower's cfg; letting it
|
| 541 |
+
# override would save compacted weights under the pre-compaction shape,
|
| 542 |
+
# and the next load would reshape-crash.
|
| 543 |
+
(path / "tower.json").write_text(json.dumps(
|
| 544 |
+
{**meta, "cfg": model.cfg}, indent=2))
|
| 545 |
+
|
| 546 |
+
|
| 547 |
+
def load_tower(path):
|
| 548 |
+
path = Path(path)
|
| 549 |
+
meta = json.loads((path / "tower.json").read_text())
|
| 550 |
+
model = ContextTower(**meta["cfg"])
|
| 551 |
+
z = np.load(path / "tower.npz")
|
| 552 |
+
_load_params(model, {k: z[k] for k in z.files})
|
| 553 |
+
model.set_training(False)
|
| 554 |
+
return model, ContextCodec.load(path / "codec.npz"), meta
|
python/elidedb/data/ssv2_classes.txt
ADDED
|
@@ -0,0 +1,174 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Approaching something with your camera
|
| 2 |
+
Attaching something to something
|
| 3 |
+
Bending something so that it deforms
|
| 4 |
+
Bending something until it breaks
|
| 5 |
+
Burying something in something
|
| 6 |
+
Closing something
|
| 7 |
+
Covering something with something
|
| 8 |
+
Digging something out of something
|
| 9 |
+
Dropping something behind something
|
| 10 |
+
Dropping something in front of something
|
| 11 |
+
Dropping something into something
|
| 12 |
+
Dropping something next to something
|
| 13 |
+
Dropping something onto something
|
| 14 |
+
Failing to put something into something because something does not fit
|
| 15 |
+
Folding something
|
| 16 |
+
Hitting something with something
|
| 17 |
+
Holding something
|
| 18 |
+
Holding something behind something
|
| 19 |
+
Holding something in front of something
|
| 20 |
+
Holding something next to something
|
| 21 |
+
Holding something over something
|
| 22 |
+
Laying something on the table on its side, not upright
|
| 23 |
+
Letting something roll along a flat surface
|
| 24 |
+
Letting something roll down a slanted surface
|
| 25 |
+
Letting something roll up a slanted surface, so it rolls back down
|
| 26 |
+
Lifting a surface with something on it but not enough for it to slide down
|
| 27 |
+
Lifting a surface with something on it until it starts sliding down
|
| 28 |
+
Lifting something up completely without letting it drop down
|
| 29 |
+
Lifting something up completely, then letting it drop down
|
| 30 |
+
Lifting something with something on it
|
| 31 |
+
Lifting up one end of something without letting it drop down
|
| 32 |
+
Lifting up one end of something, then letting it drop down
|
| 33 |
+
Moving away from something with your camera
|
| 34 |
+
Moving part of something
|
| 35 |
+
Moving something across a surface until it falls down
|
| 36 |
+
Moving something across a surface without it falling down
|
| 37 |
+
Moving something and something away from each other
|
| 38 |
+
Moving something and something closer to each other
|
| 39 |
+
Moving something and something so they collide with each other
|
| 40 |
+
Moving something and something so they pass each other
|
| 41 |
+
Moving something away from something
|
| 42 |
+
Moving something away from the camera
|
| 43 |
+
Moving something closer to something
|
| 44 |
+
Moving something down
|
| 45 |
+
Moving something towards the camera
|
| 46 |
+
Moving something up
|
| 47 |
+
Opening something
|
| 48 |
+
Picking something up
|
| 49 |
+
Piling something up
|
| 50 |
+
Plugging something into something
|
| 51 |
+
Plugging something into something but pulling it right out as you remove your hand
|
| 52 |
+
Poking a hole into some substance
|
| 53 |
+
Poking a hole into something soft
|
| 54 |
+
Poking a stack of something so the stack collapses
|
| 55 |
+
Poking a stack of something without the stack collapsing
|
| 56 |
+
Poking something so it slightly moves
|
| 57 |
+
Poking something so lightly that it doesn't or almost doesn't move
|
| 58 |
+
Poking something so that it falls over
|
| 59 |
+
Poking something so that it spins around
|
| 60 |
+
Pouring something into something
|
| 61 |
+
Pouring something into something until it overflows
|
| 62 |
+
Pouring something onto something
|
| 63 |
+
Pouring something out of something
|
| 64 |
+
Pretending or failing to wipe something off of something
|
| 65 |
+
Pretending or trying and failing to twist something
|
| 66 |
+
Pretending to be tearing something that is not tearable
|
| 67 |
+
Pretending to close something without actually closing it
|
| 68 |
+
Pretending to open something without actually opening it
|
| 69 |
+
Pretending to pick something up
|
| 70 |
+
Pretending to poke something
|
| 71 |
+
Pretending to pour something out of something, but something is empty
|
| 72 |
+
Pretending to put something behind something
|
| 73 |
+
Pretending to put something into something
|
| 74 |
+
Pretending to put something next to something
|
| 75 |
+
Pretending to put something on a surface
|
| 76 |
+
Pretending to put something onto something
|
| 77 |
+
Pretending to put something underneath something
|
| 78 |
+
Pretending to scoop something up with something
|
| 79 |
+
Pretending to spread air onto something
|
| 80 |
+
Pretending to sprinkle air onto something
|
| 81 |
+
Pretending to squeeze something
|
| 82 |
+
Pretending to take something from somewhere
|
| 83 |
+
Pretending to take something out of something
|
| 84 |
+
Pretending to throw something
|
| 85 |
+
Pretending to turn something upside down
|
| 86 |
+
Pulling something from behind of something
|
| 87 |
+
Pulling something from left to right
|
| 88 |
+
Pulling something from right to left
|
| 89 |
+
Pulling something onto something
|
| 90 |
+
Pulling something out of something
|
| 91 |
+
Pulling two ends of something but nothing happens
|
| 92 |
+
Pulling two ends of something so that it gets stretched
|
| 93 |
+
Pulling two ends of something so that it separates into two pieces
|
| 94 |
+
Pushing something from left to right
|
| 95 |
+
Pushing something from right to left
|
| 96 |
+
Pushing something off of something
|
| 97 |
+
Pushing something onto something
|
| 98 |
+
Pushing something so it spins
|
| 99 |
+
Pushing something so that it almost falls off but doesn't
|
| 100 |
+
Pushing something so that it falls off the table
|
| 101 |
+
Pushing something so that it slightly moves
|
| 102 |
+
Pushing something with something
|
| 103 |
+
Putting number of something onto something
|
| 104 |
+
Putting something and something on the table
|
| 105 |
+
Putting something behind something
|
| 106 |
+
Putting something in front of something
|
| 107 |
+
Putting something into something
|
| 108 |
+
Putting something next to something
|
| 109 |
+
Putting something on a flat surface without letting it roll
|
| 110 |
+
Putting something on a surface
|
| 111 |
+
Putting something on the edge of something so it is not supported and falls down
|
| 112 |
+
Putting something onto a slanted surface but it doesn't glide down
|
| 113 |
+
Putting something onto something
|
| 114 |
+
Putting something onto something else that cannot support it so it falls down
|
| 115 |
+
Putting something similar to other things that are already on the table
|
| 116 |
+
Putting something that can't roll onto a slanted surface, so it slides down
|
| 117 |
+
Putting something that can't roll onto a slanted surface, so it stays where it is
|
| 118 |
+
Putting something that cannot actually stand upright upright on the table, so it falls on its side
|
| 119 |
+
Putting something underneath something
|
| 120 |
+
Putting something upright on the table
|
| 121 |
+
Putting something, something and something on the table
|
| 122 |
+
Removing something, revealing something behind
|
| 123 |
+
Rolling something on a flat surface
|
| 124 |
+
Scooping something up with something
|
| 125 |
+
Showing a photo of something to the camera
|
| 126 |
+
Showing something behind something
|
| 127 |
+
Showing something next to something
|
| 128 |
+
Showing something on top of something
|
| 129 |
+
Showing something to the camera
|
| 130 |
+
Showing that something is empty
|
| 131 |
+
Showing that something is inside something
|
| 132 |
+
Something being deflected from something
|
| 133 |
+
Something colliding with something and both are being deflected
|
| 134 |
+
Something colliding with something and both come to a halt
|
| 135 |
+
Something falling like a feather or paper
|
| 136 |
+
Something falling like a rock
|
| 137 |
+
Spilling something behind something
|
| 138 |
+
Spilling something next to something
|
| 139 |
+
Spilling something onto something
|
| 140 |
+
Spinning something so it continues spinning
|
| 141 |
+
Spinning something that quickly stops spinning
|
| 142 |
+
Spreading something onto something
|
| 143 |
+
Sprinkling something onto something
|
| 144 |
+
Squeezing something
|
| 145 |
+
Stacking number of something
|
| 146 |
+
Stuffing something into something
|
| 147 |
+
Taking one of many similar things on the table
|
| 148 |
+
Taking something from somewhere
|
| 149 |
+
Taking something out of something
|
| 150 |
+
Tearing something into two pieces
|
| 151 |
+
Tearing something just a little bit
|
| 152 |
+
Throwing something
|
| 153 |
+
Throwing something against something
|
| 154 |
+
Throwing something in the air and catching it
|
| 155 |
+
Throwing something in the air and letting it fall
|
| 156 |
+
Throwing something onto a surface
|
| 157 |
+
Tilting something with something on it slightly so it doesn't fall down
|
| 158 |
+
Tilting something with something on it until it falls off
|
| 159 |
+
Tipping something over
|
| 160 |
+
Tipping something with something in it over, so something in it falls out
|
| 161 |
+
Touching (without moving) part of something
|
| 162 |
+
Trying but failing to attach something to something because it doesn't stick
|
| 163 |
+
Trying to bend something unbendable so nothing happens
|
| 164 |
+
Trying to pour something into something, but missing so it spills next to it
|
| 165 |
+
Turning something upside down
|
| 166 |
+
Turning the camera downwards while filming something
|
| 167 |
+
Turning the camera left while filming something
|
| 168 |
+
Turning the camera right while filming something
|
| 169 |
+
Turning the camera upwards while filming something
|
| 170 |
+
Twisting (wringing) something wet until water comes out
|
| 171 |
+
Twisting something
|
| 172 |
+
Uncovering something
|
| 173 |
+
Unfolding something
|
| 174 |
+
Wiping something off of something
|
python/elidedb/derive.py
ADDED
|
@@ -0,0 +1,148 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Corpus-derived replacements for hand-authored constants.
|
| 2 |
+
|
| 3 |
+
Every function here answers a question the code used to answer with a
|
| 4 |
+
literal. The rule they serve: nothing in code may encode what the data
|
| 5 |
+
CONTAINS or what the task IS. A constant is legitimate when it is a
|
| 6 |
+
tuning dial measured on the corpus or a property of a model; it is a
|
| 7 |
+
violation when it is a fact about the domain, because then the system is
|
| 8 |
+
being told the answer rather than recognising it.
|
| 9 |
+
|
| 10 |
+
The replacements have to work for any physical-AI corpus, not just this
|
| 11 |
+
one. A table-top manipulation set, a driving log and a warehouse camera
|
| 12 |
+
have nothing in common at the level of "open" or "put_into" - but they
|
| 13 |
+
all have objects that move, appear, vanish, approach and separate, and
|
| 14 |
+
those are measurable without naming anything.
|
| 15 |
+
|
| 16 |
+
WHAT WAS REPLACED
|
| 17 |
+
-----------------
|
| 18 |
+
MIN_BLOB = 120 (and 40, 60, 60, 30 elsewhere) -> fit_cut on the
|
| 19 |
+
observed blob-area distribution, in frame fractions so the
|
| 20 |
+
number transfers across resolutions.
|
| 21 |
+
REL_MIN / CAV_MIN thresholds -> fit_cut, or gone
|
| 22 |
+
entirely where the quantity feeds clustering instead of an
|
| 23 |
+
if/else.
|
| 24 |
+
hand-listed vocabularies -> attested(), words
|
| 25 |
+
the corpus actually produced.
|
| 26 |
+
hand-listed antonym pairs -> opposite_pairs(),
|
| 27 |
+
found by reflection in an embedding space.
|
| 28 |
+
"""
|
| 29 |
+
from __future__ import annotations
|
| 30 |
+
|
| 31 |
+
import numpy as np
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def fit_cut(values, q=None, mode="knee"):
|
| 35 |
+
"""A threshold from the data instead of from a keyboard.
|
| 36 |
+
|
| 37 |
+
mode="knee" the largest gap in the sorted values - where the
|
| 38 |
+
distribution itself separates, if it does.
|
| 39 |
+
mode="quantile" the q-th percentile, for "tighter than all but q% of
|
| 40 |
+
what I can prove is background".
|
| 41 |
+
mode="otsu" the between-class variance maximiser, for a
|
| 42 |
+
distribution with two modes and no prior on where
|
| 43 |
+
the split sits.
|
| 44 |
+
|
| 45 |
+
Returns (cut, diagnostics). The diagnostics matter: a knee found in
|
| 46 |
+
a unimodal distribution is an artefact, and `separation` says so, so
|
| 47 |
+
a caller can refuse to threshold a quantity that has no structure.
|
| 48 |
+
"""
|
| 49 |
+
v = np.asarray(values, np.float64)
|
| 50 |
+
v = v[np.isfinite(v)]
|
| 51 |
+
if len(v) < 8:
|
| 52 |
+
return None, {"reason": "too few samples", "n": int(len(v))}
|
| 53 |
+
s = np.sort(v)
|
| 54 |
+
if mode == "quantile":
|
| 55 |
+
cut = float(np.percentile(s, q if q is not None else 99.0))
|
| 56 |
+
elif mode == "otsu":
|
| 57 |
+
hist, edges = np.histogram(s, bins=min(256, max(16, len(s) // 8)))
|
| 58 |
+
p = hist / hist.sum()
|
| 59 |
+
c = np.cumsum(p)
|
| 60 |
+
m = np.cumsum(p * np.arange(len(p)))
|
| 61 |
+
mt = m[-1]
|
| 62 |
+
with np.errstate(invalid="ignore", divide="ignore"):
|
| 63 |
+
var = (mt * c - m) ** 2 / (c * (1 - c))
|
| 64 |
+
i = int(np.nanargmax(var))
|
| 65 |
+
cut = float(edges[i + 1])
|
| 66 |
+
else:
|
| 67 |
+
# knee: the widest gap between consecutive order statistics,
|
| 68 |
+
# ignoring the extreme tails where single outliers dominate
|
| 69 |
+
lo, hi = int(0.05 * len(s)), int(0.95 * len(s))
|
| 70 |
+
seg = s[lo:hi]
|
| 71 |
+
if len(seg) < 4:
|
| 72 |
+
return float(np.median(s)), {"reason": "degenerate", "n": len(s)}
|
| 73 |
+
gaps = np.diff(seg)
|
| 74 |
+
i = int(np.argmax(gaps))
|
| 75 |
+
cut = float((seg[i] + seg[i + 1]) / 2)
|
| 76 |
+
below, above = s[s <= cut], s[s > cut]
|
| 77 |
+
sep = 0.0
|
| 78 |
+
if len(below) and len(above):
|
| 79 |
+
spread = s.std() + 1e-12
|
| 80 |
+
sep = float((above.mean() - below.mean()) / spread)
|
| 81 |
+
return cut, {"n": int(len(s)), "cut": cut, "separation": round(sep, 3),
|
| 82 |
+
"frac_below": round(float(len(below) / len(s)), 3),
|
| 83 |
+
"mode": mode}
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def attested(texts, min_count=2, max_frac=0.5):
|
| 87 |
+
"""Vocabulary the CORPUS produced, not a list someone typed.
|
| 88 |
+
|
| 89 |
+
Drops words too rare to be a category and words so common they
|
| 90 |
+
cannot discriminate. No stopword list: "the" is excluded because it
|
| 91 |
+
appears in most documents, which is a measurement, not an opinion
|
| 92 |
+
about English.
|
| 93 |
+
"""
|
| 94 |
+
from collections import Counter
|
| 95 |
+
docs = [str(t).lower().split() for t in texts if t]
|
| 96 |
+
if not docs:
|
| 97 |
+
return []
|
| 98 |
+
df = Counter()
|
| 99 |
+
for d in docs:
|
| 100 |
+
df.update(set(d))
|
| 101 |
+
n = len(docs)
|
| 102 |
+
return sorted(w for w, c in df.items()
|
| 103 |
+
if c >= min_count and c / n <= max_frac and w.isalpha())
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def opposite_pairs(words, vec, top=1, min_sim=0.25):
|
| 107 |
+
"""Antonyms by REFLECTION, not by a table of 81 hand-written pairs.
|
| 108 |
+
|
| 109 |
+
In an embedding space trained on natural text, an antonym pair tends
|
| 110 |
+
to be the two ends of one axis: a and b are close in topic and
|
| 111 |
+
opposed in direction once their shared component is removed. So for
|
| 112 |
+
each word, remove the corpus mean, and look for the word whose
|
| 113 |
+
residual points most nearly the other way.
|
| 114 |
+
|
| 115 |
+
This finds whatever oppositions the corpus has. On a driving log it
|
| 116 |
+
would find accelerate/brake without anyone having thought of them;
|
| 117 |
+
on this one it finds open/close. That is the whole point - the
|
| 118 |
+
system stops depending on someone having anticipated the domain.
|
| 119 |
+
"""
|
| 120 |
+
W = [w for w in words if w]
|
| 121 |
+
if len(W) < 4:
|
| 122 |
+
return []
|
| 123 |
+
V = np.stack([vec(w) for w in W]).astype(np.float32)
|
| 124 |
+
V = V - V.mean(0) # the shared component is topic
|
| 125 |
+
V /= np.linalg.norm(V, axis=1, keepdims=True) + 1e-8
|
| 126 |
+
S = V @ V.T
|
| 127 |
+
np.fill_diagonal(S, 0.0)
|
| 128 |
+
out, seen = [], set()
|
| 129 |
+
for i, w in enumerate(W):
|
| 130 |
+
j = int(np.argmin(S[i])) # most opposed residual
|
| 131 |
+
if S[i, j] > -min_sim:
|
| 132 |
+
continue
|
| 133 |
+
key = tuple(sorted((i, j)))
|
| 134 |
+
if key in seen:
|
| 135 |
+
continue
|
| 136 |
+
seen.add(key)
|
| 137 |
+
out.append((w, W[j], round(float(-S[i, j]), 3)))
|
| 138 |
+
return sorted(out, key=lambda r: -r[2])[:max(top * len(W), 1)]
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def frame_fraction(px, width, height):
|
| 142 |
+
"""Pixel counts are resolution-specific; fractions are not.
|
| 143 |
+
|
| 144 |
+
MIN_BLOB=120 meant one thing at 640x480 and something else at
|
| 145 |
+
256x256, which is why five files carried five different values for
|
| 146 |
+
the same idea.
|
| 147 |
+
"""
|
| 148 |
+
return float(px) / float(max(width * height, 1))
|
python/elidedb/desk.py
ADDED
|
@@ -0,0 +1,1723 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""ElideDB Desk — local database browser (the Atlas/Compass role).
|
| 3 |
+
|
| 4 |
+
Zero-dependency server: stdlib http.server + the elidedb package. All state
|
| 5 |
+
lives in the stores themselves; Desk only reads. Thumbnails are decoded on
|
| 6 |
+
demand through the same byte-range path queries use — nothing is pre-baked,
|
| 7 |
+
so every embedded window can always show its frame.
|
| 8 |
+
|
| 9 |
+
elidedb desk [--root lake] [--port 8787] [--open]
|
| 10 |
+
"""
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import argparse
|
| 14 |
+
import io
|
| 15 |
+
import json
|
| 16 |
+
import os
|
| 17 |
+
import sys
|
| 18 |
+
import threading
|
| 19 |
+
import time
|
| 20 |
+
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
from urllib.parse import parse_qs, urlparse
|
| 23 |
+
|
| 24 |
+
ROOT = Path.cwd()
|
| 25 |
+
|
| 26 |
+
import numpy as np # noqa: E402
|
| 27 |
+
from elidedb import Store # noqa: E402
|
| 28 |
+
|
| 29 |
+
STORES: dict[str, Store] = {}
|
| 30 |
+
LAKE = ROOT / "lake"
|
| 31 |
+
_CACHE: dict = {}
|
| 32 |
+
import os # noqa: E402
|
| 33 |
+
READONLY = os.environ.get("DESK_READONLY", "") == "1"
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def discover():
|
| 37 |
+
STORES.clear()
|
| 38 |
+
if LAKE.is_dir():
|
| 39 |
+
for p in sorted(LAKE.iterdir()):
|
| 40 |
+
if (p / "_store.json").exists():
|
| 41 |
+
try:
|
| 42 |
+
STORES[p.name] = Store.open(p)
|
| 43 |
+
except Exception:
|
| 44 |
+
pass
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def _covered_seconds(db, desc):
|
| 48 |
+
"""Recorded CONTENT time: the union of row-group ts ranges of the
|
| 49 |
+
frame tables, from parquet footers only. Extent (max minus min) lies
|
| 50 |
+
on sparse stores; a store holding three hours of episodes spread
|
| 51 |
+
over three days of wall clock covers three hours, not three days."""
|
| 52 |
+
import pyarrow.parquet as pq
|
| 53 |
+
frame_tabs = [d["table"] for d in desc
|
| 54 |
+
if d["kind"] == "frame_index" and d["rows"]]
|
| 55 |
+
if not frame_tabs:
|
| 56 |
+
frame_tabs = [d["table"] for d in desc
|
| 57 |
+
if d["table"] == "frames" and d["rows"]]
|
| 58 |
+
# exact when available: every ingest commit recorded its fps, and
|
| 59 |
+
# frame count over fps is the recorded duration regardless of how
|
| 60 |
+
# the rows pack into row groups
|
| 61 |
+
total_s = 0.0
|
| 62 |
+
exact = False
|
| 63 |
+
for name in frame_tabs:
|
| 64 |
+
try:
|
| 65 |
+
for c in db.table(name).history():
|
| 66 |
+
m = c.get("meta") or {}
|
| 67 |
+
fps = m.get("fps")
|
| 68 |
+
if fps and c.get("added_rows"):
|
| 69 |
+
total_s += c["added_rows"] / float(fps)
|
| 70 |
+
exact = True
|
| 71 |
+
except Exception:
|
| 72 |
+
continue
|
| 73 |
+
if exact:
|
| 74 |
+
return total_s
|
| 75 |
+
# episode-organized stores: the episodes table IS the content list
|
| 76 |
+
if any(d["table"] == "episodes" and d["rows"] for d in desc):
|
| 77 |
+
try:
|
| 78 |
+
t = db.table("episodes").scan(columns=["ts", "t1"])
|
| 79 |
+
a = t.column("ts").to_numpy()
|
| 80 |
+
b = t.column("t1").to_numpy()
|
| 81 |
+
return float((b - a).sum() / 1e9)
|
| 82 |
+
except Exception:
|
| 83 |
+
pass
|
| 84 |
+
spans = []
|
| 85 |
+
for name in frame_tabs:
|
| 86 |
+
try:
|
| 87 |
+
st = db.table(name).state()
|
| 88 |
+
for f in st.files:
|
| 89 |
+
pf = pq.ParquetFile(db.dir / "tables" / name / f.path)
|
| 90 |
+
md = pf.metadata
|
| 91 |
+
names = md.schema.names
|
| 92 |
+
if "ts" not in names:
|
| 93 |
+
continue
|
| 94 |
+
ti = names.index("ts")
|
| 95 |
+
for g in range(md.num_row_groups):
|
| 96 |
+
s = md.row_group(g).column(ti).statistics
|
| 97 |
+
if s and s.min is not None:
|
| 98 |
+
spans.append((int(s.min), int(s.max)))
|
| 99 |
+
except Exception:
|
| 100 |
+
continue
|
| 101 |
+
if not spans:
|
| 102 |
+
return None
|
| 103 |
+
spans.sort()
|
| 104 |
+
total, cur_a, cur_b = 0, spans[0][0], spans[0][1]
|
| 105 |
+
for a, b in spans[1:]:
|
| 106 |
+
if a > cur_b:
|
| 107 |
+
total += cur_b - cur_a
|
| 108 |
+
cur_a, cur_b = a, b
|
| 109 |
+
else:
|
| 110 |
+
cur_b = max(cur_b, b)
|
| 111 |
+
total += cur_b - cur_a
|
| 112 |
+
return total / 1e9
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def _raw_source_bytes(db, desc):
|
| 116 |
+
"""Bytes of the ORIGINAL ingested sources, from the ingest metadata
|
| 117 |
+
each append recorded. This is the honest numerator of the
|
| 118 |
+
compression story; a missing original is reported as unknown, not
|
| 119 |
+
guessed."""
|
| 120 |
+
originals = set()
|
| 121 |
+
for d in desc:
|
| 122 |
+
if d["kind"] != "frame_index" and d["table"] != "frames":
|
| 123 |
+
continue
|
| 124 |
+
try:
|
| 125 |
+
for c in db.table(d["table"]).history():
|
| 126 |
+
o = (c.get("meta") or {}).get("original")
|
| 127 |
+
if o:
|
| 128 |
+
originals.add(o)
|
| 129 |
+
except Exception:
|
| 130 |
+
continue
|
| 131 |
+
known = missing = 0
|
| 132 |
+
total = 0
|
| 133 |
+
for o in originals:
|
| 134 |
+
p = Path(o)
|
| 135 |
+
if p.exists():
|
| 136 |
+
total += p.stat().st_size
|
| 137 |
+
known += 1
|
| 138 |
+
else:
|
| 139 |
+
missing += 1
|
| 140 |
+
return {"bytes": total, "files": known, "missing": missing}
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def store_summary(key: str):
|
| 144 |
+
db = STORES[key]
|
| 145 |
+
desc = db.describe()
|
| 146 |
+
ck = ("summary", key,
|
| 147 |
+
tuple(sorted((d["table"], d["version"]) for d in desc)))
|
| 148 |
+
if ck in _CACHE:
|
| 149 |
+
return _CACHE[ck]
|
| 150 |
+
total_rows = sum(d["rows"] for d in desc)
|
| 151 |
+
total_bytes = sum(d["bytes"] for d in desc)
|
| 152 |
+
span_tabs = [d for d in desc if d["rows"] and d["table"] != "centroids"]
|
| 153 |
+
lo = min((d["min_ts"] for d in span_tabs), default=0)
|
| 154 |
+
hi = max((d["max_ts"] for d in span_tabs), default=0)
|
| 155 |
+
emb = next((d for d in desc if d["table"] == "embeddings"), None)
|
| 156 |
+
# PHYSICAL directory bytes (lstat): a symlink is a path entry, not
|
| 157 |
+
# the target's bytes. Media reached through symlinks is counted as
|
| 158 |
+
# linked_bytes instead, because the store STOPS WORKING if those
|
| 159 |
+
# targets go away, and the standalone claim must be earned.
|
| 160 |
+
media_bytes = linked_bytes = 0
|
| 161 |
+
mdir = db.dir / "media"
|
| 162 |
+
if mdir.is_dir():
|
| 163 |
+
for p in mdir.glob("*"):
|
| 164 |
+
if p.is_symlink():
|
| 165 |
+
try:
|
| 166 |
+
linked_bytes += p.stat().st_size
|
| 167 |
+
except OSError:
|
| 168 |
+
pass
|
| 169 |
+
elif p.is_file():
|
| 170 |
+
media_bytes += p.lstat().st_size
|
| 171 |
+
db_bytes = sum(p.lstat().st_size for p in db.dir.rglob("*")
|
| 172 |
+
if p.is_file() and not p.is_symlink())
|
| 173 |
+
external_bytes = 0
|
| 174 |
+
for d in desc:
|
| 175 |
+
if d["kind"] != "frame_index":
|
| 176 |
+
continue
|
| 177 |
+
t = db.table(d["table"]).scan(columns=["source"])
|
| 178 |
+
for s in set(t.column("source").to_pylist()):
|
| 179 |
+
if not s.startswith("@") and Path(s).exists():
|
| 180 |
+
external_bytes += Path(s).stat().st_size
|
| 181 |
+
raw = _raw_source_bytes(db, desc)
|
| 182 |
+
vec_rows = sum(d["rows"] for d in desc if d["kind"] == "embeddings")
|
| 183 |
+
vec_tables = sum(1 for d in desc
|
| 184 |
+
if d["kind"] == "embeddings" and d["rows"])
|
| 185 |
+
out = {
|
| 186 |
+
"key": key, "name": db.name, "path": str(db.dir),
|
| 187 |
+
"tables": desc, "rows": total_rows, "bytes": total_bytes,
|
| 188 |
+
"db_bytes": db_bytes, "media_bytes": media_bytes,
|
| 189 |
+
"emb_bytes": emb["bytes"] if emb else 0,
|
| 190 |
+
"external_bytes": external_bytes,
|
| 191 |
+
"linked_bytes": linked_bytes,
|
| 192 |
+
"standalone": external_bytes == 0 and linked_bytes == 0,
|
| 193 |
+
"raw_bytes": raw["bytes"], "raw_files": raw["files"],
|
| 194 |
+
"raw_missing": raw["missing"],
|
| 195 |
+
"covered_s": _covered_seconds(db, desc),
|
| 196 |
+
"min_ts": lo, "max_ts": hi,
|
| 197 |
+
"windows": emb["rows"] if emb else 0,
|
| 198 |
+
"vec_rows": vec_rows, "vec_tables": vec_tables,
|
| 199 |
+
"model": (emb or {}).get("meta", {}).get("model", ""),
|
| 200 |
+
"display": db.meta.get("display", {}),
|
| 201 |
+
}
|
| 202 |
+
if len(_CACHE) > 32:
|
| 203 |
+
_CACHE.clear()
|
| 204 |
+
_CACHE[ck] = out
|
| 205 |
+
return out
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def api_storage(key: str, table: str):
|
| 209 |
+
"""The Parquet format, made visible: this table's commit log plus every
|
| 210 |
+
active file's row-group layout straight from the Parquet footers."""
|
| 211 |
+
import pyarrow.parquet as pq
|
| 212 |
+
db = STORES[key]
|
| 213 |
+
tab = db.table(table)
|
| 214 |
+
st = tab.state()
|
| 215 |
+
files = []
|
| 216 |
+
for f in st.files:
|
| 217 |
+
pf = pq.ParquetFile(db.dir / "tables" / table / f.path)
|
| 218 |
+
md = pf.metadata
|
| 219 |
+
names = md.schema.names
|
| 220 |
+
ts_i = names.index("ts") if "ts" in names else 0
|
| 221 |
+
rgs = []
|
| 222 |
+
for g in range(md.num_row_groups):
|
| 223 |
+
rg = md.row_group(g)
|
| 224 |
+
s = rg.column(ts_i).statistics
|
| 225 |
+
rgs.append({"rows": rg.num_rows,
|
| 226 |
+
"bytes": sum(rg.column(c).total_compressed_size
|
| 227 |
+
for c in range(rg.num_columns)),
|
| 228 |
+
"min_ts": s.min if s else None,
|
| 229 |
+
"max_ts": s.max if s else None})
|
| 230 |
+
files.append({"name": f.path, "bytes": f.bytes, "rows": f.rows,
|
| 231 |
+
"min_ts": f.min_ts, "max_ts": f.max_ts,
|
| 232 |
+
"footer_bytes": md.serialized_size,
|
| 233 |
+
"columns": names, "row_groups": rgs})
|
| 234 |
+
return {"table": table, "kind": st.kind, "version": st.version,
|
| 235 |
+
"schema": st.schema, "history": tab.history(), "files": files}
|
| 236 |
+
|
| 237 |
+
|
| 238 |
+
MAP_MAX_POINTS = 6000
|
| 239 |
+
|
| 240 |
+
|
| 241 |
+
def api_map(key: str):
|
| 242 |
+
"""2D layout of the embeddings table. UMAP over a bounded SAMPLE,
|
| 243 |
+
cached beside the log (a derived view; never used for retrieval).
|
| 244 |
+
|
| 245 |
+
The unsampled version killed the app at pilot scale, three ways at
|
| 246 |
+
once (measured on the 100 h store): to_pylist() over 180k x 1152
|
| 247 |
+
vectors is ~25 GB of Python floats (4 -> 28 GB RSS), UMAP over 180k
|
| 248 |
+
points runs for minutes, and a 180k-point JSON payload crushes the
|
| 249 |
+
WebView. A map is an OVERVIEW: an even time-stride sample of a few
|
| 250 |
+
thousand windows shows the same structure, reads from the mmap
|
| 251 |
+
sidecar, and stays bounded no matter how large the corpus grows."""
|
| 252 |
+
ck = ("map", key)
|
| 253 |
+
if ck in _CACHE:
|
| 254 |
+
return _CACHE[ck]
|
| 255 |
+
db = STORES[key]
|
| 256 |
+
from elidedb.embeddings import _vec_table
|
| 257 |
+
try:
|
| 258 |
+
t, vecs = _vec_table(db, "embeddings")
|
| 259 |
+
except Exception:
|
| 260 |
+
return {"points": []}
|
| 261 |
+
n = len(t)
|
| 262 |
+
if n == 0:
|
| 263 |
+
return {"points": []}
|
| 264 |
+
stride = max(1, n // MAP_MAX_POINTS)
|
| 265 |
+
idx = np.arange(0, n, stride)
|
| 266 |
+
sample = np.asarray(vecs[idx], np.float32)
|
| 267 |
+
|
| 268 |
+
cache_file = db.dir / "tables" / "embeddings" / "_desk_umap.json"
|
| 269 |
+
st = db.table("embeddings").state()
|
| 270 |
+
xy = None
|
| 271 |
+
if cache_file.exists():
|
| 272 |
+
c = json.loads(cache_file.read_text())
|
| 273 |
+
if c.get("version") == st.version and len(c["xy"]) == len(idx):
|
| 274 |
+
xy = np.array(c["xy"], np.float32)
|
| 275 |
+
if xy is None:
|
| 276 |
+
try:
|
| 277 |
+
import umap
|
| 278 |
+
from sklearn.decomposition import PCA
|
| 279 |
+
red = PCA(n_components=min(50, len(sample), sample.shape[1]),
|
| 280 |
+
random_state=0).fit_transform(sample)
|
| 281 |
+
xy = umap.UMAP(n_components=2, random_state=0,
|
| 282 |
+
low_memory=True).fit_transform(red)
|
| 283 |
+
except Exception:
|
| 284 |
+
from sklearn.decomposition import PCA
|
| 285 |
+
xy = PCA(n_components=2, random_state=0).fit_transform(sample)
|
| 286 |
+
cache_file.write_text(json.dumps(
|
| 287 |
+
{"version": st.version, "xy": np.round(xy, 3).tolist()}))
|
| 288 |
+
labels = (t.column("cluster").to_pylist()
|
| 289 |
+
if "cluster" in t.column_names else None)
|
| 290 |
+
# NO cluster column: every point got c=0 and the map was one flat
|
| 291 |
+
# colour with a legend reading "cluster 0 · 3891" - information-free.
|
| 292 |
+
# Cluster the 2D LAYOUT for display, cached beside the UMAP cache.
|
| 293 |
+
# DISPLAY ONLY, and the UI says so: UMAP preserves neighbourhoods,
|
| 294 |
+
# not distances, so these coordinates never touch a retrieval
|
| 295 |
+
# decision - that rule is why the store has no cluster column in the
|
| 296 |
+
# first place.
|
| 297 |
+
kmeans_display = False
|
| 298 |
+
if labels is None and len(xy) >= 24:
|
| 299 |
+
kc = db.dir / "tables" / "embeddings" / "_desk_kmeans.json"
|
| 300 |
+
got = None
|
| 301 |
+
if kc.exists():
|
| 302 |
+
try:
|
| 303 |
+
j = json.loads(kc.read_text())
|
| 304 |
+
if j.get("version") == st.version and len(j["lab"]) == n:
|
| 305 |
+
got = j["lab"]
|
| 306 |
+
except Exception:
|
| 307 |
+
got = None
|
| 308 |
+
if got is None:
|
| 309 |
+
P = np.asarray(xy, np.float32)
|
| 310 |
+
K = min(10, max(2, len(P) // 40))
|
| 311 |
+
rng = np.random.default_rng(0)
|
| 312 |
+
C = P[rng.choice(len(P), K, replace=False)]
|
| 313 |
+
for _ in range(25): # Lloyd, few rounds
|
| 314 |
+
d2 = ((P[:, None, :] - C[None]) ** 2).sum(-1)
|
| 315 |
+
a = d2.argmin(1)
|
| 316 |
+
for k in range(K):
|
| 317 |
+
m = a == k
|
| 318 |
+
if m.any():
|
| 319 |
+
C[k] = P[m].mean(0)
|
| 320 |
+
sub = a.astype(int).tolist()
|
| 321 |
+
got = [0] * n
|
| 322 |
+
for j, i in enumerate(idx):
|
| 323 |
+
got[i] = sub[j]
|
| 324 |
+
try:
|
| 325 |
+
kc.write_text(json.dumps({"version": st.version,
|
| 326 |
+
"lab": got}))
|
| 327 |
+
except Exception:
|
| 328 |
+
pass
|
| 329 |
+
labels = got
|
| 330 |
+
kmeans_display = True
|
| 331 |
+
ss = t.column("stream").to_pylist()
|
| 332 |
+
ta = t.column("ts").to_pylist()
|
| 333 |
+
tb = t.column("t1").to_pylist()
|
| 334 |
+
out = {"points": [
|
| 335 |
+
{"x": float(xy[j][0]), "y": float(xy[j][1]),
|
| 336 |
+
"c": int(labels[i]) if labels else 0,
|
| 337 |
+
"s": ss[i], "t0": ta[i], "t1": tb[i]}
|
| 338 |
+
for j, i in enumerate(idx)],
|
| 339 |
+
"sampled_of": n, "stride": int(stride),
|
| 340 |
+
"cluster_source": ("stored cluster column" if not kmeans_display
|
| 341 |
+
else "k-means over the 2D layout — DISPLAY "
|
| 342 |
+
"ONLY, never a retrieval decision")}
|
| 343 |
+
_CACHE[ck] = out
|
| 344 |
+
return out
|
| 345 |
+
|
| 346 |
+
|
| 347 |
+
def api_architecture(key: str):
|
| 348 |
+
"""The store as a living schematic: every node and edge derived from
|
| 349 |
+
what is ACTUALLY on disk — parquet footers for schemas, _meta.json for
|
| 350 |
+
state, table meta for lineage, _cache for mmap sidecars. Nothing here
|
| 351 |
+
is drawn from documentation; a table that vanished vanishes from the
|
| 352 |
+
drawing, a channel appears the moment its table exists."""
|
| 353 |
+
import pyarrow.parquet as _pq
|
| 354 |
+
db = STORES[key]
|
| 355 |
+
nodes, edges = [], []
|
| 356 |
+
|
| 357 |
+
def table_node(name):
|
| 358 |
+
tab = db.table(name)
|
| 359 |
+
try:
|
| 360 |
+
st = tab.state()
|
| 361 |
+
except Exception:
|
| 362 |
+
return None
|
| 363 |
+
if not st.files:
|
| 364 |
+
return None
|
| 365 |
+
meta_json = {}
|
| 366 |
+
mf = tab.dir / "_meta.json"
|
| 367 |
+
if mf.exists():
|
| 368 |
+
try:
|
| 369 |
+
meta_json = json.loads(mf.read_text())
|
| 370 |
+
except Exception:
|
| 371 |
+
pass
|
| 372 |
+
fields = []
|
| 373 |
+
try:
|
| 374 |
+
sch = _pq.ParquetFile(tab.dir / st.files[0].path).schema_arrow
|
| 375 |
+
for f in sch:
|
| 376 |
+
t = str(f.type)
|
| 377 |
+
t = t.replace("fixed_size_list<item: float>", "f32vec")
|
| 378 |
+
fields.append({"name": f.name, "type": t})
|
| 379 |
+
except Exception:
|
| 380 |
+
pass
|
| 381 |
+
sidecars = sorted(p.name for p in (tab.dir / "_cache").glob("*.npy")) \
|
| 382 |
+
if (tab.dir / "_cache").is_dir() else []
|
| 383 |
+
rows = sum(f.rows for f in st.files)
|
| 384 |
+
return {"id": name, "kind": st.kind, "rows": rows,
|
| 385 |
+
"bytes": sum(f.bytes for f in st.files),
|
| 386 |
+
"files": len(st.files), "version": st.version,
|
| 387 |
+
"min_ts": min((f.min_ts for f in st.files), default=None),
|
| 388 |
+
"max_ts": max((f.max_ts for f in st.files), default=None),
|
| 389 |
+
"fields": fields, "meta": st.meta or {},
|
| 390 |
+
"meta_json": bool(meta_json), "sidecars": sidecars}
|
| 391 |
+
|
| 392 |
+
names = db.tables()
|
| 393 |
+
for n in names:
|
| 394 |
+
nd = table_node(n)
|
| 395 |
+
if nd:
|
| 396 |
+
nodes.append(nd)
|
| 397 |
+
have = {n["id"] for n in nodes}
|
| 398 |
+
|
| 399 |
+
# media + models are first-class citizens of the drawing
|
| 400 |
+
mdir = db.dir / "media"
|
| 401 |
+
if mdir.is_dir():
|
| 402 |
+
fs = list(mdir.glob("*"))
|
| 403 |
+
nodes.append({"id": "media", "kind": "media",
|
| 404 |
+
"rows": len(fs),
|
| 405 |
+
"bytes": sum(f.stat().st_size for f in fs
|
| 406 |
+
if f.is_file()),
|
| 407 |
+
"files": len(fs), "fields": [], "meta": {},
|
| 408 |
+
"note": "transcoded H.264, byte-range decoded"})
|
| 409 |
+
for m in sorted((db.dir / "models").glob("*")) \
|
| 410 |
+
if (db.dir / "models").is_dir() else []:
|
| 411 |
+
if m.is_dir():
|
| 412 |
+
nodes.append({"id": f"model:{m.name}", "kind": "model",
|
| 413 |
+
"bytes": sum(f.stat().st_size
|
| 414 |
+
for f in m.rglob("*") if f.is_file()),
|
| 415 |
+
"fields": [], "meta": {},
|
| 416 |
+
"rows": None, "files": None})
|
| 417 |
+
|
| 418 |
+
# lineage: explicit table meta first, then structural conventions
|
| 419 |
+
def edge(a, b, label):
|
| 420 |
+
if a in have or a in ("media", "SOURCE", "QUERIES") or \
|
| 421 |
+
a.startswith("model:"):
|
| 422 |
+
edges.append({"from": a, "to": b, "label": label})
|
| 423 |
+
for n in nodes:
|
| 424 |
+
meta, nid = n.get("meta", {}), n["id"]
|
| 425 |
+
if meta.get("source_table"):
|
| 426 |
+
edge(meta["source_table"], nid,
|
| 427 |
+
meta.get("built_by", "derived"))
|
| 428 |
+
if meta.get("events_table"):
|
| 429 |
+
edge(meta["events_table"], nid, "spans")
|
| 430 |
+
if meta.get("teacher"):
|
| 431 |
+
edge("frames", nid, meta.get("model", "encoder"))
|
| 432 |
+
conventions = {
|
| 433 |
+
"frame_vectors": ("frames", "encoder, every frame"),
|
| 434 |
+
"object_vectors": ("frames", "FastSAM regions + crops"),
|
| 435 |
+
"context_captions": ("frames", "VLM captions"),
|
| 436 |
+
"context_events": ("frames", "gate segmentation"),
|
| 437 |
+
"vlm_verdicts": ("frames", "2B/7B judgments"),
|
| 438 |
+
"frames": ("media", "frame index"),
|
| 439 |
+
}
|
| 440 |
+
done = {(e["from"], e["to"]) for e in edges}
|
| 441 |
+
for nid, (src, lab) in conventions.items():
|
| 442 |
+
if nid in have and (src, nid) not in done and \
|
| 443 |
+
(src in have or src == "media"):
|
| 444 |
+
edges.append({"from": src, "to": nid, "label": lab})
|
| 445 |
+
|
| 446 |
+
# query channels: present iff their table exists on disk. This list
|
| 447 |
+
# is the CURRENT set path, nothing else; a channel appears the
|
| 448 |
+
# moment its vectors are ingested and vanishes with them.
|
| 449 |
+
channels = []
|
| 450 |
+
def chan(cid, label, need, note):
|
| 451 |
+
ok = need() if callable(need) else need in have
|
| 452 |
+
if ok:
|
| 453 |
+
channels.append({"id": cid, "label": label, "note": note})
|
| 454 |
+
chan("app", "appearance", "embeddings", "window vectors, exact scan")
|
| 455 |
+
chan("pe", "perception", "pe_vectors", "PE-Core text to frame")
|
| 456 |
+
chan("sig2", "fine-grained", "sig2_vectors", "SigLIP 2 frame space")
|
| 457 |
+
chan("conj", "conjunction", "sig2_vectors",
|
| 458 |
+
"every noun phrase must find its own frame")
|
| 459 |
+
chan("iv2", "video-text", "iv2_vectors",
|
| 460 |
+
"InternVideo2, 4 frames encoded together")
|
| 461 |
+
chan("act", "action", "action_probs", "V-JEPA 2 verb posteriors")
|
| 462 |
+
chan("vid", "clip-text", "xclip_vectors", "X-CLIP pooled clips")
|
| 463 |
+
chan("obj", "objects", "object_vectors", "region crops, conjunctive")
|
| 464 |
+
chan("mot", "motion", "motion_vectors",
|
| 465 |
+
"delta appearance against the antonym")
|
| 466 |
+
chan("prf", "feedback", "vjepa_vectors",
|
| 467 |
+
"Rocchio anchors mined from this corpus")
|
| 468 |
+
chan("itm", "cross-encoder", lambda: (
|
| 469 |
+
os.environ.get("ELIDEDB_ITM") == "1"
|
| 470 |
+
and (db.dir / "_cache/itm_tokens").exists()),
|
| 471 |
+
"InternVideo2 1B, rerank STAGE over the top-N - never a channel "
|
| 472 |
+
"(as a weighted voter it cost 0.38 -> 0.27, RRF discards the "
|
| 473 |
+
"margin's scale). Cost-gated: 0.4s/episode")
|
| 474 |
+
chan("anchor", "transition anchor", "motion_vectors",
|
| 475 |
+
"the query names a transition, the corpus defines its DIRECTION "
|
| 476 |
+
"in motion space. Text cannot ask for it: cos(opens, closes) = "
|
| 477 |
+
"0.957. Motion separates them 0.983 held-out")
|
| 478 |
+
|
| 479 |
+
# selection pipeline: the fitted stages every result passes
|
| 480 |
+
# through, in order. Fitted values come from the per-store
|
| 481 |
+
# artifact; stages render even unfitted (neutral defaults).
|
| 482 |
+
fitted = {}
|
| 483 |
+
swp = db.dir / "_set_weights.json"
|
| 484 |
+
if swp.exists():
|
| 485 |
+
try:
|
| 486 |
+
fitted = json.loads(swp.read_text())
|
| 487 |
+
except Exception:
|
| 488 |
+
pass
|
| 489 |
+
pipeline = [
|
| 490 |
+
{"id": "fuse", "label": "Fitted fusion",
|
| 491 |
+
"note": "weighted rank consensus, per-store weights"},
|
| 492 |
+
{"id": "prf", "label": "Pseudo-relevance",
|
| 493 |
+
"note": "the head of pass 1 re-queries the corpus"},
|
| 494 |
+
{"id": "itm", "label": "Cross-encoder cascade",
|
| 495 |
+
"note": "distribution-preserving: it PERMUTES the candidates "
|
| 496 |
+
"and returns the same sorted scores in the new order, "
|
| 497 |
+
"because the cut downstream is fitted to RRF's scale"},
|
| 498 |
+
{"id": "anchor", "label": "Transition anchor",
|
| 499 |
+
"note": "direction from the corpus, weighted by a reliability "
|
| 500 |
+
"each kind earns unsupervised (close 0.52, open 0.35, "
|
| 501 |
+
"put_on 0.00 - so it cannot damage what it cannot help)"},
|
| 502 |
+
{"id": "evk", "label": "Event corroboration",
|
| 503 |
+
"note": "does this demo carry the asked-for transition at all"},
|
| 504 |
+
{"id": "dens", "label": "Motion density",
|
| 505 |
+
"note": "15-NN agreement in motion space"},
|
| 506 |
+
{"id": "gate", "label": "No-match gate",
|
| 507 |
+
"note": "abstains when the corpus lacks the action"},
|
| 508 |
+
{"id": "filter", "label": "Contrast filter",
|
| 509 |
+
"note": "direction evidence, fitted quantile"},
|
| 510 |
+
{"id": "nms", "label": "Event dedup",
|
| 511 |
+
"note": "one clip per event, fitted radius"},
|
| 512 |
+
{"id": "cut", "label": "Confidence cut",
|
| 513 |
+
"note": "the set ends where confidence does"},
|
| 514 |
+
{"id": "audit", "label": "Geometry audit",
|
| 515 |
+
"note": "SAM 3 tracker verification, opt-in tier"},
|
| 516 |
+
]
|
| 517 |
+
|
| 518 |
+
# THE FIVE ELEMENTS — what a clip is decomposed into. Each renders
|
| 519 |
+
# only if the table that carries it exists in THIS store, so an
|
| 520 |
+
# un-elemented store shows an empty list rather than a promise.
|
| 521 |
+
def _rows(t):
|
| 522 |
+
try:
|
| 523 |
+
return db.table(t).scan().num_rows if t in have else 0
|
| 524 |
+
except Exception:
|
| 525 |
+
return 0
|
| 526 |
+
|
| 527 |
+
ev_t = "events" if "events" in have else (
|
| 528 |
+
"events_s" if "events_s" in have else None)
|
| 529 |
+
elements = []
|
| 530 |
+
if ev_t:
|
| 531 |
+
import collections as _c
|
| 532 |
+
kinds = _c.Counter(db.table(ev_t).scan().column("kind").to_pylist())
|
| 533 |
+
elements = [
|
| 534 |
+
{"id": "scene", "label": "scene", "note": "demo gist vector",
|
| 535 |
+
"n": _rows("answers2")},
|
| 536 |
+
{"id": "agent", "label": "agent",
|
| 537 |
+
"note": "the self-moving thing, from flow",
|
| 538 |
+
"n": kinds.get("agent", 0)},
|
| 539 |
+
{"id": "participants", "label": "participants",
|
| 540 |
+
"note": "what the agent contacts, in order - by CAUSALITY "
|
| 541 |
+
"(motion onset adjacent to the agent), not pixel "
|
| 542 |
+
"change", "n": kinds.get("contact", 0)},
|
| 543 |
+
{"id": "events", "label": "events",
|
| 544 |
+
"note": "typed, TIMESTAMPED transitions: " + ", ".join(
|
| 545 |
+
f"{k} {v}" for k, v in kinds.most_common()
|
| 546 |
+
if k not in ("agent", "contact", "release")),
|
| 547 |
+
"n": sum(v for k, v in kinds.items()
|
| 548 |
+
if k not in ("agent", "contact", "release"))},
|
| 549 |
+
{"id": "answer", "label": "answer",
|
| 550 |
+
"note": "initial -> final diff, attributed to the agent",
|
| 551 |
+
"n": _rows("answers2")},
|
| 552 |
+
]
|
| 553 |
+
|
| 554 |
+
# MODELS — the manifest is the authority, and it records the env
|
| 555 |
+
# needed to reproduce its own numbers (see teacher_v2.json).
|
| 556 |
+
models = []
|
| 557 |
+
for tag in ("teacher_v2", "teacher_v1"):
|
| 558 |
+
p = Path(__file__).resolve().parents[2] / f"models/{tag}.json"
|
| 559 |
+
if p.exists():
|
| 560 |
+
try:
|
| 561 |
+
m = json.loads(p.read_text())
|
| 562 |
+
models.append({
|
| 563 |
+
"id": tag, "kind": "teacher",
|
| 564 |
+
"yield": m["metric"].get("mean_yield"),
|
| 565 |
+
"prec": m["metric"].get("mean_prec"),
|
| 566 |
+
"env": m.get("env", {}),
|
| 567 |
+
"reproduce": m.get("reproduce", ""),
|
| 568 |
+
"note": "cosine channels -> RRF -> PRF -> ITM cascade "
|
| 569 |
+
"-> anchor + event gate + density -> cut"})
|
| 570 |
+
break
|
| 571 |
+
except Exception:
|
| 572 |
+
pass
|
| 573 |
+
sp = Path(__file__).resolve().parents[2] / "models/student_v1/meta.json"
|
| 574 |
+
if sp.exists():
|
| 575 |
+
try:
|
| 576 |
+
sm = json.loads(sp.read_text())
|
| 577 |
+
models.append({
|
| 578 |
+
"id": "student_v1", "kind": "student",
|
| 579 |
+
"params_M": round(sm.get("params", 0) / 1e6, 2),
|
| 580 |
+
"read_ms": 27, "write_s_per_demo": 0.85,
|
| 581 |
+
"note": "two-tower bi-encoder + listwise rerank head. "
|
| 582 |
+
"Produces the five elements ITSELF; stage order "
|
| 583 |
+
"PRF -> gate -> cascade was measured, not "
|
| 584 |
+
"inherited (the teacher's order scored worse)"})
|
| 585 |
+
except Exception:
|
| 586 |
+
pass
|
| 587 |
+
|
| 588 |
+
# ARTIFACTS lane: the model files and fitted state this store's query
|
| 589 |
+
# path actually loads. It was empty because nothing ever emitted a
|
| 590 |
+
# node of kind "model" - the lane existed with nothing to put in it.
|
| 591 |
+
RT = Path(__file__).resolve().parents[2]
|
| 592 |
+
|
| 593 |
+
def _sz(p):
|
| 594 |
+
p = Path(p)
|
| 595 |
+
if p.is_dir():
|
| 596 |
+
return sum(f.stat().st_size for f in p.rglob("*") if f.is_file())
|
| 597 |
+
return p.stat().st_size if p.exists() else 0
|
| 598 |
+
|
| 599 |
+
for path, label, note in (
|
| 600 |
+
("models/teacher_v2.json", "teacher_v2", "manifest: tables, "
|
| 601 |
+
"digests, stages, env, measured yield/prec"),
|
| 602 |
+
("models/student_v1", "student_v1", "student.pt (two-tower + "
|
| 603 |
+
"rerank head), episode_emb.npz, namer.pt"),
|
| 604 |
+
("models/iv2_stage2_1b", "iv2_stage2_1b", "InternVideo2-Stage2 "
|
| 605 |
+
"1B — the ITM cross-encoder, cost-gated"),
|
| 606 |
+
("models/fdnnv", "fdnnv", "FDNN-V encoder — embeds every frame "
|
| 607 |
+
"at ingest, 1,109 fps"),
|
| 608 |
+
("artifacts/verbs_v2.json", "verbs_v2", "geometry verb partition"),
|
| 609 |
+
("artifacts/cavity.json", "cavity", "cavity/articulation thresholds"),
|
| 610 |
+
):
|
| 611 |
+
b = _sz(RT / path)
|
| 612 |
+
if b:
|
| 613 |
+
nodes.append({"id": label, "kind": "model", "rows": None,
|
| 614 |
+
"bytes": b, "note": note, "path": path})
|
| 615 |
+
for a, note in (("_set_weights.json", "fitted selection weights "
|
| 616 |
+
"(+ .prev rotation, .loqo holdout)"),
|
| 617 |
+
("_channel_weights.json", "per-channel fusion weights"),
|
| 618 |
+
("_vocab.json", "corpus-attested vocabulary cache")):
|
| 619 |
+
b = _sz(db.dir / a)
|
| 620 |
+
if b:
|
| 621 |
+
nodes.append({"id": a, "kind": "model", "rows": None,
|
| 622 |
+
"bytes": b, "note": note, "path": a})
|
| 623 |
+
|
| 624 |
+
# THE TRAINING / RETRAINING LOOP — how more data becomes a better
|
| 625 |
+
# model. This is a closed loop and the store is inside it.
|
| 626 |
+
training = [
|
| 627 |
+
{"id": "ingest", "label": "1 · INGEST",
|
| 628 |
+
"note": "raw video → frames (byte-range index) + FDNN-V vectors "
|
| 629 |
+
"for EVERY frame. 63.5 s for 3.91 h, 1,109 fps. No "
|
| 630 |
+
"labels, no metadata: the store never ingests task "
|
| 631 |
+
"strings."},
|
| 632 |
+
{"id": "elements", "label": "2 · WRITE-PATH ELEMENTS",
|
| 633 |
+
"note": "geometry over the frames produces the five elements "
|
| 634 |
+
"per demo — agent from flow, participants by causality, "
|
| 635 |
+
"events from cavity + displacement, answer rollup. The "
|
| 636 |
+
"student's Namer head predicts the participant NAME "
|
| 637 |
+
"VECTOR the 7B VLM would have produced."},
|
| 638 |
+
{"id": "teacher", "label": "3 · TEACHER RANKS",
|
| 639 |
+
"note": "the expensive path (channels → RRF → PRF → ITM 1B "
|
| 640 |
+
"cross-encoder → anchor + gate) ranks the corpus for "
|
| 641 |
+
"generated queries. 2–75 s/query, so it is never the "
|
| 642 |
+
"serving path — it exists to produce labels."},
|
| 643 |
+
{"id": "distill", "label": "4 · DISTILL",
|
| 644 |
+
"note": "teacher rankings over 200 corpus-vocabulary queries "
|
| 645 |
+
"become listwise targets for the student's two-tower "
|
| 646 |
+
"+ rerank head. 645 s to label, 8 s to train, 3.28M "
|
| 647 |
+
"params. The eval truthset is NEVER trained on."},
|
| 648 |
+
{"id": "serve", "label": "5 · SERVE",
|
| 649 |
+
"note": "student answers in 27 ms: one text forward + one "
|
| 650 |
+
"matmul over precomputed episode vectors, then the "
|
| 651 |
+
"structural gate reads columns the write path already "
|
| 652 |
+
"produced."},
|
| 653 |
+
{"id": "grow", "label": "6 · MORE DATA → BETTER MODEL",
|
| 654 |
+
"note": "adding episodes widens every corpus statistic the "
|
| 655 |
+
"system is built on: the transition anchors get more "
|
| 656 |
+
"attesting episodes so their earned reliability rises, "
|
| 657 |
+
"PRF anchors sharpen, the attested vocabulary grows, "
|
| 658 |
+
"and the teacher has more to label. Re-running step 3–4 "
|
| 659 |
+
"is the retrain; nothing here needs human annotation."},
|
| 660 |
+
]
|
| 661 |
+
|
| 662 |
+
total_rows = sum(n.get("rows") or 0 for n in nodes)
|
| 663 |
+
total_bytes = sum(n.get("bytes") or 0 for n in nodes)
|
| 664 |
+
return {"store": db.name, "key": key, "nodes": nodes, "edges": edges,
|
| 665 |
+
"channels": channels, "pipeline": pipeline,
|
| 666 |
+
"elements": elements, "models": models, "training": training,
|
| 667 |
+
"fitted": {k: fitted[k] for k in
|
| 668 |
+
("set_weights_dir", "set_weights", "cut_alpha_dir",
|
| 669 |
+
"cut_alpha", "nms_r_dir", "nms_r", "loqo_mean")
|
| 670 |
+
if k in fitted},
|
| 671 |
+
"totals": {"rows": total_rows, "bytes": total_bytes,
|
| 672 |
+
"tables": len([n for n in nodes
|
| 673 |
+
if n["kind"] not in
|
| 674 |
+
("media", "model")])}}
|
| 675 |
+
|
| 676 |
+
|
| 677 |
+
def api_bytes(key: str):
|
| 678 |
+
"""WHERE THE BYTES ACTUALLY ARE, and how much compression bought.
|
| 679 |
+
|
| 680 |
+
`du` on this store says 3.9 GB while the tables hold 249 MB, and the
|
| 681 |
+
gap is not a compression failure - it is three different things that
|
| 682 |
+
a single directory size silently adds together:
|
| 683 |
+
|
| 684 |
+
live the parquet the LOG currently points at. The database.
|
| 685 |
+
orphaned parquet superseded by a replace/compact commit. Removed
|
| 686 |
+
from the active set, still on disk until vacuum.
|
| 687 |
+
media managed video renditions, byte-range decoded.
|
| 688 |
+
cache _cache/ - the ITM cross-encoder's vision tokens, 3 GB
|
| 689 |
+
of it. DISPOSABLE: deleting it costs recompute time and
|
| 690 |
+
never correctness, and it is not the database.
|
| 691 |
+
|
| 692 |
+
Models are NOT in the store. They live in models/ at the repo root
|
| 693 |
+
(iv2_stage2_1b alone is 2.8 GB), which is why "the store" and "what
|
| 694 |
+
this system needs on disk" are different questions.
|
| 695 |
+
"""
|
| 696 |
+
import pyarrow.parquet as pq
|
| 697 |
+
db = STORES[key]
|
| 698 |
+
live = orphan = on_disk = 0
|
| 699 |
+
rows = []
|
| 700 |
+
for t in sorted(db.tables()):
|
| 701 |
+
st = db.table(t).state()
|
| 702 |
+
keep = {f.path for f in st.files}
|
| 703 |
+
c = u = 0
|
| 704 |
+
encs = set()
|
| 705 |
+
comp = set()
|
| 706 |
+
d = db.dir / "tables" / t
|
| 707 |
+
for f in d.glob("*.parquet"):
|
| 708 |
+
b = f.stat().st_size
|
| 709 |
+
on_disk += b
|
| 710 |
+
if f.name not in keep:
|
| 711 |
+
orphan += b
|
| 712 |
+
continue
|
| 713 |
+
try:
|
| 714 |
+
md = pq.ParquetFile(f).metadata
|
| 715 |
+
for g in range(md.num_row_groups):
|
| 716 |
+
rg = md.row_group(g)
|
| 717 |
+
for j in range(rg.num_columns):
|
| 718 |
+
col = rg.column(j)
|
| 719 |
+
c += col.total_compressed_size
|
| 720 |
+
u += col.total_uncompressed_size
|
| 721 |
+
for e in (col.encodings or ()):
|
| 722 |
+
encs.add(str(e))
|
| 723 |
+
if col.compression:
|
| 724 |
+
comp.add(str(col.compression))
|
| 725 |
+
except Exception:
|
| 726 |
+
pass
|
| 727 |
+
live += st.bytes
|
| 728 |
+
rows.append({"table": t, "rows": st.rows, "compressed": c,
|
| 729 |
+
"uncompressed": u,
|
| 730 |
+
"ratio": round(u / max(c, 1), 2),
|
| 731 |
+
"encodings": sorted(encs), "codec": sorted(comp),
|
| 732 |
+
"files": len(st.files), "version": st.version})
|
| 733 |
+
rows.sort(key=lambda r: -r["compressed"])
|
| 734 |
+
media = sum(p.lstat().st_size for p in (db.dir / "media").glob("*")
|
| 735 |
+
if p.is_file()) if (db.dir / "media").is_dir() else 0
|
| 736 |
+
cdir = db.dir / "_cache"
|
| 737 |
+
cache = sum(f.stat().st_size for f in cdir.rglob("*")
|
| 738 |
+
if f.is_file()) if cdir.is_dir() else 0
|
| 739 |
+
cparts = ([{"name": p.name,
|
| 740 |
+
"bytes": sum(f.stat().st_size for f in p.rglob("*")
|
| 741 |
+
if f.is_file())}
|
| 742 |
+
for p in cdir.iterdir() if p.is_dir()] if cdir.is_dir()
|
| 743 |
+
else [])
|
| 744 |
+
tc = sum(r["compressed"] for r in rows)
|
| 745 |
+
tu = sum(r["uncompressed"] for r in rows)
|
| 746 |
+
raw = _raw_source_bytes(db, db.describe())
|
| 747 |
+
return {"store": db.name, "key": key, "tables": rows,
|
| 748 |
+
"live": live, "orphaned": orphan, "on_disk": on_disk,
|
| 749 |
+
"media": media, "cache": cache, "cache_parts": cparts,
|
| 750 |
+
"total": on_disk + media + cache,
|
| 751 |
+
"compressed": tc, "uncompressed": tu,
|
| 752 |
+
"ratio": round(tu / max(tc, 1), 2),
|
| 753 |
+
"raw_source": raw["bytes"], "raw_files": raw["files"]}
|
| 754 |
+
|
| 755 |
+
|
| 756 |
+
def api_dbinternals(key: str, table: str | None = None):
|
| 757 |
+
"""The storage engine, as it actually is on disk.
|
| 758 |
+
|
| 759 |
+
Not a diagram: every number here is read from the transaction log
|
| 760 |
+
and the Parquet footers of this store, right now. Four things a
|
| 761 |
+
storage engine has to be able to show:
|
| 762 |
+
|
| 763 |
+
1. the LOG — the table is the fold of an append-only list of
|
| 764 |
+
JSON commits, each carrying its own file list
|
| 765 |
+
2. FILES — with the zone map (min_ts/max_ts) that lets a
|
| 766 |
+
query drop a whole file without opening it
|
| 767 |
+
3. PAGES — row groups inside a file, each with its own ts
|
| 768 |
+
statistics and per-column chunk sizes
|
| 769 |
+
4. PRUNING — what a real window query touches, in bytes,
|
| 770 |
+
across both layers
|
| 771 |
+
"""
|
| 772 |
+
import pyarrow.parquet as pq
|
| 773 |
+
db = STORES[key]
|
| 774 |
+
names = sorted(db.tables())
|
| 775 |
+
table = table if table in names else (
|
| 776 |
+
"frames" if "frames" in names else names[0])
|
| 777 |
+
t = db.table(table)
|
| 778 |
+
st = t.state()
|
| 779 |
+
|
| 780 |
+
# ---- 1. the log: manifest commits, newest last
|
| 781 |
+
# read the RAW commit files: Table.history() summarises and drops the
|
| 782 |
+
# removal count, which is the whole point of an op=replace entry.
|
| 783 |
+
log = []
|
| 784 |
+
ldir = db.dir / "tables" / table / "_log"
|
| 785 |
+
for f in sorted(ldir.glob("*.json")) if ldir.is_dir() else []:
|
| 786 |
+
try:
|
| 787 |
+
e = json.loads(f.read_text())
|
| 788 |
+
except Exception:
|
| 789 |
+
continue
|
| 790 |
+
add = e.get("add") or []
|
| 791 |
+
log.append({
|
| 792 |
+
"version": int(f.stem), "op": e.get("op"), "kind": e.get("kind"),
|
| 793 |
+
"added": len(add),
|
| 794 |
+
"removed": len(e.get("remove") or []),
|
| 795 |
+
"added_rows": sum(a.get("rows", 0) for a in add),
|
| 796 |
+
"added_bytes": sum(a.get("bytes", 0) for a in add),
|
| 797 |
+
"ts": e.get("ts_utc"),
|
| 798 |
+
"meta": {k: v for k, v in (e.get("meta") or {}).items()
|
| 799 |
+
if not isinstance(v, (list, dict))},
|
| 800 |
+
})
|
| 801 |
+
|
| 802 |
+
# ---- 2/3. files and their row groups (pages)
|
| 803 |
+
files, pages = [], []
|
| 804 |
+
total_rg = 0
|
| 805 |
+
for f in st.files:
|
| 806 |
+
p = db.dir / "tables" / table / f.path
|
| 807 |
+
row = {"path": f.path, "rows": f.rows, "bytes": f.bytes,
|
| 808 |
+
"min_ts": f.min_ts, "max_ts": f.max_ts, "row_groups": None,
|
| 809 |
+
"footer_bytes": None}
|
| 810 |
+
try:
|
| 811 |
+
pf = pq.ParquetFile(p)
|
| 812 |
+
md = pf.metadata
|
| 813 |
+
row["row_groups"] = md.num_row_groups
|
| 814 |
+
row["footer_bytes"] = md.serialized_size
|
| 815 |
+
total_rg += md.num_row_groups
|
| 816 |
+
ts_i = (md.schema.names.index("ts")
|
| 817 |
+
if "ts" in md.schema.names else 0)
|
| 818 |
+
if len(pages) < 24: # a readable sample, not all
|
| 819 |
+
for g in range(min(md.num_row_groups, 8)):
|
| 820 |
+
rg = md.row_group(g)
|
| 821 |
+
s = rg.column(ts_i).statistics
|
| 822 |
+
cols = []
|
| 823 |
+
for c in range(rg.num_columns):
|
| 824 |
+
col = rg.column(c)
|
| 825 |
+
cols.append({
|
| 826 |
+
"name": (md.schema.names[c]
|
| 827 |
+
if c < len(md.schema.names) else "?"),
|
| 828 |
+
"compressed": col.total_compressed_size,
|
| 829 |
+
"uncompressed": col.total_uncompressed_size,
|
| 830 |
+
"encodings": [str(e) for e in
|
| 831 |
+
(col.encodings or [])][:3],
|
| 832 |
+
})
|
| 833 |
+
cols.sort(key=lambda x: -x["compressed"])
|
| 834 |
+
pages.append({
|
| 835 |
+
"file": f.path[:18], "group": g,
|
| 836 |
+
"rows": rg.num_rows, "bytes": rg.total_byte_size,
|
| 837 |
+
"min_ts": getattr(s, "min", None),
|
| 838 |
+
"max_ts": getattr(s, "max", None),
|
| 839 |
+
"columns": cols[:6],
|
| 840 |
+
"n_columns": rg.num_columns})
|
| 841 |
+
except Exception:
|
| 842 |
+
pass
|
| 843 |
+
files.append(row)
|
| 844 |
+
|
| 845 |
+
# ---- 4. pruning, executed for real on a 2% slice of the span
|
| 846 |
+
prune = None
|
| 847 |
+
if st.files and st.min_ts is not None:
|
| 848 |
+
# ANCHOR THE WINDOW ON REAL DATA. Taking the midpoint of
|
| 849 |
+
# min_ts..max_ts lands in one of the 60 s gaps this store puts
|
| 850 |
+
# between demos and returns zero rows, which measures nothing.
|
| 851 |
+
# A row group's own ts statistics are, by construction, a range
|
| 852 |
+
# that contains rows - so the demo window is the middle row
|
| 853 |
+
# group of the middle file.
|
| 854 |
+
t0, t1 = st.min_ts, st.min_ts + max(
|
| 855 |
+
(st.max_ts - st.min_ts) // 50, 1)
|
| 856 |
+
try:
|
| 857 |
+
mid = st.files[len(st.files) // 2]
|
| 858 |
+
pf = pq.ParquetFile(db.dir / "tables" / table / mid.path)
|
| 859 |
+
md = pf.metadata
|
| 860 |
+
ts_i = (md.schema.names.index("ts")
|
| 861 |
+
if "ts" in md.schema.names else 0)
|
| 862 |
+
g = md.row_group(md.num_row_groups // 2)
|
| 863 |
+
s = g.column(ts_i).statistics
|
| 864 |
+
if s is not None and s.min is not None:
|
| 865 |
+
t0, t1 = int(s.min), int(s.max)
|
| 866 |
+
except Exception:
|
| 867 |
+
pass
|
| 868 |
+
from .store import QueryStats
|
| 869 |
+
qs = QueryStats()
|
| 870 |
+
try:
|
| 871 |
+
got = t.scan(t0, t1, stats=qs)
|
| 872 |
+
prune = {
|
| 873 |
+
"window_ns": int(t1 - t0),
|
| 874 |
+
"files_total": qs.files_total,
|
| 875 |
+
"files_touched": qs.files_touched,
|
| 876 |
+
"corpus_bytes": qs.corpus_bytes,
|
| 877 |
+
"bytes_touched": qs.bytes_touched,
|
| 878 |
+
"rows_returned": len(got),
|
| 879 |
+
"elided_pct": (round(100.0 * (qs.corpus_bytes
|
| 880 |
+
- qs.bytes_touched)
|
| 881 |
+
/ max(qs.corpus_bytes, 1), 2)),
|
| 882 |
+
}
|
| 883 |
+
except Exception as e:
|
| 884 |
+
prune = {"error": f"{type(e).__name__}: {e}"[:120]}
|
| 885 |
+
|
| 886 |
+
# ---- THE LADDER: the containment hierarchy, one rung per level,
|
| 887 |
+
# with this store's real numbers on each. Ordered smallest first.
|
| 888 |
+
lake_dir = db.dir.parent
|
| 889 |
+
n_stores = sum(1 for p in lake_dir.iterdir()
|
| 890 |
+
if p.is_dir() and (p / "_store.json").exists())
|
| 891 |
+
f0 = files[0] if files else None # the dict built above
|
| 892 |
+
rg0 = pages[0] if pages else None
|
| 893 |
+
ch0 = (rg0 or {}).get("columns", [{}])[0] if rg0 else {}
|
| 894 |
+
PAGE_TARGET = 1 << 20 # parquet default 1 MiB
|
| 895 |
+
est_pages = max(1, round((ch0.get("compressed") or 0) / PAGE_TARGET)) \
|
| 896 |
+
if ch0 else None
|
| 897 |
+
ladder = [
|
| 898 |
+
{"id": "value", "label": "VALUE / ROW",
|
| 899 |
+
"n": f"{st.rows:,} rows",
|
| 900 |
+
"sub": f"{(rg0 or {}).get('n_columns', 0)} columns",
|
| 901 |
+
"meta": "—",
|
| 902 |
+
"why": "One cell. Rows are never stored contiguously: inside a "
|
| 903 |
+
"row group the data is laid out COLUMN BY COLUMN, which "
|
| 904 |
+
"is what lets a query read one column and skip the rest."},
|
| 905 |
+
{"id": "page", "label": "PAGE",
|
| 906 |
+
"n": (f"~{est_pages} per column chunk" if est_pages else "—"),
|
| 907 |
+
"sub": "~1 MiB target",
|
| 908 |
+
"meta": "page header: encoding, value count, (optional) stats",
|
| 909 |
+
"why": "The atomic unit of compression and decode. You cannot "
|
| 910 |
+
"read half a page - it is decompressed whole - so page "
|
| 911 |
+
"size is the floor on random-access cost."},
|
| 912 |
+
{"id": "chunk", "label": "COLUMN CHUNK",
|
| 913 |
+
"n": f"{(rg0 or {}).get('n_columns', 0)} per row group",
|
| 914 |
+
"sub": (f"{ch0.get('name','')} "
|
| 915 |
+
f"{(ch0.get('encodings') or [None])[0] or ''}"),
|
| 916 |
+
"meta": "offset, size, encodings, compression, min/max/nulls",
|
| 917 |
+
"why": "All the pages of ONE column within ONE row group, "
|
| 918 |
+
"contiguous on disk. This is the unit PROJECTION "
|
| 919 |
+
"pushdown skips: ask for 2 of 11 columns and the other "
|
| 920 |
+
"9 chunks are never read."},
|
| 921 |
+
{"id": "rowgroup", "label": "ROW GROUP",
|
| 922 |
+
"n": f"{total_rg} in this table",
|
| 923 |
+
"sub": (f"{rg0['rows']:,} rows · {rg0['bytes']:,} B"
|
| 924 |
+
if rg0 else "—"),
|
| 925 |
+
"meta": "per-column statistics: min, max, null_count",
|
| 926 |
+
"why": "A horizontal slice of rows holding every column's "
|
| 927 |
+
"chunk. Its ts statistics are what PREDICATE pushdown "
|
| 928 |
+
"tests, so a non-overlapping group is never "
|
| 929 |
+
"decompressed. Sized by BYTES (8 MB target), not rows."},
|
| 930 |
+
{"id": "footer", "label": "FOOTER (FileMetaData)",
|
| 931 |
+
"n": (f"{f0.get('footer_bytes') or 0:,} B" if f0 else "—"),
|
| 932 |
+
"sub": "at the END of the file",
|
| 933 |
+
"meta": "THE schema + every row group's metadata + offsets",
|
| 934 |
+
"why": "Written last so the file streams out in one pass, read "
|
| 935 |
+
"first so one seek reveals the whole layout. NOTE: row "
|
| 936 |
+
"group metadata lives HERE, inside the same file - not "
|
| 937 |
+
"in a separate meta file. Separate meta files start one "
|
| 938 |
+
"level up."},
|
| 939 |
+
{"id": "file", "label": "PARQUET FILE",
|
| 940 |
+
"n": f"{len(st.files)} in this table",
|
| 941 |
+
"sub": (f"{f0['rows']:,} rows · {f0['bytes']:,} B" if f0 else "—"),
|
| 942 |
+
"meta": "immutable; never edited in place",
|
| 943 |
+
"why": "The unit of atomic addition and removal. Rewriting is "
|
| 944 |
+
"how you 'edit', which is what makes snapshots cheap."},
|
| 945 |
+
{"id": "commit", "label": "COMMIT (manifest file)",
|
| 946 |
+
"n": f"{len(log)} in tables/{table}/_log/",
|
| 947 |
+
"sub": (f"latest: v{log[-1]['version']} {log[-1]['op']} "
|
| 948 |
+
f"+{log[-1]['added']}"
|
| 949 |
+
f"{' −' + str(log[-1]['removed']) if log[-1]['removed'] else ''}"
|
| 950 |
+
if log else "—"),
|
| 951 |
+
"meta": "op, schema, file list WITH zone maps (min_ts/max_ts)",
|
| 952 |
+
"why": "THE separate meta file. One JSON per commit, listing "
|
| 953 |
+
"the files this version contains and each file's time "
|
| 954 |
+
"range - so a window query drops whole files here, "
|
| 955 |
+
"before opening a single footer. Iceberg calls this a "
|
| 956 |
+
"manifest; Delta calls it a log entry."},
|
| 957 |
+
{"id": "log", "label": "LOG → TABLE STATE",
|
| 958 |
+
"n": f"version {st.version}",
|
| 959 |
+
"sub": f"state = fold of {len(log)} commits",
|
| 960 |
+
"meta": "the fold: adds minus removes, in order",
|
| 961 |
+
"why": "A table is NOT the files in its directory - it is the "
|
| 962 |
+
"result of replaying this log. That is what buys "
|
| 963 |
+
"snapshot isolation (v{N} is immutable forever), atomic "
|
| 964 |
+
"multi-file commits, and time travel."},
|
| 965 |
+
{"id": "table", "label": "TABLE (+ schema)",
|
| 966 |
+
"n": f"{len(names)} in this store",
|
| 967 |
+
"sub": f"{table}: {st.rows:,} rows · {st.bytes:,} B",
|
| 968 |
+
"meta": "schema travels with each commit (schema-on-log)",
|
| 969 |
+
"why": "Schema evolution is an append, never a rewrite. Every "
|
| 970 |
+
"table must carry ts (int64 ns) sorted within a file - "
|
| 971 |
+
"the one schema law, and why time is the primary axis."},
|
| 972 |
+
{"id": "store", "label": "STORE",
|
| 973 |
+
"n": f"{len(names)} tables",
|
| 974 |
+
"sub": f"{db.name} · {sum((db.table(x).state().bytes) for x in names):,} B",
|
| 975 |
+
"meta": "_store.json + fitted artifacts + caches + media/",
|
| 976 |
+
"why": "A directory of tables plus sidecar state. Sidecars are "
|
| 977 |
+
"graded: authoritative (_store.json), fitted "
|
| 978 |
+
"(_set_weights.json), cache (_vocab.json), disposable "
|
| 979 |
+
"(_cache/). Raw media is referenced in place, never "
|
| 980 |
+
"copied in."},
|
| 981 |
+
{"id": "lake", "label": "LAKE",
|
| 982 |
+
"n": f"{n_stores} store{'' if n_stores == 1 else 's'}",
|
| 983 |
+
"sub": "lake/ (" + lake_dir.name + ")",
|
| 984 |
+
"meta": "plain directories — no catalog service",
|
| 985 |
+
"why": "Many stores side by side. Nothing above this is "
|
| 986 |
+
"needed: the format is open, so any engine (DuckDB, "
|
| 987 |
+
"Spark, pandas) reads these files directly without "
|
| 988 |
+
"going through us."},
|
| 989 |
+
]
|
| 990 |
+
|
| 991 |
+
return {
|
| 992 |
+
"store": db.name, "key": key, "table": table, "tables": names,
|
| 993 |
+
"version": st.version, "kind": st.kind, "rows": st.rows,
|
| 994 |
+
"bytes": st.bytes, "n_files": len(st.files),
|
| 995 |
+
"n_row_groups": total_rg, "ladder": ladder,
|
| 996 |
+
"schema": [str(x) for x in str(st.schema).split("\n") if x][:24],
|
| 997 |
+
"log": log[-12:], "log_total": len(log),
|
| 998 |
+
"files": files[:12], "pages": pages, "prune": prune,
|
| 999 |
+
"log_dir": f"tables/{table}/_log/",
|
| 1000 |
+
"store_files": sorted(
|
| 1001 |
+
p.name for p in db.dir.glob("*")
|
| 1002 |
+
if p.is_file())[:20],
|
| 1003 |
+
}
|
| 1004 |
+
|
| 1005 |
+
|
| 1006 |
+
def api_analytics(key: str):
|
| 1007 |
+
"""Operational analytics, general to ANY store: everything here is
|
| 1008 |
+
derived from transaction logs, parquet footers, and table meta.
|
| 1009 |
+
No data pages are read and nothing is dataset-specific; a store of
|
| 1010 |
+
factory video, dashcam runs, or plain sensor CSVs renders the same
|
| 1011 |
+
panels."""
|
| 1012 |
+
import pyarrow.parquet as pq
|
| 1013 |
+
ck = ("analytics", key,
|
| 1014 |
+
tuple(sorted((d["table"], d["version"])
|
| 1015 |
+
for d in STORES[key].describe())))
|
| 1016 |
+
if ck in _CACHE:
|
| 1017 |
+
return _CACHE[ck]
|
| 1018 |
+
db = STORES[key]
|
| 1019 |
+
desc = db.describe()
|
| 1020 |
+
lo = min((d["min_ts"] for d in desc
|
| 1021 |
+
if d["rows"] and d.get("min_ts")), default=0)
|
| 1022 |
+
hi = max((d["max_ts"] for d in desc
|
| 1023 |
+
if d["rows"] and d.get("max_ts")), default=lo + 1)
|
| 1024 |
+
span = max(hi - lo, 1)
|
| 1025 |
+
|
| 1026 |
+
# write history straight off the transaction logs
|
| 1027 |
+
commits = []
|
| 1028 |
+
for d in desc:
|
| 1029 |
+
try:
|
| 1030 |
+
for c in db.table(d["table"]).history():
|
| 1031 |
+
commits.append({"table": d["table"],
|
| 1032 |
+
"version": c.get("version"),
|
| 1033 |
+
"op": c.get("op", ""),
|
| 1034 |
+
"rows": c.get("added_rows", 0),
|
| 1035 |
+
"ts_utc": c.get("ts_utc", "")})
|
| 1036 |
+
except Exception:
|
| 1037 |
+
pass
|
| 1038 |
+
commits.sort(key=lambda c: c["ts_utc"])
|
| 1039 |
+
|
| 1040 |
+
# temporal density from ROW-GROUP footer stats only: rows per time
|
| 1041 |
+
# bucket per table. The row group is the pruning unit, so this is
|
| 1042 |
+
# literally the elision map a range query sees.
|
| 1043 |
+
buckets_n = 64
|
| 1044 |
+
density = {}
|
| 1045 |
+
for d in desc:
|
| 1046 |
+
if not d["rows"] or not d.get("min_ts"):
|
| 1047 |
+
continue
|
| 1048 |
+
st = db.table(d["table"]).state()
|
| 1049 |
+
buckets = [0.0] * buckets_n
|
| 1050 |
+
try:
|
| 1051 |
+
for f in st.files:
|
| 1052 |
+
pf = pq.ParquetFile(db.dir / "tables" / d["table"]
|
| 1053 |
+
/ f.path)
|
| 1054 |
+
md = pf.metadata
|
| 1055 |
+
names = md.schema.names
|
| 1056 |
+
if "ts" not in names:
|
| 1057 |
+
continue
|
| 1058 |
+
ti = names.index("ts")
|
| 1059 |
+
for g in range(md.num_row_groups):
|
| 1060 |
+
rg = md.row_group(g)
|
| 1061 |
+
s = rg.column(ti).statistics
|
| 1062 |
+
if not s or s.min is None:
|
| 1063 |
+
continue
|
| 1064 |
+
a = int((s.min - lo) * buckets_n // span)
|
| 1065 |
+
b = int((s.max - lo) * buckets_n // span)
|
| 1066 |
+
a = min(max(a, 0), buckets_n - 1)
|
| 1067 |
+
b = min(max(b, a), buckets_n - 1)
|
| 1068 |
+
per = rg.num_rows / (b - a + 1)
|
| 1069 |
+
for i in range(a, b + 1):
|
| 1070 |
+
buckets[i] += per
|
| 1071 |
+
except Exception:
|
| 1072 |
+
continue
|
| 1073 |
+
if sum(buckets) > 0:
|
| 1074 |
+
density[d["table"]] = [int(round(x)) for x in buckets]
|
| 1075 |
+
|
| 1076 |
+
# vector inventory: every embeddings-kind table, with coverage
|
| 1077 |
+
# against the store's episode base when one exists
|
| 1078 |
+
base_rows = next((d["rows"] for d in desc
|
| 1079 |
+
if d["table"] == "episodes" and d["rows"]), None)
|
| 1080 |
+
vectors = []
|
| 1081 |
+
for d in desc:
|
| 1082 |
+
if d["kind"] != "embeddings" or not d["rows"]:
|
| 1083 |
+
continue
|
| 1084 |
+
meta = d.get("meta") or {}
|
| 1085 |
+
vectors.append({
|
| 1086 |
+
"table": d["table"], "rows": d["rows"], "bytes": d["bytes"],
|
| 1087 |
+
"dim": meta.get("dim"),
|
| 1088 |
+
"model": str(meta.get("model", ""))[:60],
|
| 1089 |
+
"per_base": (round(d["rows"] / base_rows, 2)
|
| 1090 |
+
if base_rows else None)})
|
| 1091 |
+
|
| 1092 |
+
fitted = None
|
| 1093 |
+
swp = db.dir / "_set_weights.json"
|
| 1094 |
+
if swp.exists():
|
| 1095 |
+
try:
|
| 1096 |
+
fitted = json.loads(swp.read_text())
|
| 1097 |
+
except Exception:
|
| 1098 |
+
pass
|
| 1099 |
+
|
| 1100 |
+
ix = api_indexes(key)
|
| 1101 |
+
out = {
|
| 1102 |
+
"span": {"lo": lo, "hi": hi},
|
| 1103 |
+
"tables": [{"table": d["table"], "kind": d["kind"],
|
| 1104 |
+
"rows": d["rows"], "bytes": d["bytes"],
|
| 1105 |
+
"files": d["files"], "version": d["version"],
|
| 1106 |
+
"bpr": (round(d["bytes"] / d["rows"], 1)
|
| 1107 |
+
if d["rows"] else None)} for d in desc],
|
| 1108 |
+
"commits": commits[-48:],
|
| 1109 |
+
"commit_total": len(commits),
|
| 1110 |
+
"density": density, "buckets": buckets_n,
|
| 1111 |
+
"vectors": vectors, "base_rows": base_rows,
|
| 1112 |
+
"fitted": fitted,
|
| 1113 |
+
"indexes": {"bptree": ix["bptree"], "ann": ix["ann"]},
|
| 1114 |
+
}
|
| 1115 |
+
if len(_CACHE) > 32:
|
| 1116 |
+
_CACHE.clear()
|
| 1117 |
+
_CACHE[ck] = out
|
| 1118 |
+
return out
|
| 1119 |
+
|
| 1120 |
+
|
| 1121 |
+
def api_geo(key: str):
|
| 1122 |
+
"""Generic geo panel: any timeseries table with latitude+longitude."""
|
| 1123 |
+
db = STORES[key]
|
| 1124 |
+
for d in db.describe():
|
| 1125 |
+
if d["kind"] != "timeseries":
|
| 1126 |
+
continue
|
| 1127 |
+
cols = db.table(d["table"]).scan(columns=None)
|
| 1128 |
+
if {"latitude", "longitude"} <= set(cols.column_names):
|
| 1129 |
+
la = cols.column("latitude").to_numpy()
|
| 1130 |
+
lo = cols.column("longitude").to_numpy()
|
| 1131 |
+
ts = cols.column("ts").to_numpy()
|
| 1132 |
+
step = max(1, len(la) // 2500)
|
| 1133 |
+
return {"table": d["table"],
|
| 1134 |
+
"points": [{"la": float(la[i]), "lo": float(lo[i]),
|
| 1135 |
+
"t": int(ts[i])}
|
| 1136 |
+
for i in range(0, len(la), step)]}
|
| 1137 |
+
return {"points": []}
|
| 1138 |
+
|
| 1139 |
+
|
| 1140 |
+
def _as_frameset(db, fs):
|
| 1141 |
+
"""Tolerate frame tables whose registered kind is not frame_index
|
| 1142 |
+
(seen on stores assembled by filtering another store): window()
|
| 1143 |
+
then returns a plain Arrow table, which still carries the frame
|
| 1144 |
+
index columns and decodes fine once wrapped."""
|
| 1145 |
+
if fs is None or hasattr(fs, "decode"):
|
| 1146 |
+
return fs
|
| 1147 |
+
try:
|
| 1148 |
+
from .video import FrameSet
|
| 1149 |
+
return FrameSet(db, "frames", fs)
|
| 1150 |
+
except Exception:
|
| 1151 |
+
return None
|
| 1152 |
+
|
| 1153 |
+
|
| 1154 |
+
def api_thumb(key: str, stream: str, t: int, width: int = 360):
|
| 1155 |
+
from PIL import Image
|
| 1156 |
+
db = STORES[key]
|
| 1157 |
+
win, _ = db.window(t - 2_000_000_000, t + 2_000_000_000, tables=["frames"])
|
| 1158 |
+
fs = _as_frameset(db, win.get("frames"))
|
| 1159 |
+
if fs is None or len(fs) == 0:
|
| 1160 |
+
return None
|
| 1161 |
+
decoded = fs.decode(stream=stream or None, width=width, limit=1)
|
| 1162 |
+
if not decoded:
|
| 1163 |
+
return None
|
| 1164 |
+
img = Image.fromarray(decoded[0][1])
|
| 1165 |
+
rot = db.meta.get("display", {}).get("rotate", 0)
|
| 1166 |
+
if rot:
|
| 1167 |
+
img = img.rotate(rot, expand=True)
|
| 1168 |
+
buf = io.BytesIO()
|
| 1169 |
+
img.save(buf, "JPEG", quality=80)
|
| 1170 |
+
return buf.getvalue()
|
| 1171 |
+
|
| 1172 |
+
|
| 1173 |
+
def _audio_for_stream(db, stream: str, t0: int, t1: int):
|
| 1174 |
+
"""Find the sensor's audio table (e.g. 'Sensor 108/cam0' →
|
| 1175 |
+
sensor_108_audio), return mono WAV bytes for the window, or None."""
|
| 1176 |
+
import struct as _struct
|
| 1177 |
+
prefix = stream.split("/")[0].replace("/", "_").replace(" ", "_").lower()
|
| 1178 |
+
cand = f"{prefix}_audio"
|
| 1179 |
+
if cand not in db.tables():
|
| 1180 |
+
return None
|
| 1181 |
+
t = db.table(cand).scan(t0, t1, columns=["ts", "ch0"])
|
| 1182 |
+
if len(t) < 100:
|
| 1183 |
+
return None
|
| 1184 |
+
ts = t.column("ts").to_numpy()
|
| 1185 |
+
rate = int(round((len(ts) - 1) * 1e9 / max(int(ts[-1] - ts[0]), 1)))
|
| 1186 |
+
pcm = t.column("ch0").to_numpy().astype("<i2").tobytes()
|
| 1187 |
+
hdr = b"RIFF" + _struct.pack("<I", 36 + len(pcm)) + b"WAVEfmt " + \
|
| 1188 |
+
_struct.pack("<IHHIIHH", 16, 1, 1, rate, rate * 2, 2, 16) + \
|
| 1189 |
+
b"data" + _struct.pack("<I", len(pcm))
|
| 1190 |
+
return hdr + pcm
|
| 1191 |
+
|
| 1192 |
+
|
| 1193 |
+
def api_clip(key: str, stream: str, t0: int, t1: int, width: int = 640):
|
| 1194 |
+
"""Playable clip: byte-range decode of the window's frames → H.264 MP4
|
| 1195 |
+
(plus the sensor's microphone track when the store has one). Cached by
|
| 1196 |
+
parameters; the source media is only ever read, never touched."""
|
| 1197 |
+
import hashlib
|
| 1198 |
+
import subprocess
|
| 1199 |
+
import tempfile
|
| 1200 |
+
# Merged segments can be minutes long; the player previews the first 30 s
|
| 1201 |
+
# rather than refusing (full-range export belongs to the Python API).
|
| 1202 |
+
t1 = min(t1, t0 + 30_000_000_000)
|
| 1203 |
+
cache_dir = Path(tempfile.gettempdir()) / "elidedb_clips"
|
| 1204 |
+
cache_dir.mkdir(exist_ok=True)
|
| 1205 |
+
ck = hashlib.sha1(f"{key}|{stream}|{t0}|{t1}|{width}".encode()).hexdigest()
|
| 1206 |
+
out_path = cache_dir / f"{ck}.mp4"
|
| 1207 |
+
if out_path.exists():
|
| 1208 |
+
return out_path.read_bytes()
|
| 1209 |
+
|
| 1210 |
+
from PIL import Image
|
| 1211 |
+
db = STORES[key]
|
| 1212 |
+
win, _ = db.window(t0, t1, tables=["frames"])
|
| 1213 |
+
fs = _as_frameset(db, win.get("frames"))
|
| 1214 |
+
# Every failure below names itself. A <video> element cannot render an
|
| 1215 |
+
# error body, so the player fetches the clip and shows these strings —
|
| 1216 |
+
# "could not build a clip" with no reason is not a diagnosis.
|
| 1217 |
+
if fs is None or len(fs) == 0:
|
| 1218 |
+
return {"error": "no frames indexed in this window",
|
| 1219 |
+
"detail": f"{stream or 'all streams'} "
|
| 1220 |
+
f"{(t1 - t0) / 1e9:.2f}s window"}
|
| 1221 |
+
have = fs.streams()
|
| 1222 |
+
if stream and stream not in have:
|
| 1223 |
+
return {"error": f"stream '{stream}' has no frames here",
|
| 1224 |
+
"detail": f"streams present in this window: "
|
| 1225 |
+
f"{', '.join(have) or 'none'}"}
|
| 1226 |
+
try:
|
| 1227 |
+
decoded = fs.decode(stream=stream or None, width=width)
|
| 1228 |
+
except Exception as e:
|
| 1229 |
+
return {"error": f"decode failed: {type(e).__name__}", "detail": str(e)}
|
| 1230 |
+
if len(decoded) < 2:
|
| 1231 |
+
return {"error": "not enough decodable frames for a clip",
|
| 1232 |
+
"detail": f"{len(decoded)} frame(s) decoded from "
|
| 1233 |
+
f"{len(fs)} indexed"}
|
| 1234 |
+
rot = db.meta.get("display", {}).get("rotate", 0)
|
| 1235 |
+
span_s = max((decoded[-1][0] - decoded[0][0]) / 1e9, 0.1)
|
| 1236 |
+
fps = max(round((len(decoded) - 1) / span_s, 2), 1)
|
| 1237 |
+
|
| 1238 |
+
wav = _audio_for_stream(db, stream, decoded[0][0], decoded[-1][0])
|
| 1239 |
+
wav_path = None
|
| 1240 |
+
if wav:
|
| 1241 |
+
wav_path = cache_dir / f"{ck}.wav"
|
| 1242 |
+
wav_path.write_bytes(wav)
|
| 1243 |
+
from .fftools import find
|
| 1244 |
+
cmd = [find("ffmpeg"), "-v", "error", "-y",
|
| 1245 |
+
"-f", "image2pipe", "-framerate", str(fps), "-i", "-"]
|
| 1246 |
+
if wav_path:
|
| 1247 |
+
cmd += ["-i", str(wav_path)]
|
| 1248 |
+
cmd += ["-c:v", "libx264", "-preset", "veryfast", "-crf", "23",
|
| 1249 |
+
"-pix_fmt", "yuv420p", "-movflags", "+faststart"]
|
| 1250 |
+
if wav_path:
|
| 1251 |
+
cmd += ["-c:a", "aac", "-b:a", "96k", "-shortest"]
|
| 1252 |
+
cmd += [str(out_path)]
|
| 1253 |
+
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE,
|
| 1254 |
+
stderr=subprocess.PIPE)
|
| 1255 |
+
try:
|
| 1256 |
+
for (_ts, arr) in decoded:
|
| 1257 |
+
img = Image.fromarray(arr)
|
| 1258 |
+
if rot:
|
| 1259 |
+
img = img.rotate(rot, expand=True)
|
| 1260 |
+
# libx264 yuv420 needs even dimensions
|
| 1261 |
+
if img.width % 2 or img.height % 2:
|
| 1262 |
+
img = img.crop((0, 0, img.width & ~1, img.height & ~1))
|
| 1263 |
+
img.save(proc.stdin, "JPEG", quality=90)
|
| 1264 |
+
proc.stdin.close()
|
| 1265 |
+
except BrokenPipeError:
|
| 1266 |
+
pass # ffmpeg died early; the stderr below explains why
|
| 1267 |
+
err = proc.stderr.read().decode(errors="replace").strip()
|
| 1268 |
+
proc.wait()
|
| 1269 |
+
if wav_path:
|
| 1270 |
+
wav_path.unlink(missing_ok=True)
|
| 1271 |
+
if proc.returncode != 0 or not out_path.exists():
|
| 1272 |
+
return {"error": "ffmpeg could not mux the clip",
|
| 1273 |
+
"detail": err or "no output produced"}
|
| 1274 |
+
return out_path.read_bytes()
|
| 1275 |
+
|
| 1276 |
+
|
| 1277 |
+
def build_id() -> str:
|
| 1278 |
+
"""Identity of the code this server is actually running.
|
| 1279 |
+
|
| 1280 |
+
The launcher reuses whatever already listens on the port, so a server
|
| 1281 |
+
started from an older checkout keeps serving stale code forever — Python
|
| 1282 |
+
caches modules at import, so editing files changes nothing until restart.
|
| 1283 |
+
That is invisible from the browser and produces bug reports about
|
| 1284 |
+
behaviour that no longer exists in the source. The launcher now compares
|
| 1285 |
+
this against the on-disk files and restarts on a mismatch.
|
| 1286 |
+
"""
|
| 1287 |
+
import hashlib
|
| 1288 |
+
h = hashlib.sha1()
|
| 1289 |
+
for f in sorted(Path(__file__).parent.glob("*.py")) + \
|
| 1290 |
+
[Path(__file__).parent / "desk_ui.html"]:
|
| 1291 |
+
if f.exists():
|
| 1292 |
+
st = f.stat()
|
| 1293 |
+
h.update(f"{f.name}:{st.st_size}:{int(st.st_mtime)}".encode())
|
| 1294 |
+
return h.hexdigest()[:12]
|
| 1295 |
+
|
| 1296 |
+
|
| 1297 |
+
def api_schema(key: str, table: str | None = None):
|
| 1298 |
+
"""What the data actually IS: columns, types, and real sample rows.
|
| 1299 |
+
The first thing anyone opening a database wants to see."""
|
| 1300 |
+
import pyarrow as pa
|
| 1301 |
+
import pyarrow.parquet as pq
|
| 1302 |
+
db = STORES[key]
|
| 1303 |
+
names = [table] if table else db.tables()
|
| 1304 |
+
out = []
|
| 1305 |
+
for name in names:
|
| 1306 |
+
st = db.table(name).state()
|
| 1307 |
+
if not st.files:
|
| 1308 |
+
continue
|
| 1309 |
+
pf = pq.ParquetFile(db.dir / "tables" / name / st.files[0].path)
|
| 1310 |
+
schema = pf.schema_arrow
|
| 1311 |
+
cols = []
|
| 1312 |
+
for f in schema:
|
| 1313 |
+
t = str(f.type)
|
| 1314 |
+
if pa.types.is_fixed_size_list(f.type):
|
| 1315 |
+
t = f"vector[{f.type.list_size}]"
|
| 1316 |
+
cols.append({"name": f.name, "type": t})
|
| 1317 |
+
sample = []
|
| 1318 |
+
if table: # only materialise rows for the focused table
|
| 1319 |
+
head = db.table(name).scan(columns=[c["name"] for c in cols
|
| 1320 |
+
if not c["type"].startswith("vector")]
|
| 1321 |
+
).slice(0, 8).to_pylist()
|
| 1322 |
+
for row in head:
|
| 1323 |
+
sample.append({k: (round(v, 6) if isinstance(v, float) else v)
|
| 1324 |
+
for k, v in row.items()})
|
| 1325 |
+
out.append({"table": name, "kind": st.kind, "rows": st.rows,
|
| 1326 |
+
"bytes": st.bytes, "version": st.version,
|
| 1327 |
+
"min_ts": st.min_ts, "max_ts": st.max_ts,
|
| 1328 |
+
"columns": cols, "sample": sample})
|
| 1329 |
+
return out
|
| 1330 |
+
|
| 1331 |
+
|
| 1332 |
+
def api_indexes(key: str):
|
| 1333 |
+
"""Every index in the store: B+ trees per table + ANN tiers on
|
| 1334 |
+
embeddings, with size and the data version each was built for."""
|
| 1335 |
+
db = STORES[key]
|
| 1336 |
+
out = {"bptree": [], "ann": []}
|
| 1337 |
+
for name in db.tables():
|
| 1338 |
+
ixdir = db.dir / "tables" / name / "_index"
|
| 1339 |
+
if not ixdir.is_dir():
|
| 1340 |
+
continue
|
| 1341 |
+
for p in ixdir.glob("*.bpt"):
|
| 1342 |
+
col, ver = p.stem.rsplit(".v", 1)
|
| 1343 |
+
out["bptree"].append({"table": name, "column": col,
|
| 1344 |
+
"version": int(ver),
|
| 1345 |
+
"bytes": p.stat().st_size})
|
| 1346 |
+
for p in list(ixdir.glob("hnsw.v*.bin")) + list(ixdir.glob("ivfpq.v*.npz")):
|
| 1347 |
+
kind = "hnsw" if p.name.startswith("hnsw") else "ivfpq"
|
| 1348 |
+
ver = int(p.name.split(".v")[1].split(".")[0])
|
| 1349 |
+
cur = db.table("embeddings").state().version
|
| 1350 |
+
out["ann"].append({"kind": kind, "version": ver,
|
| 1351 |
+
"current": cur, "stale": ver != cur,
|
| 1352 |
+
"bytes": p.stat().st_size})
|
| 1353 |
+
# candidate numeric columns for B+ indexing
|
| 1354 |
+
out["indexable"] = {}
|
| 1355 |
+
for name in db.tables():
|
| 1356 |
+
st = db.table(name).state()
|
| 1357 |
+
if st.kind != "timeseries" or not st.files:
|
| 1358 |
+
continue
|
| 1359 |
+
import pyarrow.parquet as pq
|
| 1360 |
+
schema = pq.ParquetFile(db.dir / "tables" / name / st.files[0].path) \
|
| 1361 |
+
.schema_arrow
|
| 1362 |
+
import pyarrow as pa
|
| 1363 |
+
cols = [f.name for f in schema
|
| 1364 |
+
if f.name != "ts" and (pa.types.is_integer(f.type)
|
| 1365 |
+
or pa.types.is_floating(f.type))]
|
| 1366 |
+
if cols:
|
| 1367 |
+
out["indexable"][name] = cols
|
| 1368 |
+
out["has_embeddings"] = "embeddings" in db.tables() and \
|
| 1369 |
+
db.table("embeddings").state().rows > 0
|
| 1370 |
+
return out
|
| 1371 |
+
|
| 1372 |
+
|
| 1373 |
+
def api_build_index(key: str, body: dict):
|
| 1374 |
+
db = STORES[key]
|
| 1375 |
+
t0 = time.perf_counter()
|
| 1376 |
+
if body.get("ann"):
|
| 1377 |
+
from . import ann
|
| 1378 |
+
r = (ann.build_hnsw(db) if body["ann"] == "hnsw"
|
| 1379 |
+
else ann.build_ivfpq(db))
|
| 1380 |
+
r["ms"] = round((time.perf_counter() - t0) * 1e3)
|
| 1381 |
+
return r
|
| 1382 |
+
r = db.table(body["table"]).create_index(body["column"])
|
| 1383 |
+
r["ms"] = round((time.perf_counter() - t0) * 1e3)
|
| 1384 |
+
return r
|
| 1385 |
+
|
| 1386 |
+
|
| 1387 |
+
def api_maintenance(key: str, body: dict):
|
| 1388 |
+
db = STORES[key]
|
| 1389 |
+
op = body["op"]
|
| 1390 |
+
t0 = time.perf_counter()
|
| 1391 |
+
if op == "compact":
|
| 1392 |
+
r = db.table(body["table"]).compact()
|
| 1393 |
+
elif op == "vacuum":
|
| 1394 |
+
r = db.vacuum(retain_versions=int(body.get("retain", 3)),
|
| 1395 |
+
dry_run=bool(body.get("dry_run", False)))
|
| 1396 |
+
elif op == "delete_range":
|
| 1397 |
+
r = db.table(body["table"]).delete_range(int(body["t0"]),
|
| 1398 |
+
int(body["t1"]))
|
| 1399 |
+
else:
|
| 1400 |
+
raise ValueError(f"unknown maintenance op {op!r}")
|
| 1401 |
+
r["ms"] = round((time.perf_counter() - t0) * 1e3)
|
| 1402 |
+
return r
|
| 1403 |
+
|
| 1404 |
+
|
| 1405 |
+
def api_query(key: str, body: dict):
|
| 1406 |
+
db = STORES[key]
|
| 1407 |
+
kind = body.get("type")
|
| 1408 |
+
t_start = time.perf_counter()
|
| 1409 |
+
if kind == "sql":
|
| 1410 |
+
df = db.sql(body["sql"]).head(200)
|
| 1411 |
+
return {"columns": list(df.columns),
|
| 1412 |
+
"rows": json.loads(df.to_json(orient="values")),
|
| 1413 |
+
"ms": round((time.perf_counter() - t_start) * 1e3, 1)}
|
| 1414 |
+
if kind == "text":
|
| 1415 |
+
floor = body.get("floor")
|
| 1416 |
+
kw = {}
|
| 1417 |
+
if body.get("floor_mode") == "percentile" and floor is not None:
|
| 1418 |
+
kw["percentile"] = float(floor)
|
| 1419 |
+
elif body.get("floor_mode") == "min_score" and floor is not None:
|
| 1420 |
+
kw["min_score"] = float(floor)
|
| 1421 |
+
hits, stats = db.search(
|
| 1422 |
+
body["text"], k=int(body.get("k", 8)),
|
| 1423 |
+
nprobe=int(body.get("nprobe", 3)),
|
| 1424 |
+
method=body.get("method", "auto"),
|
| 1425 |
+
neg_weight=float(body.get("neg_weight", 0.5)),
|
| 1426 |
+
t0=body.get("t0"), t1=body.get("t1"),
|
| 1427 |
+
streams=body.get("streams") or None,
|
| 1428 |
+
rerank=bool(body.get("rerank")),
|
| 1429 |
+
rerank_top=int(body.get("rerank_top", 10)), **kw)
|
| 1430 |
+
return {"hits": hits, "stats": stats,
|
| 1431 |
+
"ms": round((time.perf_counter() - t_start) * 1e3, 1)}
|
| 1432 |
+
if kind == "context":
|
| 1433 |
+
# PRODUCT SURFACE = MEASURED SURFACE: this is the exact
|
| 1434 |
+
# search_set the ledger benchmarks (stable-benchmark
|
| 1435 |
+
# directive — the old captioned search_context served here
|
| 1436 |
+
# while acceptance was measured elsewhere; never again).
|
| 1437 |
+
from elidedb.scenario import search_set
|
| 1438 |
+
want = int(body.get("k", 10))
|
| 1439 |
+
scoped = bool(body.get("streams") or body.get("t0") is not None
|
| 1440 |
+
or body.get("t1") is not None)
|
| 1441 |
+
# scoped queries must rank a large pool BEFORE the scope
|
| 1442 |
+
# filters below — filtering the unscoped top-k starves the
|
| 1443 |
+
# result. Cap sizing differs by tier: audit_n (search_set's
|
| 1444 |
+
# "stratified sample" doc) is DEAD — never passed to
|
| 1445 |
+
# _binding_audit, which loops the SAM-3.1 tracker over every
|
| 1446 |
+
# element of `chosen` (scenario.py:390-392, 509-512), so audit
|
| 1447 |
+
# cost grows with k_max, not a fixed sample. A scoped+audited
|
| 1448 |
+
# query therefore gets a small pool; scoped+fast can afford 400.
|
| 1449 |
+
pool = (400 if not body.get("rerank") else max(4 * want, 40)) \
|
| 1450 |
+
if scoped else want
|
| 1451 |
+
# ENGINE. The Desk called search_set unconditionally, which is
|
| 1452 |
+
# the TEACHER: 59 s cold, 8.4 s for a warm NEW query, because it
|
| 1453 |
+
# loads PE + SigLIP2 + IV2 + V-JEPA + X-CLIP and runs every
|
| 1454 |
+
# channel over the corpus. The teacher is a LABELLER, not a
|
| 1455 |
+
# serving path. The student answers the same question in ~27 ms
|
| 1456 |
+
# from precomputed columns and is now the default.
|
| 1457 |
+
from .student import available as _st_ok, search_student
|
| 1458 |
+
engine = body.get("engine") or ("student" if _st_ok()
|
| 1459 |
+
else "teacher")
|
| 1460 |
+
if engine == "student" and _st_ok() and not body.get("rerank"):
|
| 1461 |
+
r = search_student(db, body["text"], k_max=pool)
|
| 1462 |
+
else:
|
| 1463 |
+
engine = "teacher"
|
| 1464 |
+
r = search_set(db, body["text"],
|
| 1465 |
+
purity="audited" if body.get("rerank") else "fast",
|
| 1466 |
+
k_max=pool)
|
| 1467 |
+
hits = [{"stream": c["stream"], "t0": c["t0"], "t1": c["t1"],
|
| 1468 |
+
"score": c["score"]} for c in r["clips"]]
|
| 1469 |
+
if body.get("streams"):
|
| 1470 |
+
hits = [h for h in hits if h["stream"] in body["streams"]]
|
| 1471 |
+
lo, hi = body.get("t0"), body.get("t1")
|
| 1472 |
+
if lo is not None:
|
| 1473 |
+
hits = [h for h in hits if h["t1"] >= int(lo)]
|
| 1474 |
+
if hi is not None:
|
| 1475 |
+
hits = [h for h in hits if h["t0"] <= int(hi)]
|
| 1476 |
+
hits = hits[:want]
|
| 1477 |
+
stats = {"channels": r.get("channels", []),
|
| 1478 |
+
"scored": r.get("scored", 0),
|
| 1479 |
+
"direction_filtered": r.get("direction_filtered", 0),
|
| 1480 |
+
"no_match": bool(r.get("no_match")),
|
| 1481 |
+
"engine": engine,
|
| 1482 |
+
"set_ms": r.get("ms")}
|
| 1483 |
+
return {"hits": hits, "stats": stats,
|
| 1484 |
+
"ms": round((time.perf_counter() - t_start) * 1e3, 1)}
|
| 1485 |
+
if kind == "predicate":
|
| 1486 |
+
from elidedb.store import QueryStats
|
| 1487 |
+
qs = QueryStats()
|
| 1488 |
+
tab = db.table(body["table"])
|
| 1489 |
+
vals = [float(body["value"])] if body.get("value2") in (None, "") \
|
| 1490 |
+
else [float(body["value"]), float(body["value2"])]
|
| 1491 |
+
out = tab.where(body["column"], body["op"], *vals, stats=qs).to_pandas()
|
| 1492 |
+
return {"columns": list(out.columns[:8]),
|
| 1493 |
+
"rows": json.loads(out.head(50).iloc[:, :8].to_json(
|
| 1494 |
+
orient="values")),
|
| 1495 |
+
"count": len(out),
|
| 1496 |
+
"stats": {"files": f"{qs.files_touched}/{qs.files_total}",
|
| 1497 |
+
"bytes_touched": qs.bytes_touched,
|
| 1498 |
+
"corpus_bytes": qs.corpus_bytes,
|
| 1499 |
+
"elided_pct": round(qs.elided_pct, 3)},
|
| 1500 |
+
"ms": round((time.perf_counter() - t_start) * 1e3, 1)}
|
| 1501 |
+
if kind == "clip":
|
| 1502 |
+
hits, stats = db.search_clip(body["stream"], int(body["t0"]),
|
| 1503 |
+
int(body["t1"]), k=int(body.get("k", 8)),
|
| 1504 |
+
method=body.get("method", "auto"))
|
| 1505 |
+
return {"hits": hits, "stats": stats,
|
| 1506 |
+
"ms": round((time.perf_counter() - t_start) * 1e3, 1)}
|
| 1507 |
+
if kind == "window":
|
| 1508 |
+
w, stats = db.window(int(body["t0"]), int(body["t1"]),
|
| 1509 |
+
tables=body.get("tables") or None)
|
| 1510 |
+
out = {"tables": {}, "ms": round(stats.wall_ms, 1),
|
| 1511 |
+
"stats": {"files": f"{stats.files_touched}/{stats.files_total}",
|
| 1512 |
+
"bytes_touched": stats.bytes_touched,
|
| 1513 |
+
"corpus_bytes": stats.corpus_bytes,
|
| 1514 |
+
"elided_pct": round(stats.elided_pct, 3),
|
| 1515 |
+
"rows": stats.rows_returned}}
|
| 1516 |
+
from elidedb.video import FrameSet
|
| 1517 |
+
for name, v in w.items():
|
| 1518 |
+
if isinstance(v, FrameSet):
|
| 1519 |
+
out["tables"][name] = {"kind": "frames", "count": len(v),
|
| 1520 |
+
"streams": v.streams()}
|
| 1521 |
+
else:
|
| 1522 |
+
df = v.to_pandas().head(6)
|
| 1523 |
+
out["tables"][name] = {
|
| 1524 |
+
"kind": "rows", "count": len(v),
|
| 1525 |
+
"columns": list(df.columns),
|
| 1526 |
+
"head": json.loads(df.to_json(orient="values"))}
|
| 1527 |
+
return out
|
| 1528 |
+
raise ValueError(f"unknown query type {kind!r}")
|
| 1529 |
+
|
| 1530 |
+
|
| 1531 |
+
class Handler(BaseHTTPRequestHandler):
|
| 1532 |
+
def log_message(self, *a): # quiet
|
| 1533 |
+
pass
|
| 1534 |
+
|
| 1535 |
+
def _send(self, code, body, ctype="application/json", cache=False):
|
| 1536 |
+
self.send_response(code)
|
| 1537 |
+
self.send_header("Content-Type", ctype)
|
| 1538 |
+
self.send_header("Content-Length", str(len(body)))
|
| 1539 |
+
if READONLY:
|
| 1540 |
+
# the public demo is pinged cross-origin by the landing
|
| 1541 |
+
# page to tell "awake" from "cold starting"
|
| 1542 |
+
self.send_header("Access-Control-Allow-Origin", "*")
|
| 1543 |
+
if cache:
|
| 1544 |
+
self.send_header("Cache-Control", "max-age=3600")
|
| 1545 |
+
self.end_headers()
|
| 1546 |
+
self.wfile.write(body)
|
| 1547 |
+
|
| 1548 |
+
def _send_media(self, body, ctype):
|
| 1549 |
+
"""Range-aware send: <video> elements (Safari especially) seek via
|
| 1550 |
+
byte-range requests; a byte-range database ought to honor them."""
|
| 1551 |
+
rng = self.headers.get("Range")
|
| 1552 |
+
total = len(body)
|
| 1553 |
+
if rng and rng.startswith("bytes="):
|
| 1554 |
+
spec = rng[6:].split("-")
|
| 1555 |
+
a = int(spec[0]) if spec[0] else 0
|
| 1556 |
+
b = int(spec[1]) if len(spec) > 1 and spec[1] else total - 1
|
| 1557 |
+
b = min(b, total - 1)
|
| 1558 |
+
chunk = body[a:b + 1]
|
| 1559 |
+
self.send_response(206)
|
| 1560 |
+
self.send_header("Content-Type", ctype)
|
| 1561 |
+
self.send_header("Content-Range", f"bytes {a}-{b}/{total}")
|
| 1562 |
+
self.send_header("Accept-Ranges", "bytes")
|
| 1563 |
+
self.send_header("Content-Length", str(len(chunk)))
|
| 1564 |
+
self.send_header("Cache-Control", "max-age=3600")
|
| 1565 |
+
self.end_headers()
|
| 1566 |
+
self.wfile.write(chunk)
|
| 1567 |
+
return
|
| 1568 |
+
self.send_response(200)
|
| 1569 |
+
self.send_header("Content-Type", ctype)
|
| 1570 |
+
self.send_header("Accept-Ranges", "bytes")
|
| 1571 |
+
self.send_header("Content-Length", str(total))
|
| 1572 |
+
self.send_header("Cache-Control", "max-age=3600")
|
| 1573 |
+
self.end_headers()
|
| 1574 |
+
self.wfile.write(body)
|
| 1575 |
+
|
| 1576 |
+
def _json(self, obj, code=200):
|
| 1577 |
+
self._send(code, json.dumps(obj).encode())
|
| 1578 |
+
|
| 1579 |
+
def do_GET(self):
|
| 1580 |
+
u = urlparse(self.path)
|
| 1581 |
+
q = {k: v[0] for k, v in parse_qs(u.query).items()}
|
| 1582 |
+
try:
|
| 1583 |
+
if u.path == "/" or u.path == "/index.html":
|
| 1584 |
+
html = (Path(__file__).parent / "desk_ui.html").read_bytes()
|
| 1585 |
+
return self._send(200, html, "text/html; charset=utf-8")
|
| 1586 |
+
if u.path == "/api/stores":
|
| 1587 |
+
discover()
|
| 1588 |
+
return self._json([store_summary(k) for k in STORES])
|
| 1589 |
+
if u.path == "/api/architecture":
|
| 1590 |
+
return self._json(api_architecture(q["store"]))
|
| 1591 |
+
if u.path == "/api/bytes":
|
| 1592 |
+
return self._json(api_bytes(q["store"]))
|
| 1593 |
+
if u.path == "/api/dbinternals":
|
| 1594 |
+
return self._json(api_dbinternals(q["store"],
|
| 1595 |
+
q.get("table")))
|
| 1596 |
+
if u.path == "/api/map":
|
| 1597 |
+
return self._json(api_map(q["store"]))
|
| 1598 |
+
if u.path == "/api/geo":
|
| 1599 |
+
return self._json(api_geo(q["store"]))
|
| 1600 |
+
if u.path == "/api/analytics":
|
| 1601 |
+
return self._json(api_analytics(q["store"]))
|
| 1602 |
+
if u.path == "/api/history":
|
| 1603 |
+
db = STORES[q["store"]]
|
| 1604 |
+
return self._json(db.table(q["table"]).history())
|
| 1605 |
+
if u.path == "/api/storage":
|
| 1606 |
+
return self._json(api_storage(q["store"], q["table"]))
|
| 1607 |
+
if u.path == "/api/indexes":
|
| 1608 |
+
return self._json(api_indexes(q["store"]))
|
| 1609 |
+
if u.path == "/api/schema":
|
| 1610 |
+
return self._json(api_schema(q["store"], q.get("table")))
|
| 1611 |
+
if u.path == "/api/thumb":
|
| 1612 |
+
jpg = api_thumb(q["store"], q.get("stream", ""),
|
| 1613 |
+
int(q["t"]), int(q.get("w", "360")))
|
| 1614 |
+
if jpg is None:
|
| 1615 |
+
return self._json({"error": "no frame"}, 404)
|
| 1616 |
+
return self._send(200, jpg, "image/jpeg", cache=True)
|
| 1617 |
+
if u.path == "/api/version":
|
| 1618 |
+
return self._json({"build": build_id(),
|
| 1619 |
+
"readonly": READONLY})
|
| 1620 |
+
if u.path == "/api/clip":
|
| 1621 |
+
mp4 = api_clip(q["store"], q.get("stream", ""),
|
| 1622 |
+
int(q["t0"]), int(q["t1"]),
|
| 1623 |
+
int(q.get("w", "640")))
|
| 1624 |
+
if isinstance(mp4, dict): # structured failure
|
| 1625 |
+
return self._json(mp4, 422)
|
| 1626 |
+
if mp4 is None:
|
| 1627 |
+
return self._json({"error": "no frames in window"}, 404)
|
| 1628 |
+
return self._send_media(mp4, "video/mp4")
|
| 1629 |
+
return self._json({"error": "not found"}, 404)
|
| 1630 |
+
except Exception as e: # surface, don't die
|
| 1631 |
+
return self._json({"error": f"{type(e).__name__}: {e}"}, 500)
|
| 1632 |
+
|
| 1633 |
+
def do_POST(self):
|
| 1634 |
+
u = urlparse(self.path)
|
| 1635 |
+
try:
|
| 1636 |
+
n = int(self.headers.get("Content-Length", 0))
|
| 1637 |
+
body = json.loads(self.rfile.read(n) or b"{}")
|
| 1638 |
+
if u.path == "/api/query":
|
| 1639 |
+
return self._json(api_query(body["store"], body))
|
| 1640 |
+
# mutating operations: refused in the public demo. Queries
|
| 1641 |
+
# stay open; the stores stay exactly as shipped.
|
| 1642 |
+
if READONLY and u.path in ("/api/build_index",
|
| 1643 |
+
"/api/maintenance"):
|
| 1644 |
+
return self._json(
|
| 1645 |
+
{"error": "this deployment is read only"}, 403)
|
| 1646 |
+
if u.path == "/api/build_index":
|
| 1647 |
+
return self._json(api_build_index(body["store"], body))
|
| 1648 |
+
if u.path == "/api/maintenance":
|
| 1649 |
+
return self._json(api_maintenance(body["store"], body))
|
| 1650 |
+
return self._json({"error": "not found"}, 404)
|
| 1651 |
+
except Exception as e:
|
| 1652 |
+
return self._json({"error": f"{type(e).__name__}: {e}"}, 500)
|
| 1653 |
+
|
| 1654 |
+
|
| 1655 |
+
def _warm():
|
| 1656 |
+
"""Pre-build the per-store matrix caches AND load the text tower so the
|
| 1657 |
+
FIRST query of a session is as fast as the hundredth. Off the request
|
| 1658 |
+
path; measured: an unwarmed first query pays ~1 s of model load."""
|
| 1659 |
+
try:
|
| 1660 |
+
from elidedb.embeddings import embed_text
|
| 1661 |
+
embed_text("warmup")
|
| 1662 |
+
except Exception:
|
| 1663 |
+
pass
|
| 1664 |
+
# the STUDENT is the serving path, so warm what it needs: the PE
|
| 1665 |
+
# text tower and its own weights. Measured cold 10.2 s (almost all
|
| 1666 |
+
# of it the text encoder), warm 25-41 ms - so this thread is the
|
| 1667 |
+
# difference between the first UI query feeling broken and feeling
|
| 1668 |
+
# instant.
|
| 1669 |
+
try:
|
| 1670 |
+
from elidedb.student import available, search_student
|
| 1671 |
+
if available():
|
| 1672 |
+
for _k, _db in list(STORES.items()):
|
| 1673 |
+
search_student(_db, "the robot opens the drawer", k_max=5)
|
| 1674 |
+
break
|
| 1675 |
+
except Exception:
|
| 1676 |
+
pass
|
| 1677 |
+
for key, db in list(STORES.items()):
|
| 1678 |
+
try:
|
| 1679 |
+
from elidedb.embeddings import _vec_table
|
| 1680 |
+
_vec_table(db, "embeddings")
|
| 1681 |
+
from elidedb.verified import _recording_spans, _verdict_map
|
| 1682 |
+
_verdict_map(db)
|
| 1683 |
+
_recording_spans(db)
|
| 1684 |
+
except Exception:
|
| 1685 |
+
pass
|
| 1686 |
+
try:
|
| 1687 |
+
_vec_table(db, "motion_vectors") # builds the mmap sidecar
|
| 1688 |
+
except Exception:
|
| 1689 |
+
pass
|
| 1690 |
+
try:
|
| 1691 |
+
from elidedb.subjects import subject_prefixes
|
| 1692 |
+
subject_prefixes(db) # mines on first run, cached
|
| 1693 |
+
except Exception:
|
| 1694 |
+
pass
|
| 1695 |
+
|
| 1696 |
+
|
| 1697 |
+
def main():
|
| 1698 |
+
global LAKE
|
| 1699 |
+
ap = argparse.ArgumentParser()
|
| 1700 |
+
ap.add_argument("--root", default=os.environ.get("DESK_ROOT",
|
| 1701 |
+
str(LAKE)))
|
| 1702 |
+
ap.add_argument("--port", type=int,
|
| 1703 |
+
default=int(os.environ.get("PORT", "8787")))
|
| 1704 |
+
ap.add_argument("--host", default=os.environ.get("DESK_HOST",
|
| 1705 |
+
"127.0.0.1"))
|
| 1706 |
+
ap.add_argument("--open", action="store_true")
|
| 1707 |
+
args = ap.parse_args()
|
| 1708 |
+
LAKE = Path(args.root).resolve()
|
| 1709 |
+
discover()
|
| 1710 |
+
threading.Thread(target=_warm, daemon=True).start()
|
| 1711 |
+
srv = ThreadingHTTPServer((args.host, args.port), Handler)
|
| 1712 |
+
print(f"ElideDB Desk: http://{args.host}:{args.port} "
|
| 1713 |
+
f"({len(STORES)} stores under {LAKE})"
|
| 1714 |
+
f"{' [read only]' if READONLY else ''}")
|
| 1715 |
+
if args.open:
|
| 1716 |
+
import subprocess
|
| 1717 |
+
threading.Timer(0.4, lambda: subprocess.run(
|
| 1718 |
+
["open", f"http://localhost:{args.port}"])).start()
|
| 1719 |
+
srv.serve_forever()
|
| 1720 |
+
|
| 1721 |
+
|
| 1722 |
+
if __name__ == "__main__":
|
| 1723 |
+
main()
|
python/elidedb/desk_ui.html
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
python/elidedb/device.py
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""One device policy for every model loader.
|
| 2 |
+
|
| 3 |
+
The store format is portable; the model loaders must be too. Every
|
| 4 |
+
loader asks this module instead of probing torch itself, so a cloud
|
| 5 |
+
CPU box, an Apple laptop, and a CUDA node run the same code with the
|
| 6 |
+
right placement:
|
| 7 |
+
|
| 8 |
+
ELIDEDB_DEVICE force a device ("cpu", "mps", "cuda"); default is
|
| 9 |
+
mps when available, else cuda, else cpu
|
| 10 |
+
ELIDEDB_DTYPE force a torch dtype name; default float16 on
|
| 11 |
+
mps/cuda and bfloat16 on cpu (halves resident
|
| 12 |
+
memory against fp32, and fp16 matmuls are not a
|
| 13 |
+
real CPU option)
|
| 14 |
+
"""
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import os
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def text_only():
|
| 21 |
+
"""Serving deployments set ELIDEDB_TEXT_ONLY=1: the query path
|
| 22 |
+
encodes TEXT only (every frame vector is precomputed at ingest),
|
| 23 |
+
so the loaders drop their vision towers after load and roughly
|
| 24 |
+
halve resident memory. Ingest machines leave it unset."""
|
| 25 |
+
return os.environ.get("ELIDEDB_TEXT_ONLY", "") == "1"
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def strip_vision(model, *attrs):
|
| 29 |
+
"""Release the named submodules when serving text-only."""
|
| 30 |
+
if not text_only():
|
| 31 |
+
return model
|
| 32 |
+
import gc
|
| 33 |
+
for a in attrs:
|
| 34 |
+
if hasattr(model, a):
|
| 35 |
+
setattr(model, a, None)
|
| 36 |
+
gc.collect()
|
| 37 |
+
return model
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def pick():
|
| 41 |
+
import torch
|
| 42 |
+
dev = os.environ.get("ELIDEDB_DEVICE", "").strip()
|
| 43 |
+
if not dev:
|
| 44 |
+
if torch.backends.mps.is_available():
|
| 45 |
+
dev = "mps"
|
| 46 |
+
elif torch.cuda.is_available():
|
| 47 |
+
dev = "cuda"
|
| 48 |
+
else:
|
| 49 |
+
dev = "cpu"
|
| 50 |
+
name = os.environ.get("ELIDEDB_DTYPE", "").strip()
|
| 51 |
+
if name:
|
| 52 |
+
dtype = getattr(torch, name)
|
| 53 |
+
else:
|
| 54 |
+
dtype = torch.float16 if dev in ("mps", "cuda") \
|
| 55 |
+
else torch.bfloat16
|
| 56 |
+
return dev, dtype
|
python/elidedb/dinov3.py
ADDED
|
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""DINOv3 — the teacher's substrate. One backbone, three granularities.
|
| 2 |
+
|
| 3 |
+
Frames (every frame -> scene series), track crops (identity descriptors)
|
| 4 |
+
and, later, whatever the FDNN student distills from. Chosen over the
|
| 5 |
+
CLIP family for everything appearance-shaped because it has NO text
|
| 6 |
+
tower: it was never trained to collapse instances into nameable
|
| 7 |
+
categories, which is precisely the failure that made a red and a green
|
| 8 |
+
pepper the same object under text-aligned encoders. The instance-level
|
| 9 |
+
numbers agree (+10.9 GAP retrieval over DINOv2, which already beat CLIP
|
| 10 |
+
at instance top-1).
|
| 11 |
+
|
| 12 |
+
Gated weights: the user accepted Meta's license on HF (2026-08-01).
|
| 13 |
+
|
| 14 |
+
No resizing cleverness: everything goes to IMG px square through the
|
| 15 |
+
processor, because identity compares crops to crops and frames to
|
| 16 |
+
frames — never one to the other — so a shared canonical size keeps
|
| 17 |
+
each comparison internally consistent.
|
| 18 |
+
"""
|
| 19 |
+
from __future__ import annotations
|
| 20 |
+
|
| 21 |
+
import os
|
| 22 |
+
|
| 23 |
+
import numpy as np
|
| 24 |
+
|
| 25 |
+
_M = {}
|
| 26 |
+
|
| 27 |
+
# ConvNeXt-Tiny, not ViT-S, and the reason is numerical, not taste:
|
| 28 |
+
# the DINOv3 ViT variants OVERFLOW FP16 on MPS - every embedding NaN,
|
| 29 |
+
# and a NaN loses every >= comparison silently, so the failure mode was
|
| 30 |
+
# not an error but "12,583 tracks, 12,583 objects, zero recurrence".
|
| 31 |
+
# ViT-S works in fp32 at ~2x its fp16 time, which makes it slower than
|
| 32 |
+
# ConvNeXt-Tiny in fp16 - and ConvNeXt measured AUC 0.9994 on the
|
| 33 |
+
# corpus's proven pairs (vs yolo26n-reid 0.9047). Faster AND stronger,
|
| 34 |
+
# so the 768-d (vs 384) table cost is accepted.
|
| 35 |
+
MID = os.environ.get("ELIDEDB_DINOV3",
|
| 36 |
+
"facebook/dinov3-convnext-tiny-pretrain-lvd1689m")
|
| 37 |
+
IMG = int(os.environ.get("ELIDEDB_DINOV3_SZ", "224"))
|
| 38 |
+
BATCH = int(os.environ.get("ELIDEDB_DINOV3_BATCH", "64"))
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def _load(mid=None):
|
| 42 |
+
mid = mid or MID
|
| 43 |
+
if _M.get("mid") == mid:
|
| 44 |
+
return _M
|
| 45 |
+
import torch
|
| 46 |
+
from transformers import AutoImageProcessor, AutoModel
|
| 47 |
+
from .device import pick
|
| 48 |
+
dev, dtype = pick()
|
| 49 |
+
if "vit" in mid.rsplit("/", 1)[-1]:
|
| 50 |
+
dtype = torch.float32 # fp16 overflow, see above
|
| 51 |
+
_M.clear()
|
| 52 |
+
_M["proc"] = AutoImageProcessor.from_pretrained(mid)
|
| 53 |
+
_M["model"] = AutoModel.from_pretrained(
|
| 54 |
+
mid, dtype=dtype, low_cpu_mem_usage=True).to(dev).eval()
|
| 55 |
+
_M["dev"], _M["dtype"], _M["mid"] = dev, dtype, mid
|
| 56 |
+
_M["torch"] = torch
|
| 57 |
+
return _M
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def embed(images, batch=None, mid=None):
|
| 61 |
+
"""(N, D) unit-norm float32 embeddings for a list of HWC uint8 arrays.
|
| 62 |
+
|
| 63 |
+
Pooled (CLS) output — the global instance/scene vector, the thing
|
| 64 |
+
retrieval compares. Patch tokens exist but are a different product
|
| 65 |
+
for a different consumer; nothing here should quietly average them.
|
| 66 |
+
"""
|
| 67 |
+
m = _load(mid)
|
| 68 |
+
torch = m["torch"]
|
| 69 |
+
out = []
|
| 70 |
+
batch = batch or BATCH
|
| 71 |
+
for i in range(0, len(images), batch):
|
| 72 |
+
chunk = [np.ascontiguousarray(x[..., :3]) for x in
|
| 73 |
+
images[i:i + batch]]
|
| 74 |
+
px = m["proc"](images=chunk, return_tensors="pt",
|
| 75 |
+
size={"height": IMG, "width": IMG})
|
| 76 |
+
px = px["pixel_values"].to(m["dev"], m["dtype"])
|
| 77 |
+
with torch.no_grad():
|
| 78 |
+
r = m["model"](pixel_values=px)
|
| 79 |
+
v = r.pooler_output.float()
|
| 80 |
+
v = v / v.norm(dim=-1, keepdim=True).clamp_min(1e-8)
|
| 81 |
+
v = v.cpu().numpy().astype(np.float32)
|
| 82 |
+
# FAIL LOUD. The ViT fp16 overflow produced NaN vectors that
|
| 83 |
+
# scored 0.0 AUC without a single exception - every comparison
|
| 84 |
+
# quietly False. Garbage must stop the run, not grade it.
|
| 85 |
+
if np.isnan(v).any():
|
| 86 |
+
raise FloatingPointError(
|
| 87 |
+
f"{_M['mid']}: NaN embeddings (dtype {_M['dtype']}) - "
|
| 88 |
+
"use fp32 for ViT variants")
|
| 89 |
+
out.append(v)
|
| 90 |
+
return (np.concatenate(out) if out
|
| 91 |
+
else np.zeros((0, 384), np.float32))
|
python/elidedb/distill.py
ADDED
|
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""CHANNEL STUDENTS: an FDNN head per teacher, over vectors already computed.
|
| 2 |
+
|
| 3 |
+
The teachers - PE, SigLIP2 so400m, InternVideo2-1B, XCLIP, V-JEPA2 -
|
| 4 |
+
were built for accuracy and they are not deployable. Measured per
|
| 5 |
+
episode at write time: pe 1.17 s, sig2 0.70, iv2 1.14, xclip 0.49,
|
| 6 |
+
before V-JEPA2 runs twice more. That is the retrieval path costing an
|
| 7 |
+
order of magnitude more than the entire rest of the write, which makes
|
| 8 |
+
the teachers part of the product instead of part of the training.
|
| 9 |
+
|
| 10 |
+
A teacher's job is to be RIGHT. A student's job is to be THERE.
|
| 11 |
+
|
| 12 |
+
THE INPUT IS ALREADY PAID FOR. The write path streams every frame
|
| 13 |
+
through FDNN-V and stores a 1152-d vector per frame, at 0.020 s per
|
| 14 |
+
episode. Every channel output is a function of the same pixels, so
|
| 15 |
+
every channel can be a small head over those vectors instead of a
|
| 16 |
+
second pass over the video. The teacher is run ONCE, offline, to make
|
| 17 |
+
training pairs; after that it never touches an ingest again.
|
| 18 |
+
|
| 19 |
+
SKELETON, not a bag of layers. Each teacher pools a SPAN of frames into
|
| 20 |
+
one vector - a window for PE and SigLIP2, a whole clip for IV2, XCLIP
|
| 21 |
+
and V-JEPA2 - so the student mirrors that shape:
|
| 22 |
+
|
| 23 |
+
attention pool over the span's frame vectors (one learned query,
|
| 24 |
+
the teacher's own
|
| 25 |
+
aggregation step)
|
| 26 |
+
pre-norm residual MLP -> the teacher's dimension
|
| 27 |
+
L2 normalise (every consumer
|
| 28 |
+
compares by cosine,
|
| 29 |
+
so the norm is not
|
| 30 |
+
part of the signal)
|
| 31 |
+
|
| 32 |
+
Mean pooling was the obvious first choice and it is wrong for the same
|
| 33 |
+
reason it is wrong in the teachers: a window's meaning is carried by a
|
| 34 |
+
few frames, and averaging buries them under the static ones.
|
| 35 |
+
|
| 36 |
+
Inference is deliberately NUMPY. The head is two small matmuls; running
|
| 37 |
+
them under a training framework costs more in dispatch than the
|
| 38 |
+
arithmetic, and the write path should not have to import one.
|
| 39 |
+
"""
|
| 40 |
+
from __future__ import annotations
|
| 41 |
+
|
| 42 |
+
import json
|
| 43 |
+
from pathlib import Path
|
| 44 |
+
|
| 45 |
+
import numpy as np
|
| 46 |
+
|
| 47 |
+
# teacher output dimensions, by channel
|
| 48 |
+
DIMS = {"pe": 1024, "sig2": 1152, "iv2": 512, "xclip": 768,
|
| 49 |
+
"vjepa": 1024, "act": 174}
|
| 50 |
+
FDNNV_D = 1152
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _l2(x, axis=-1):
|
| 54 |
+
return x / (np.linalg.norm(x, axis=axis, keepdims=True) + 1e-8)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
class Head:
|
| 58 |
+
"""A distilled channel. Loads from .npz, runs in numpy."""
|
| 59 |
+
|
| 60 |
+
def __init__(self, w: dict):
|
| 61 |
+
self.q = w["q"] # (d_in,) pooling query
|
| 62 |
+
self.W1, self.b1 = w["W1"], w["b1"]
|
| 63 |
+
self.W2, self.b2 = w["W2"], w["b2"]
|
| 64 |
+
self.Wp, self.bp = w["Wp"], w["bp"]
|
| 65 |
+
self.g, self.beta = w["g"], w["beta"]
|
| 66 |
+
|
| 67 |
+
@staticmethod
|
| 68 |
+
def load(path):
|
| 69 |
+
z = np.load(path)
|
| 70 |
+
return Head({k: z[k] for k in z.files})
|
| 71 |
+
|
| 72 |
+
def pool(self, X):
|
| 73 |
+
"""Attention-pool (n_frames, d) -> (d,) with one learned query."""
|
| 74 |
+
if len(X) == 0:
|
| 75 |
+
return np.zeros(self.q.shape[0], np.float32)
|
| 76 |
+
s = X @ self.q
|
| 77 |
+
s -= s.max()
|
| 78 |
+
a = np.exp(s)
|
| 79 |
+
a /= a.sum() + 1e-8
|
| 80 |
+
return (a[:, None] * X).sum(0)
|
| 81 |
+
|
| 82 |
+
def __call__(self, X):
|
| 83 |
+
"""(n_frames, 1152) frame vectors of one span -> teacher vector."""
|
| 84 |
+
h = self.pool(np.asarray(X, np.float32))
|
| 85 |
+
# pre-norm: the pooled vector's scale varies with span length
|
| 86 |
+
h = self.g * (h - h.mean()) / (h.std() + 1e-6) + self.beta
|
| 87 |
+
z = np.maximum(h @ self.W1 + self.b1, 0) # ReLU
|
| 88 |
+
h = h + z @ self.W2 + self.b2 # residual
|
| 89 |
+
return _l2(h @ self.Wp + self.bp)
|
| 90 |
+
|
| 91 |
+
def batch(self, spans):
|
| 92 |
+
"""Many spans at once - the write path's actual call shape."""
|
| 93 |
+
return np.stack([self(X) for X in spans]) if spans else \
|
| 94 |
+
np.zeros((0, self.Wp.shape[1]), np.float32)
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def pairs(store, channel, frames_table="frame_vectors"):
|
| 98 |
+
"""(list of frame-vector spans, teacher matrix) for one channel.
|
| 99 |
+
|
| 100 |
+
A teacher row carries ts..t1; the student sees exactly the frame
|
| 101 |
+
vectors inside that span, which is the same evidence the teacher
|
| 102 |
+
had and no more. Spans with no frames are dropped rather than
|
| 103 |
+
zero-filled - a student taught to map emptiness to a real vector
|
| 104 |
+
learns to hallucinate.
|
| 105 |
+
"""
|
| 106 |
+
import pyarrow.compute as pc
|
| 107 |
+
tv = store.table(f"{channel}_vectors").scan()
|
| 108 |
+
fv = store.table(frames_table).scan()
|
| 109 |
+
fts = np.asarray(fv.column("ts").to_pylist(), np.int64)
|
| 110 |
+
fst = np.asarray([str(s) for s in fv.column("stream").to_pylist()])
|
| 111 |
+
F = np.asarray(fv.column("vector").to_pylist(), np.float32)
|
| 112 |
+
|
| 113 |
+
order = np.argsort(fts, kind="stable")
|
| 114 |
+
fts, fst, F = fts[order], fst[order], F[order]
|
| 115 |
+
|
| 116 |
+
t0 = np.asarray(tv.column("ts").to_pylist(), np.int64)
|
| 117 |
+
t1 = np.asarray(tv.column("t1").to_pylist(), np.int64)
|
| 118 |
+
tst = np.asarray([str(s) for s in tv.column("stream").to_pylist()])
|
| 119 |
+
T = _l2(np.asarray(tv.column("vector").to_pylist(), np.float32))
|
| 120 |
+
|
| 121 |
+
X, Y, keys = [], [], []
|
| 122 |
+
for i in range(len(t0)):
|
| 123 |
+
lo = np.searchsorted(fts, t0[i], "left")
|
| 124 |
+
hi = np.searchsorted(fts, t1[i], "right")
|
| 125 |
+
if hi <= lo:
|
| 126 |
+
continue
|
| 127 |
+
sel = slice(lo, hi)
|
| 128 |
+
m = fst[sel] == tst[i]
|
| 129 |
+
if not m.any():
|
| 130 |
+
continue
|
| 131 |
+
X.append(F[sel][m])
|
| 132 |
+
Y.append(T[i])
|
| 133 |
+
keys.append((tst[i], int(t0[i])))
|
| 134 |
+
return X, (np.stack(Y) if Y else np.zeros((0, 1), np.float32)), keys
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def init(d_in=FDNNV_D, d_out=1024, hidden=512, seed=0):
|
| 138 |
+
r = np.random.default_rng(seed)
|
| 139 |
+
def n(*s):
|
| 140 |
+
return (r.normal(size=s) / np.sqrt(s[0])).astype(np.float32)
|
| 141 |
+
return {"q": n(d_in), "W1": n(d_in, hidden), "b1": np.zeros(hidden, np.float32),
|
| 142 |
+
"W2": n(hidden, d_in), "b2": np.zeros(d_in, np.float32),
|
| 143 |
+
"Wp": n(d_in, d_out), "bp": np.zeros(d_out, np.float32),
|
| 144 |
+
"g": np.ones(d_in, np.float32), "beta": np.zeros(d_in, np.float32)}
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def fidelity(head, X, Y):
|
| 148 |
+
"""Cosine to the teacher, ITS TRIVIAL BASELINE, and rank agreement.
|
| 149 |
+
|
| 150 |
+
Cosine alone is not weak, it is actively misleading here, and the
|
| 151 |
+
first PE student proved it: 0.9168 test cosine, which reads like a
|
| 152 |
+
working student until you compute what a CONSTANT prediction of the
|
| 153 |
+
corpus mean scores - 0.8629. The teacher's space is anisotropic
|
| 154 |
+
(mean pairwise cosine 0.885), so almost all of that 0.92 is the
|
| 155 |
+
shared mean and almost none of it is the episode. Nearest-neighbour
|
| 156 |
+
agreement was 0.005.
|
| 157 |
+
|
| 158 |
+
So every report carries `mean_baseline` next to `cosine`, and the
|
| 159 |
+
number that decides whether a student ships is rank agreement -
|
| 160 |
+
ranking is the only thing a retrieval channel is ever used for.
|
| 161 |
+
"""
|
| 162 |
+
P = head.batch(X)
|
| 163 |
+
cos = float(np.mean(np.sum(P * Y, 1)))
|
| 164 |
+
mu = _l2(Y.mean(0))
|
| 165 |
+
base = float(np.mean(Y @ mu))
|
| 166 |
+
n = min(len(P), 400)
|
| 167 |
+
Sp, St = P[:n] @ P[:n].T, Y[:n] @ Y[:n].T
|
| 168 |
+
np.fill_diagonal(Sp, -9); np.fill_diagonal(St, -9)
|
| 169 |
+
top1 = float(np.mean(Sp.argmax(1) == St.argmax(1)))
|
| 170 |
+
k = min(10, n - 1)
|
| 171 |
+
rp = np.argsort(-Sp, 1)[:, :k]
|
| 172 |
+
rt = np.argsort(-St, 1)[:, :k]
|
| 173 |
+
rec = float(np.mean([len(set(a) & set(b)) / k for a, b in zip(rp, rt)]))
|
| 174 |
+
return {"cosine": round(cos, 4), "mean_baseline": round(base, 4),
|
| 175 |
+
"lift_over_mean": round(cos - base, 4),
|
| 176 |
+
"nn_top1": round(top1, 4), "nn_recall@10": round(rec, 4)}
|
python/elidedb/embeddings.py
ADDED
|
@@ -0,0 +1,744 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Semantic layer over Parquet: embeddings and centroids are tables like
|
| 2 |
+
everything else — versioned by the log, readable by DuckDB, no sidecar
|
| 3 |
+
binary formats. Vectors are fixed_size_list<float32>[d] columns.
|
| 4 |
+
|
| 5 |
+
Retrieval is two-stage (learned-cell IVF): rank cluster centroids, scan the
|
| 6 |
+
top nprobe cells exactly in full space; HDBSCAN noise is always scanned so
|
| 7 |
+
the prune can cost recall nothing. All ranking is one numpy matmul — exact,
|
| 8 |
+
simple, and fast far past 100k windows."""
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import time
|
| 12 |
+
|
| 13 |
+
import numpy as np
|
| 14 |
+
import pyarrow as pa
|
| 15 |
+
import pyarrow.compute as pc
|
| 16 |
+
import pyarrow.parquet as pq
|
| 17 |
+
|
| 18 |
+
_MODEL_CACHE = {}
|
| 19 |
+
DEFAULT_MODEL = "mlx-community/siglip-so400m-patch14-384"
|
| 20 |
+
|
| 21 |
+
# Ingest cost is ALL model, not storage: measured on an M-series machine,
|
| 22 |
+
# byte-range decode runs at 2.7 ms/frame while the 384px tower runs at
|
| 23 |
+
# 90.3 ms/frame — 97% of ingest is the encoder. So the encoder is a choice,
|
| 24 |
+
# not a constant.
|
| 25 |
+
#
|
| 26 |
+
# quality siglip-so400m-patch14-384 90.3 ms/frame 1152-d (default)
|
| 27 |
+
# fast siglip-so400m-patch14-224 27.7 ms/frame 1152-d 3.3x faster
|
| 28 |
+
#
|
| 29 |
+
# `fast` is the SAME model and the same output space, fed 224px instead of
|
| 30 |
+
# 384px, so it is 256 patches per image instead of 729. Vectors from the two
|
| 31 |
+
# are NOT interchangeable — an index must be built and queried with one of
|
| 32 |
+
# them, which is why the model id is recorded in the table's metadata and the
|
| 33 |
+
# query path reads it back.
|
| 34 |
+
MODELS = {
|
| 35 |
+
"quality": "mlx-community/siglip-so400m-patch14-384",
|
| 36 |
+
"fast": "mlx-community/siglip-so400m-patch14-224",
|
| 37 |
+
}
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def resolve_model(name):
|
| 41 |
+
"""Accept a preset name ('fast'/'quality') or a raw HF model id."""
|
| 42 |
+
return MODELS.get(name, name) if name else DEFAULT_MODEL
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
# PORTABILITY: the mlx build of siglip-so400m is a straight conversion
|
| 46 |
+
# of the google checkpoint, so the SAME weights run through
|
| 47 |
+
# transformers on any machine and land in the same embedding space.
|
| 48 |
+
# Backend is auto-detected (mlx where it imports, torch elsewhere) and
|
| 49 |
+
# can be forced with ELIDEDB_TEXT_BACKEND=torch for parity testing.
|
| 50 |
+
_HF_EQUIV = {
|
| 51 |
+
"mlx-community/siglip-so400m-patch14-384":
|
| 52 |
+
"google/siglip-so400m-patch14-384",
|
| 53 |
+
"mlx-community/siglip-so400m-patch14-224":
|
| 54 |
+
"google/siglip-so400m-patch14-224",
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def _backend():
|
| 59 |
+
if "backend" not in _MODEL_CACHE:
|
| 60 |
+
import os
|
| 61 |
+
forced = os.environ.get("ELIDEDB_TEXT_BACKEND", "").strip()
|
| 62 |
+
if forced:
|
| 63 |
+
_MODEL_CACHE["backend"] = forced
|
| 64 |
+
else:
|
| 65 |
+
try:
|
| 66 |
+
import mlx_embeddings # noqa: F401
|
| 67 |
+
_MODEL_CACHE["backend"] = "mlx"
|
| 68 |
+
except ImportError:
|
| 69 |
+
_MODEL_CACHE["backend"] = "torch"
|
| 70 |
+
return _MODEL_CACHE["backend"]
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def _load_model(model_id):
|
| 74 |
+
if model_id not in _MODEL_CACHE:
|
| 75 |
+
from mlx_embeddings.utils import load
|
| 76 |
+
_MODEL_CACHE[model_id] = load(model_id)
|
| 77 |
+
return _MODEL_CACHE[model_id]
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def _load_torch(model_id):
|
| 81 |
+
key = ("torch", model_id)
|
| 82 |
+
if key not in _MODEL_CACHE:
|
| 83 |
+
from transformers import AutoModel, AutoProcessor
|
| 84 |
+
|
| 85 |
+
from .device import pick, strip_vision
|
| 86 |
+
dev, dtype = pick()
|
| 87 |
+
hf = _HF_EQUIV.get(model_id, model_id)
|
| 88 |
+
m = AutoModel.from_pretrained(
|
| 89 |
+
hf, dtype=dtype, low_cpu_mem_usage=True).to(dev).eval()
|
| 90 |
+
m = strip_vision(m, "vision_model")
|
| 91 |
+
_MODEL_CACHE[key] = (m, AutoProcessor.from_pretrained(hf), dev)
|
| 92 |
+
return _MODEL_CACHE[key]
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def _embed_images(images, model_id):
|
| 96 |
+
if _backend() == "torch":
|
| 97 |
+
import torch
|
| 98 |
+
model, processor, dev = _load_torch(resolve_model(model_id))
|
| 99 |
+
iv = processor(images=images, return_tensors="pt")
|
| 100 |
+
with torch.no_grad():
|
| 101 |
+
out = model.get_image_features(
|
| 102 |
+
pixel_values=iv["pixel_values"].to(
|
| 103 |
+
dev, model.dtype))
|
| 104 |
+
out = out.float().cpu().numpy().astype(np.float32)
|
| 105 |
+
return out / np.linalg.norm(out, axis=1, keepdims=True)
|
| 106 |
+
import mlx.core as mx
|
| 107 |
+
model, processor = _load_model(resolve_model(model_id))
|
| 108 |
+
iv = processor(images=images, return_tensors="np")
|
| 109 |
+
out = np.array(model.get_image_features(mx.array(iv["pixel_values"])),
|
| 110 |
+
dtype=np.float32)
|
| 111 |
+
return out / np.linalg.norm(out, axis=1, keepdims=True)
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def embed_text(text, model_id=DEFAULT_MODEL):
|
| 115 |
+
if _backend() == "torch":
|
| 116 |
+
import torch
|
| 117 |
+
model, processor, dev = _load_torch(resolve_model(model_id))
|
| 118 |
+
try:
|
| 119 |
+
max_len = int(
|
| 120 |
+
model.config.text_config.max_position_embeddings)
|
| 121 |
+
except AttributeError:
|
| 122 |
+
max_len = 64
|
| 123 |
+
ti = processor(text=[text], padding="max_length",
|
| 124 |
+
max_length=max_len, truncation=True,
|
| 125 |
+
return_tensors="pt")
|
| 126 |
+
with torch.no_grad():
|
| 127 |
+
v = model.get_text_features(
|
| 128 |
+
input_ids=ti["input_ids"].to(dev))
|
| 129 |
+
v = v[0].float().cpu().numpy().astype(np.float32)
|
| 130 |
+
return v / np.linalg.norm(v)
|
| 131 |
+
import mlx.core as mx
|
| 132 |
+
model, processor = _load_model(resolve_model(model_id))
|
| 133 |
+
# Each checkpoint has its own text context (so400m-384: 64 tokens,
|
| 134 |
+
# so400m-224: 16). Ask the model rather than assuming.
|
| 135 |
+
try:
|
| 136 |
+
max_len = int(model.config.text_config.max_position_embeddings)
|
| 137 |
+
except AttributeError:
|
| 138 |
+
max_len = 64
|
| 139 |
+
ti = processor(text=[text], padding="max_length", max_length=max_len,
|
| 140 |
+
truncation=True, return_tensors="np")
|
| 141 |
+
v = np.array(model.get_text_features(mx.array(ti["input_ids"])),
|
| 142 |
+
dtype=np.float32)[0]
|
| 143 |
+
return v / np.linalg.norm(v)
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
_MAT_CACHE: dict = {}
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def _vec_table(store, name="embeddings", version=None, column="vector"):
|
| 150 |
+
"""Vector table + MEMORY-MAPPED matrix, cached per log version.
|
| 151 |
+
|
| 152 |
+
Two generations of this function materialized the matrix in RAM. At
|
| 153 |
+
pilot scale that broke: bridge-full's 180k x 1152 embeddings are 0.83 GB
|
| 154 |
+
of data but cost +4.3 GB peak RSS to load (parquet decode + Arrow
|
| 155 |
+
chunks + combine copy + the cache holding table AND matrix), and the
|
| 156 |
+
desk warms every store — 7 GB before the first query. The fix is the
|
| 157 |
+
store's own law applied to vectors: mmap for reads. The matrix is
|
| 158 |
+
materialized ONCE per (table, version) into a raw .npy sidecar, then
|
| 159 |
+
every process maps it — RSS is only the pages a query touches, startup
|
| 160 |
+
costs a file open, and the OS page cache decides residency.
|
| 161 |
+
|
| 162 |
+
Row alignment: the sidecar is written from the same scan() that serves
|
| 163 |
+
the meta columns; scan's ts sort is stable over a deterministic file
|
| 164 |
+
order, so a later projected scan yields the identical permutation. A
|
| 165 |
+
length mismatch (e.g. sidecar from a dead version) forces a rebuild.
|
| 166 |
+
The parquet remains the source of truth — a sidecar is disposable.
|
| 167 |
+
"""
|
| 168 |
+
import os
|
| 169 |
+
import uuid as _uuid
|
| 170 |
+
ver = store.table(name).state().version if version is None else version
|
| 171 |
+
key = (str(store.dir), name, column, ver)
|
| 172 |
+
if key in _MAT_CACHE:
|
| 173 |
+
return _MAT_CACHE[key]
|
| 174 |
+
|
| 175 |
+
tab = store.table(name)
|
| 176 |
+
cache_dir = tab.dir / "_cache"
|
| 177 |
+
npy = cache_dir / f"{column}-v{ver}.npy"
|
| 178 |
+
|
| 179 |
+
t = vecs = None
|
| 180 |
+
if npy.exists():
|
| 181 |
+
st = tab.state(version)
|
| 182 |
+
if st.files:
|
| 183 |
+
names = pq.ParquetFile(
|
| 184 |
+
tab.dir / st.files[0].path).schema_arrow.names
|
| 185 |
+
meta_cols = [c for c in names if c != column]
|
| 186 |
+
t = tab.scan(version=version, columns=meta_cols)
|
| 187 |
+
vecs = np.load(npy, mmap_mode="r")
|
| 188 |
+
if len(vecs) != len(t):
|
| 189 |
+
t = vecs = None # stale sidecar: rebuild
|
| 190 |
+
|
| 191 |
+
if vecs is None:
|
| 192 |
+
t_full = tab.scan(version=version)
|
| 193 |
+
if len(t_full) == 0:
|
| 194 |
+
extra = ""
|
| 195 |
+
try:
|
| 196 |
+
if store.table("frame_vectors").state().files:
|
| 197 |
+
extra = (" Per-frame vectors already exist, so this "
|
| 198 |
+
"costs a numpy mean, not a GPU pass.")
|
| 199 |
+
except Exception:
|
| 200 |
+
pass
|
| 201 |
+
raise RuntimeError(
|
| 202 |
+
f"store '{store.name}' has no '{name}' table — run "
|
| 203 |
+
f"store.embed_windows() first.{extra}")
|
| 204 |
+
col = t_full.column(column)
|
| 205 |
+
if isinstance(col, pa.ChunkedArray):
|
| 206 |
+
col = col.combine_chunks()
|
| 207 |
+
try: # FixedSizeList: flat buffer reshape
|
| 208 |
+
mat = col.values.to_numpy(zero_copy_only=False) \
|
| 209 |
+
.astype(np.float32, copy=False).reshape(len(t_full), -1)
|
| 210 |
+
except Exception: # any other layout: the slow road
|
| 211 |
+
mat = np.stack([np.asarray(v, dtype=np.float32)
|
| 212 |
+
for v in col.to_pylist()])
|
| 213 |
+
cache_dir.mkdir(parents=True, exist_ok=True)
|
| 214 |
+
tmp = cache_dir / f".{_uuid.uuid4().hex[:8]}.npy"
|
| 215 |
+
np.save(tmp, np.ascontiguousarray(mat))
|
| 216 |
+
os.replace(tmp, npy) # atomic: readers see whole files
|
| 217 |
+
t = t_full.drop_columns([column]) # meta only — no double storage
|
| 218 |
+
del t_full, mat, col
|
| 219 |
+
vecs = np.load(npy, mmap_mode="r")
|
| 220 |
+
|
| 221 |
+
if len(_MAT_CACHE) > 8:
|
| 222 |
+
_MAT_CACHE.clear()
|
| 223 |
+
_MAT_CACHE[key] = (t, vecs)
|
| 224 |
+
return t, vecs
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def pool_windows(store, window_s=2.0, stride_s=None, table="frame_vectors"):
|
| 228 |
+
"""Build the `embeddings` table by POOLING existing per-frame vectors.
|
| 229 |
+
|
| 230 |
+
A window embedding is the mean of its frame embeddings. If `frame_vectors`
|
| 231 |
+
already exists there is nothing to compute with a model: decoding every
|
| 232 |
+
frame again and re-running SigLIP to reach the same answer is pure waste —
|
| 233 |
+
on the Bridge store that was 25 minutes of GPU to reproduce a number a
|
| 234 |
+
numpy mean gives in under a second.
|
| 235 |
+
|
| 236 |
+
This is the ordinary database move: two indexes over one scan, not two
|
| 237 |
+
scans.
|
| 238 |
+
"""
|
| 239 |
+
fv = store.table(table).scan()
|
| 240 |
+
if len(fv) == 0:
|
| 241 |
+
raise RuntimeError(f"'{table}' is empty — run embed_frames() first")
|
| 242 |
+
win = int(window_s * 1e9)
|
| 243 |
+
stride = int((stride_s or window_s) * 1e9)
|
| 244 |
+
t_start = time.time()
|
| 245 |
+
rows = {"ts": [], "t1": [], "stream": [], "vector": []}
|
| 246 |
+
for s in sorted(set(fv.column("stream").to_pylist())):
|
| 247 |
+
sub = fv.filter(pc.equal(fv.column("stream"), s))
|
| 248 |
+
ts = sub.column("ts").to_numpy()
|
| 249 |
+
order = np.argsort(ts)
|
| 250 |
+
ts = ts[order]
|
| 251 |
+
# zero-copy reshape, NOT to_pylist(): at 1.6M frames the Python-list
|
| 252 |
+
# road needs ~50 GB; the FixedSizeList buffer is already the matrix
|
| 253 |
+
col = sub.column("vector")
|
| 254 |
+
if isinstance(col, pa.ChunkedArray):
|
| 255 |
+
col = col.combine_chunks()
|
| 256 |
+
try:
|
| 257 |
+
vecs = col.values.to_numpy(zero_copy_only=False) \
|
| 258 |
+
.astype(np.float32, copy=False).reshape(len(sub), -1)[order]
|
| 259 |
+
except Exception:
|
| 260 |
+
vecs = np.asarray(col.to_pylist(), dtype=np.float32)[order]
|
| 261 |
+
t = int(ts[0])
|
| 262 |
+
while t <= int(ts[-1]):
|
| 263 |
+
lo, hi = np.searchsorted(ts, [t, t + win])
|
| 264 |
+
if hi > lo:
|
| 265 |
+
v = vecs[lo:hi].mean(axis=0)
|
| 266 |
+
v /= np.linalg.norm(v) + 1e-8
|
| 267 |
+
rows["ts"].append(t)
|
| 268 |
+
rows["t1"].append(min(t + win - 1, int(ts[-1])))
|
| 269 |
+
rows["stream"].append(s)
|
| 270 |
+
rows["vector"].append(v)
|
| 271 |
+
t += stride
|
| 272 |
+
dim = len(rows["vector"][0])
|
| 273 |
+
# FixedSizeListArray straight from the flat float32 buffer. The
|
| 274 |
+
# tolist() road materialises n*dim PYTHON floats — at 1.8M frames /
|
| 275 |
+
# 180k windows that was tens of GB and the process died by jetsam
|
| 276 |
+
# (exit 137) on the very last stage of a 100 h load.
|
| 277 |
+
flat = np.ascontiguousarray(
|
| 278 |
+
np.stack(rows["vector"]).astype(np.float32)).reshape(-1)
|
| 279 |
+
vec_arr = pa.FixedSizeListArray.from_arrays(pa.array(flat), dim)
|
| 280 |
+
tbl = pa.table({
|
| 281 |
+
"ts": pa.array(rows["ts"], pa.int64()),
|
| 282 |
+
"t1": pa.array(rows["t1"], pa.int64()),
|
| 283 |
+
"stream": pa.array(rows["stream"]),
|
| 284 |
+
"vector": vec_arr,
|
| 285 |
+
})
|
| 286 |
+
st = store.table("embeddings").state()
|
| 287 |
+
# `model` must be the id of the model that defines the SPACE — the query
|
| 288 |
+
# path loads it as the text tower. Student-produced vectors live in the
|
| 289 |
+
# TEACHER's space, so when the source was written by an engine (model
|
| 290 |
+
# "fdnnv"), the space id is its `teacher` field. Writing the engine name
|
| 291 |
+
# here sent "fdnnv" to the HF loader as a repo id.
|
| 292 |
+
src = store.table(table).state().meta or {}
|
| 293 |
+
src_model = src.get("model", DEFAULT_MODEL)
|
| 294 |
+
if src_model in (None, "fdnnv"):
|
| 295 |
+
src_model = src.get("teacher", DEFAULT_MODEL)
|
| 296 |
+
meta = {"model": src_model, "built_by": "pooled from frame_vectors",
|
| 297 |
+
"dim": dim, "window_s": window_s, "source_table": table,
|
| 298 |
+
"seconds": round(time.time() - t_start, 2)}
|
| 299 |
+
if st.files:
|
| 300 |
+
import uuid as _uuid
|
| 301 |
+
|
| 302 |
+
from .log import FileEntry
|
| 303 |
+
from .store import write_parquet
|
| 304 |
+
fn = f"part-{_uuid.uuid4().hex[:12]}.parquet"
|
| 305 |
+
p = store.dir / "tables" / "embeddings" / fn
|
| 306 |
+
write_parquet(tbl, p)
|
| 307 |
+
tsv = tbl.column("ts").to_numpy()
|
| 308 |
+
version = store.table("embeddings").log.commit(
|
| 309 |
+
op="replace", kind="embeddings", schema=str(tbl.schema),
|
| 310 |
+
add=[FileEntry(fn, len(tbl), p.stat().st_size,
|
| 311 |
+
int(tsv.min()), int(tsv.max()))],
|
| 312 |
+
remove=[f.path for f in st.files], meta=meta)
|
| 313 |
+
else:
|
| 314 |
+
version = store.table("embeddings").append(tbl, kind="embeddings",
|
| 315 |
+
meta=meta)
|
| 316 |
+
return {"windows": len(tbl), "dim": dim, "version": version,
|
| 317 |
+
"seconds": meta["seconds"], "source": table}
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
def embed_windows(store, frame_table="frames", window_s=2.0,
|
| 321 |
+
frames_per_window=2, model=None, batch=16, stride_s=None,
|
| 322 |
+
incremental=True, reuse_frame_vectors=True):
|
| 323 |
+
"""Tumbling windows over every video stream → mean-pooled SigLIP vectors
|
| 324 |
+
→ one commit to the `embeddings` table. Frames come through the same
|
| 325 |
+
byte-range path queries use.
|
| 326 |
+
|
| 327 |
+
If per-frame vectors already exist, they are pooled instead of re-running
|
| 328 |
+
the model (see `pool_windows`) — same result, no GPU.
|
| 329 |
+
"""
|
| 330 |
+
if reuse_frame_vectors:
|
| 331 |
+
try:
|
| 332 |
+
if store.table("frame_vectors").state().files:
|
| 333 |
+
return pool_windows(store, window_s, stride_s)
|
| 334 |
+
except Exception:
|
| 335 |
+
pass
|
| 336 |
+
from PIL import Image # noqa: F401 (decode happens in FrameSet)
|
| 337 |
+
model = model or DEFAULT_MODEL
|
| 338 |
+
tab = store.table(frame_table)
|
| 339 |
+
st = tab.state()
|
| 340 |
+
win_ns = int(window_s * 1e9)
|
| 341 |
+
stride_ns = int((stride_s or window_s) * 1e9)
|
| 342 |
+
frames = tab.scan()
|
| 343 |
+
streams = sorted(set(frames.column("stream").to_pylist()))
|
| 344 |
+
# Incremental: only embed windows past what the embeddings table already
|
| 345 |
+
# covers per stream — adding a new day of footage costs a new day of
|
| 346 |
+
# embedding, not a re-run of history.
|
| 347 |
+
done_until = {}
|
| 348 |
+
if incremental:
|
| 349 |
+
try:
|
| 350 |
+
prev = store.table("embeddings").scan()
|
| 351 |
+
if len(prev):
|
| 352 |
+
s_arr = prev.column("stream").to_pylist()
|
| 353 |
+
t1_arr = prev.column("t1").to_pylist()
|
| 354 |
+
for s_, e_ in zip(s_arr, t1_arr):
|
| 355 |
+
done_until[s_] = max(done_until.get(s_, 0), e_)
|
| 356 |
+
except Exception:
|
| 357 |
+
pass
|
| 358 |
+
jobs = [] # (stream, t0, t1)
|
| 359 |
+
for s in streams:
|
| 360 |
+
rows = frames.filter(pc.equal(frames.column("stream"), s))
|
| 361 |
+
ts = rows.column("ts").to_numpy()
|
| 362 |
+
t = (int(ts[0]) // win_ns) * win_ns
|
| 363 |
+
while t <= ts[-1]:
|
| 364 |
+
lo, hi = np.searchsorted(ts, [t, t + win_ns])
|
| 365 |
+
if hi > lo and t >= done_until.get(s, -1):
|
| 366 |
+
jobs.append((s, max(t, int(ts[0])),
|
| 367 |
+
min(t + win_ns - 1, int(ts[-1]))))
|
| 368 |
+
t += stride_ns
|
| 369 |
+
from .video import FrameSet
|
| 370 |
+
t_start = time.time()
|
| 371 |
+
if not jobs:
|
| 372 |
+
return {"windows": 0, "dim": None, "version": None, "seconds": 0.0,
|
| 373 |
+
"note": "nothing new to embed (incremental)"}
|
| 374 |
+
recs = {"ts": [], "t1": [], "stream": [], "vector": []}
|
| 375 |
+
imgs, owners = [], []
|
| 376 |
+
|
| 377 |
+
def flush():
|
| 378 |
+
nonlocal imgs, owners
|
| 379 |
+
if not imgs:
|
| 380 |
+
return
|
| 381 |
+
vecs = _embed_images(imgs, model)
|
| 382 |
+
for (key, v) in zip(owners, vecs):
|
| 383 |
+
pooled.setdefault(key, []).append(v)
|
| 384 |
+
imgs, owners = [], []
|
| 385 |
+
|
| 386 |
+
pooled = {}
|
| 387 |
+
for (s, t0, t1) in jobs:
|
| 388 |
+
fs = FrameSet(store, frame_table,
|
| 389 |
+
frames.filter(pc.and_(
|
| 390 |
+
pc.equal(frames.column("stream"), s),
|
| 391 |
+
pc.and_(pc.greater_equal(frames.column("ts"), t0),
|
| 392 |
+
pc.less_equal(frames.column("ts"), t1)))))
|
| 393 |
+
n = len(fs)
|
| 394 |
+
picks = np.linspace(0, n - 1, min(frames_per_window, n)).round().astype(int)
|
| 395 |
+
decoded = fs.decode(width=512)
|
| 396 |
+
for p in picks:
|
| 397 |
+
if p < len(decoded):
|
| 398 |
+
from PIL import Image as PILImage
|
| 399 |
+
imgs.append(PILImage.fromarray(decoded[p][1]))
|
| 400 |
+
owners.append((s, t0, t1))
|
| 401 |
+
if len(imgs) >= batch:
|
| 402 |
+
flush()
|
| 403 |
+
flush()
|
| 404 |
+
for (s, t0, t1), vs in pooled.items():
|
| 405 |
+
v = np.mean(vs, axis=0)
|
| 406 |
+
v /= np.linalg.norm(v)
|
| 407 |
+
recs["stream"].append(s)
|
| 408 |
+
recs["ts"].append(t0)
|
| 409 |
+
recs["t1"].append(t1)
|
| 410 |
+
recs["vector"].append(v)
|
| 411 |
+
dim = len(recs["vector"][0])
|
| 412 |
+
t = pa.table({
|
| 413 |
+
"ts": pa.array(recs["ts"], pa.int64()),
|
| 414 |
+
"t1": pa.array(recs["t1"], pa.int64()),
|
| 415 |
+
"stream": pa.array(recs["stream"]),
|
| 416 |
+
"vector": pa.array([v.tolist() for v in recs["vector"]],
|
| 417 |
+
pa.list_(pa.float32(), dim)),
|
| 418 |
+
})
|
| 419 |
+
version = store.table("embeddings").append(
|
| 420 |
+
t, kind="embeddings",
|
| 421 |
+
meta={"model": model, "dim": dim, "window_s": window_s,
|
| 422 |
+
"source_table": frame_table,
|
| 423 |
+
"embedded_in_s": round(time.time() - t_start, 1)})
|
| 424 |
+
return {"windows": len(t), "dim": dim, "version": version,
|
| 425 |
+
"seconds": round(time.time() - t_start, 1)}
|
| 426 |
+
|
| 427 |
+
|
| 428 |
+
def cluster(store, pca_dims=50, min_cluster_size=8):
|
| 429 |
+
"""PCA → HDBSCAN over the embeddings table → cluster ids written back as
|
| 430 |
+
a new embeddings version + a `centroids` table (full-space, normalized:
|
| 431 |
+
the coarse stage must rank in the space the fine stage scores in)."""
|
| 432 |
+
t, vecs = _vec_table(store)
|
| 433 |
+
from sklearn.decomposition import PCA
|
| 434 |
+
import hdbscan
|
| 435 |
+
red = PCA(n_components=min(pca_dims, len(vecs), vecs.shape[1]),
|
| 436 |
+
random_state=0).fit_transform(vecs)
|
| 437 |
+
labels = hdbscan.HDBSCAN(min_cluster_size=min_cluster_size).fit_predict(red)
|
| 438 |
+
out = t.drop_columns(["cluster"]) if "cluster" in t.column_names else t
|
| 439 |
+
out = out.append_column("cluster", pa.array(labels.astype("int32")))
|
| 440 |
+
# replace = remove old files + add the re-clustered ones, one commit
|
| 441 |
+
st = store.table("embeddings").state()
|
| 442 |
+
log = store.table("embeddings").log
|
| 443 |
+
import pyarrow.parquet as pq
|
| 444 |
+
import uuid as _uuid
|
| 445 |
+
fname = f"part-{_uuid.uuid4().hex[:12]}.parquet"
|
| 446 |
+
from .store import write_parquet
|
| 447 |
+
write_parquet(out, store.dir / "tables" / "embeddings" / fname)
|
| 448 |
+
from .log import FileEntry
|
| 449 |
+
p = store.dir / "tables" / "embeddings" / fname
|
| 450 |
+
tsv = out.column("ts").to_numpy()
|
| 451 |
+
log.commit(op="recluster", kind="embeddings", schema=str(out.schema),
|
| 452 |
+
add=[FileEntry(fname, len(out), p.stat().st_size,
|
| 453 |
+
int(tsv.min()), int(tsv.max()))],
|
| 454 |
+
remove=[f.path for f in st.files],
|
| 455 |
+
meta={**st.meta, "clusters": int(labels.max() + 1),
|
| 456 |
+
"noise": int((labels < 0).sum())})
|
| 457 |
+
cents = []
|
| 458 |
+
for c in range(labels.max() + 1):
|
| 459 |
+
m = vecs[labels == c].mean(axis=0)
|
| 460 |
+
cents.append(m / np.linalg.norm(m))
|
| 461 |
+
if cents:
|
| 462 |
+
ct = pa.table({
|
| 463 |
+
"ts": pa.array([0] * len(cents), pa.int64()),
|
| 464 |
+
"cluster": pa.array(range(len(cents)), pa.int32()),
|
| 465 |
+
"vector": pa.array([c.tolist() for c in cents],
|
| 466 |
+
pa.list_(pa.float32(), vecs.shape[1])),
|
| 467 |
+
})
|
| 468 |
+
cst = store.table("centroids").state()
|
| 469 |
+
store.table("centroids").log.commit(
|
| 470 |
+
op="replace", kind="centroids", schema=str(ct.schema),
|
| 471 |
+
add=[], remove=[f.path for f in cst.files])
|
| 472 |
+
store.table("centroids").append(ct, kind="centroids")
|
| 473 |
+
return {"clusters": int(labels.max() + 1),
|
| 474 |
+
"noise": int((labels < 0).sum()), "windows": len(out)}
|
| 475 |
+
|
| 476 |
+
|
| 477 |
+
def _score_windows(vecs, idx, pos_vecs, neg_vecs, neg_weight):
|
| 478 |
+
"""Compositional scoring over a candidate set.
|
| 479 |
+
|
| 480 |
+
- ONE positive term → plain cosine (classic semantic search).
|
| 481 |
+
- MANY positive terms → the window's score is the WORST of its per-term
|
| 482 |
+
cosines (min-pool). This is the compositional AND: 'two people' AND
|
| 483 |
+
'a laptop' means a clip of two people with NO laptop scores low on the
|
| 484 |
+
laptop term and is therefore rejected — the fix for 'it returns every
|
| 485 |
+
clip with two people'.
|
| 486 |
+
- negative terms → each subtracts its cosine (weighted), so
|
| 487 |
+
'... NOT a phone' pushes phone-heavy frames down.
|
| 488 |
+
"""
|
| 489 |
+
cand = vecs[idx] # [m, d]
|
| 490 |
+
pos = cand @ pos_vecs.T # [m, n_pos]
|
| 491 |
+
score = pos.min(axis=1) # min-pool = AND
|
| 492 |
+
if neg_vecs is not None and len(neg_vecs):
|
| 493 |
+
score = score - neg_weight * (cand @ neg_vecs.T).max(axis=1)
|
| 494 |
+
return score
|
| 495 |
+
|
| 496 |
+
|
| 497 |
+
def _rank(store, q, k, nprobe, merge=True, t0=None, t1=None, streams=None,
|
| 498 |
+
method="auto", pos_vecs=None, neg_vecs=None, neg_weight=0.5,
|
| 499 |
+
min_score=None, percentile=None, table="embeddings", ctx=None,
|
| 500 |
+
column="vector"):
|
| 501 |
+
t, vecs = _vec_table(store, table, column=column)
|
| 502 |
+
if table != "embeddings":
|
| 503 |
+
# The ANN artifacts (HNSW graph, IVF-PQ codes, HDBSCAN centroids) are
|
| 504 |
+
# built over `embeddings` and index THOSE row ids. Reusing them here
|
| 505 |
+
# would return neighbours of the wrong table — silently, with
|
| 506 |
+
# plausible-looking scores. Any other table scans exactly.
|
| 507 |
+
method = "exact"
|
| 508 |
+
all_t0 = t.column("ts").to_numpy()
|
| 509 |
+
all_t1 = t.column("t1").to_numpy()
|
| 510 |
+
all_s = t.column("stream").to_numpy(zero_copy_only=False)
|
| 511 |
+
# `q` (the coarse retrieval direction) is the mean of positive terms;
|
| 512 |
+
# `pos_vecs` carries the individual terms for compositional scoring.
|
| 513 |
+
if pos_vecs is None:
|
| 514 |
+
pos_vecs = q[None, :]
|
| 515 |
+
|
| 516 |
+
# ---- hybrid retrieval: predicates pushed INTO candidate selection ------
|
| 517 |
+
# Time and stream are first-class dimensions of this database; vector
|
| 518 |
+
# search composes with them instead of post-filtering a global top-k
|
| 519 |
+
# (which silently starves filtered queries of results).
|
| 520 |
+
pred = np.ones(len(vecs), bool)
|
| 521 |
+
if t0 is not None:
|
| 522 |
+
pred &= all_t1 >= t0
|
| 523 |
+
if t1 is not None:
|
| 524 |
+
pred &= all_t0 <= t1
|
| 525 |
+
if streams:
|
| 526 |
+
pred &= np.isin(all_s, list(streams))
|
| 527 |
+
|
| 528 |
+
labels = (t.column("cluster").to_numpy()
|
| 529 |
+
if "cluster" in t.column_names else None)
|
| 530 |
+
probed = total_clusters = 0
|
| 531 |
+
used = "exact"
|
| 532 |
+
|
| 533 |
+
idx = scores = None
|
| 534 |
+
if method in ("auto", "hnsw"):
|
| 535 |
+
from . import ann
|
| 536 |
+
hx = ann.load_hnsw(store)
|
| 537 |
+
if hx is not None:
|
| 538 |
+
# overfetch beyond k so predicate filtering and segment merging
|
| 539 |
+
# still see the event's neighborhood, then score exactly
|
| 540 |
+
fetch = int(min(len(vecs), max(k * 8, 64)))
|
| 541 |
+
hx.set_ef(max(fetch, 64))
|
| 542 |
+
cand, _ = hx.knn_query(q, k=fetch)
|
| 543 |
+
cand = cand[0]
|
| 544 |
+
cand = cand[pred[cand]]
|
| 545 |
+
if len(cand) >= min(k, pred.sum()):
|
| 546 |
+
idx = np.asarray(cand)
|
| 547 |
+
scores = vecs[idx] @ q
|
| 548 |
+
used = "hnsw"
|
| 549 |
+
if idx is None and method in ("auto", "ivfpq"):
|
| 550 |
+
from . import ann
|
| 551 |
+
r = ann.search_ivfpq(store, q, k=max(k * 4, 32), nprobe=max(nprobe, 8),
|
| 552 |
+
mask=pred) if method == "ivfpq" else None
|
| 553 |
+
if r is not None and r[0]:
|
| 554 |
+
idx = np.array([i for i, _ in r[0]])
|
| 555 |
+
scores = np.array([s for _, s in r[0]])
|
| 556 |
+
used = "ivfpq"
|
| 557 |
+
if idx is None:
|
| 558 |
+
mask = pred.copy()
|
| 559 |
+
if labels is not None and nprobe > 0:
|
| 560 |
+
try:
|
| 561 |
+
_, cents = _vec_table(store, "centroids")
|
| 562 |
+
total_clusters = len(cents)
|
| 563 |
+
order = np.argsort(cents @ q)[::-1][:nprobe]
|
| 564 |
+
probed = len(order)
|
| 565 |
+
mask &= np.isin(labels, order) | (labels < 0) # noise stays
|
| 566 |
+
used = "ivf"
|
| 567 |
+
except RuntimeError:
|
| 568 |
+
pass
|
| 569 |
+
idx = np.where(mask)[0]
|
| 570 |
+
scores = None
|
| 571 |
+
scanned = len(idx)
|
| 572 |
+
# Final score is ALWAYS the compositional/exact function over the
|
| 573 |
+
# candidate set (the coarse tier only shortlists; it never answers).
|
| 574 |
+
scores = _score_windows(vecs, idx, pos_vecs, neg_vecs, neg_weight)
|
| 575 |
+
|
| 576 |
+
# ---- optional fusion with the context index ----------------------------
|
| 577 |
+
# Appearance cosines and context cosines live on different scales (SigLIP
|
| 578 |
+
# image-text similarity is squashed by the modality gap into ~0.01-0.15,
|
| 579 |
+
# while context vectors are mean-free and spread over most of [-1,1]).
|
| 580 |
+
# A raw weighted sum would therefore be governed entirely by the context
|
| 581 |
+
# term regardless of alpha. Standardising each over the CANDIDATE SET
|
| 582 |
+
# first makes alpha mean what it says.
|
| 583 |
+
if ctx is not None and len(idx):
|
| 584 |
+
def _z(a):
|
| 585 |
+
return (a - a.mean()) / (a.std() + 1e-8)
|
| 586 |
+
a = float(ctx["alpha"])
|
| 587 |
+
scores = (1.0 - a) * _z(scores) + a * _z(ctx["vecs"][idx] @ ctx["q"])
|
| 588 |
+
streams_sel = all_s[idx]
|
| 589 |
+
w_t0 = all_t0[idx]
|
| 590 |
+
w_t1 = all_t1[idx]
|
| 591 |
+
|
| 592 |
+
# ---- precision floor: an ABSOLUTE cut the user controls ----------------
|
| 593 |
+
# A percentile keeps only the strongest fraction; min_score is a hard
|
| 594 |
+
# cosine floor. Either turns "top-k of everything" into "only real hits",
|
| 595 |
+
# so a query with 6 true matches returns 6, not 50.
|
| 596 |
+
keep = np.ones(len(idx), bool)
|
| 597 |
+
if percentile is not None and len(scores):
|
| 598 |
+
keep &= scores >= np.percentile(scores, percentile)
|
| 599 |
+
if min_score is not None:
|
| 600 |
+
keep &= scores >= min_score
|
| 601 |
+
if not keep.all():
|
| 602 |
+
idx, scores = idx[keep], scores[keep]
|
| 603 |
+
streams_sel, w_t0, w_t1 = streams_sel[keep], w_t0[keep], w_t1[keep]
|
| 604 |
+
|
| 605 |
+
stats = {"scanned": scanned, "total": len(vecs), "method": used,
|
| 606 |
+
"clusters_probed": probed, "clusters_total": total_clusters,
|
| 607 |
+
"predicate_candidates": int(pred.sum()),
|
| 608 |
+
"after_floor": int(len(idx))}
|
| 609 |
+
|
| 610 |
+
if not merge:
|
| 611 |
+
order = np.argsort(scores)[::-1][:k]
|
| 612 |
+
hits = [{"stream": str(streams_sel[i]), "t0": int(w_t0[i]),
|
| 613 |
+
"t1": int(w_t1[i]), "score": float(scores[i]),
|
| 614 |
+
"windows": 1} for i in order]
|
| 615 |
+
return hits, stats
|
| 616 |
+
if len(idx) == 0:
|
| 617 |
+
stats["qualifying_windows"] = 0
|
| 618 |
+
stats["segments"] = 0
|
| 619 |
+
return [], stats
|
| 620 |
+
|
| 621 |
+
# ---- dynamic segments: merge, don't chunk -------------------------------
|
| 622 |
+
# Fixed embedding windows are an INDEXING granularity, not an answer
|
| 623 |
+
# granularity. A result is the maximal run of consecutive qualifying
|
| 624 |
+
# windows on one stream: a 20 s event comes back as ONE 20 s hit (its
|
| 625 |
+
# sub-windows are never returned separately), while a query that only
|
| 626 |
+
# matches 2 s of it comes back as that tight 2 s. "Qualifying" is decided
|
| 627 |
+
# per query from the score distribution — an absolute cutoff cannot work
|
| 628 |
+
# because SigLIP cosines live on different scales per query.
|
| 629 |
+
med = float(np.median(scores))
|
| 630 |
+
top = float(scores.max())
|
| 631 |
+
thr = med + 0.55 * (top - med)
|
| 632 |
+
stats["threshold"] = round(thr, 4)
|
| 633 |
+
qual = np.where(scores >= thr)[0]
|
| 634 |
+
order = np.lexsort((w_t0[qual], streams_sel[qual]))
|
| 635 |
+
qual = qual[order]
|
| 636 |
+
|
| 637 |
+
gap_ns = int(np.median(w_t1[qual] - w_t0[qual])) + 1 if len(qual) else 0
|
| 638 |
+
segs = []
|
| 639 |
+
for i in qual:
|
| 640 |
+
s, a, b, sc = (str(streams_sel[i]), int(w_t0[i]), int(w_t1[i]),
|
| 641 |
+
float(scores[i]))
|
| 642 |
+
last = segs[-1] if segs else None
|
| 643 |
+
if last and last["stream"] == s and a - last["t1"] <= gap_ns:
|
| 644 |
+
last["t1"] = max(last["t1"], b)
|
| 645 |
+
last["score"] = max(last["score"], sc) # peak represents the segment
|
| 646 |
+
last["mean"] = (last["mean"] * last["windows"] + sc) / (last["windows"] + 1)
|
| 647 |
+
last["windows"] += 1
|
| 648 |
+
else:
|
| 649 |
+
segs.append({"stream": s, "t0": a, "t1": b, "score": sc,
|
| 650 |
+
"mean": sc, "windows": 1})
|
| 651 |
+
segs.sort(key=lambda g: -g["score"])
|
| 652 |
+
stats["qualifying_windows"] = len(qual)
|
| 653 |
+
stats["segments"] = len(segs)
|
| 654 |
+
return segs[:k], stats
|
| 655 |
+
|
| 656 |
+
|
| 657 |
+
def _parse_query(text):
|
| 658 |
+
"""Parse a compositional query string into (positive terms, negatives).
|
| 659 |
+
|
| 660 |
+
Grammar (all optional, combinable):
|
| 661 |
+
'a AND b' — every term must match (compositional AND)
|
| 662 |
+
'a NOT b' — exclude b (also '-b' or 'a -b')
|
| 663 |
+
'a; b' — same as AND
|
| 664 |
+
Plain text with none of these is a single positive term (classic search).
|
| 665 |
+
"""
|
| 666 |
+
import re
|
| 667 |
+
neg = []
|
| 668 |
+
# split on NOT / leading-minus tokens
|
| 669 |
+
parts = re.split(r'\bNOT\b', text)
|
| 670 |
+
head = parts[0]
|
| 671 |
+
for extra in parts[1:]:
|
| 672 |
+
neg.append(extra.strip())
|
| 673 |
+
pos_raw = re.split(r'\bAND\b|;', head)
|
| 674 |
+
pos = []
|
| 675 |
+
for term in pos_raw:
|
| 676 |
+
term = term.strip()
|
| 677 |
+
# pull out inline -word exclusions
|
| 678 |
+
toks = term.split()
|
| 679 |
+
keep = []
|
| 680 |
+
for tk in toks:
|
| 681 |
+
if tk.startswith("-") and len(tk) > 1:
|
| 682 |
+
neg.append(tk[1:])
|
| 683 |
+
else:
|
| 684 |
+
keep.append(tk)
|
| 685 |
+
if keep:
|
| 686 |
+
pos.append(" ".join(keep))
|
| 687 |
+
pos = [p for p in pos if p]
|
| 688 |
+
neg = [n for n in neg if n]
|
| 689 |
+
return (pos or [text]), neg
|
| 690 |
+
|
| 691 |
+
|
| 692 |
+
def search(store, text, k=10, nprobe=3, merge=True, t0=None, t1=None,
|
| 693 |
+
streams=None, method="auto", neg_weight=0.5, min_score=None,
|
| 694 |
+
percentile=None, rerank=False, rerank_top=12, rerank_alpha=0.7):
|
| 695 |
+
"""Compositional text search. `text` may use AND / NOT / -term:
|
| 696 |
+
'two people AND a laptop NOT a phone'
|
| 697 |
+
`min_score` (absolute cosine floor) or `percentile` (keep top X%) turn
|
| 698 |
+
ranked-everything into precise retrieval."""
|
| 699 |
+
st = store.table("embeddings").state()
|
| 700 |
+
model = st.meta.get("model", DEFAULT_MODEL)
|
| 701 |
+
pos_terms, neg_terms = _parse_query(text)
|
| 702 |
+
pos_vecs = np.stack([embed_text(p, model) for p in pos_terms])
|
| 703 |
+
neg_vecs = (np.stack([embed_text(n, model) for n in neg_terms])
|
| 704 |
+
if neg_terms else None)
|
| 705 |
+
q = pos_vecs.mean(axis=0)
|
| 706 |
+
q /= np.linalg.norm(q) # coarse retrieval direction
|
| 707 |
+
hits, stats = _rank(store, q, k, nprobe, merge=merge, t0=t0, t1=t1,
|
| 708 |
+
streams=streams, method=method, pos_vecs=pos_vecs,
|
| 709 |
+
neg_vecs=neg_vecs, neg_weight=neg_weight,
|
| 710 |
+
min_score=min_score, percentile=percentile)
|
| 711 |
+
stats["positive_terms"] = pos_terms
|
| 712 |
+
stats["negative_terms"] = neg_terms
|
| 713 |
+
if rerank and hits:
|
| 714 |
+
# relational stage: the expensive operator runs LAST, on the pruned set
|
| 715 |
+
from .rerank import rerank_hits
|
| 716 |
+
hits, info = rerank_hits(store, hits, text, top_n=rerank_top,
|
| 717 |
+
alpha=rerank_alpha)
|
| 718 |
+
stats["rerank"] = info
|
| 719 |
+
return hits, stats
|
| 720 |
+
|
| 721 |
+
|
| 722 |
+
def search_text(store, text, k=10, nprobe=3, merge=True, t0=None, t1=None,
|
| 723 |
+
streams=None, method="auto", **kw):
|
| 724 |
+
# backward-compatible alias; forwards compositional kwargs too
|
| 725 |
+
return search(store, text, k=k, nprobe=nprobe, merge=merge, t0=t0, t1=t1,
|
| 726 |
+
streams=streams, method=method, **kw)
|
| 727 |
+
|
| 728 |
+
|
| 729 |
+
def search_clip(store, stream, t0, t1, k=10, nprobe=3, merge=True,
|
| 730 |
+
pt0=None, pt1=None, pstreams=None, method="auto"):
|
| 731 |
+
t, vecs = _vec_table(store)
|
| 732 |
+
s = t.column("stream").to_numpy(zero_copy_only=False)
|
| 733 |
+
a = t.column("ts").to_numpy()
|
| 734 |
+
b = t.column("t1").to_numpy()
|
| 735 |
+
sel = (s == stream) & (a <= t1) & (b >= t0)
|
| 736 |
+
if not sel.any():
|
| 737 |
+
raise ValueError(f"no embedded windows overlap {stream} [{t0},{t1}]")
|
| 738 |
+
q = vecs[sel].mean(axis=0)
|
| 739 |
+
q /= np.linalg.norm(q)
|
| 740 |
+
hits, stats = _rank(store, q, k + 8, nprobe, merge=merge,
|
| 741 |
+
t0=pt0, t1=pt1, streams=pstreams, method=method)
|
| 742 |
+
hits = [h for h in hits
|
| 743 |
+
if not (h["stream"] == stream and h["t0"] <= t1 and h["t1"] >= t0)]
|
| 744 |
+
return hits[:k], stats
|
python/elidedb/encprune.py
ADDED
|
@@ -0,0 +1,307 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FDNN cellular turnover applied to the SigLIP ENCODER — where the time is.
|
| 2 |
+
|
| 3 |
+
WHY HERE AND NOT THE LITTLE TOWER
|
| 4 |
+
---------------------------------
|
| 5 |
+
Ingest cost is measured, and it is not the database: byte-range decode runs at
|
| 6 |
+
2.7 ms/frame while the SigLIP vision tower runs at 27.7-90.3 ms/frame. So 97%
|
| 7 |
+
of ingest is one frozen 428M-parameter ViT. Pruning the 80k-parameter temporal
|
| 8 |
+
tower — which is what was done first — optimises 0.001% of the bill. The
|
| 9 |
+
encoder is the neuron population that matters.
|
| 10 |
+
|
| 11 |
+
THE NEURON, HERE
|
| 12 |
+
----------------
|
| 13 |
+
Each encoder layer is `fc1 (1152 -> 4304) -> gelu -> fc2 (4304 -> 1152)`.
|
| 14 |
+
Hidden unit *c* is a neuron in exactly FDNN's sense: it takes the residual
|
| 15 |
+
stream, applies its own nonlinearity, and writes back through its own column
|
| 16 |
+
of `fc2`. There are 27 x 4304 = 116,208 of them and they are 65% of the
|
| 17 |
+
tower's parameters.
|
| 18 |
+
|
| 19 |
+
Silencing one is `fc1.weight[c] = 0, fc1.bias[c] = 0` — gelu(0) = 0, so the
|
| 20 |
+
channel contributes nothing through fc2. Removing one for real is slicing
|
| 21 |
+
`fc1.weight[keep]`, `fc1.bias[keep]`, `fc2.weight[:, keep]`, which is a
|
| 22 |
+
genuine FLOP reduction, not a multiply by zero.
|
| 23 |
+
|
| 24 |
+
UTILIZATION IS DEFINED BY THE DATABASE'S WORKLOAD
|
| 25 |
+
-------------------------------------------------
|
| 26 |
+
FDNN measures utilization as the increase in validation loss when a neuron is
|
| 27 |
+
silenced. The equivalent here is NOT ImageNet accuracy — this encoder exists
|
| 28 |
+
to produce vectors that a retrieval index ranks with. So utilization is the
|
| 29 |
+
loss of EMBEDDING FIDELITY on frames from the actual corpus:
|
| 30 |
+
|
| 31 |
+
fidelity = mean cosine( pruned_embedding, unpruned_embedding )
|
| 32 |
+
|
| 33 |
+
i.e. the unpruned encoder is its own teacher and the calibration set is the
|
| 34 |
+
user's own data. A channel that matters for photographs of dogs but never
|
| 35 |
+
fires on a robot arm in a toy kitchen is, for this database, dead weight.
|
| 36 |
+
|
| 37 |
+
Ablating 116,208 channels one at a time is not affordable (one forward pass
|
| 38 |
+
each). FDNN's own PPO feature vector already contains the cheap surrogates —
|
| 39 |
+
activation magnitude and downstream weight norm — and their product is the
|
| 40 |
+
standard structured-pruning saliency:
|
| 41 |
+
|
| 42 |
+
saliency(c) = E_frames |act_c| * || fc2[:, c] ||
|
| 43 |
+
|
| 44 |
+
which is exactly "how much signal this neuron actually injects". Ablation is
|
| 45 |
+
still used, but per LAYER (27 measurements, affordable), to calibrate how much
|
| 46 |
+
each layer's saliency scale is worth. PPO then chooses keep-probabilities from
|
| 47 |
+
the same five features FDNN used, and the reward is the same shape:
|
| 48 |
+
|
| 49 |
+
reward = -(fidelity loss ratio) - lambda * kept_fraction
|
| 50 |
+
"""
|
| 51 |
+
from __future__ import annotations
|
| 52 |
+
|
| 53 |
+
import time
|
| 54 |
+
|
| 55 |
+
import mlx.core as mx
|
| 56 |
+
import numpy as np
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
# ===========================================================================
|
| 60 |
+
# 0. Reaching the layers
|
| 61 |
+
# ===========================================================================
|
| 62 |
+
def vision_layers(model):
|
| 63 |
+
"""The encoder layer list of a SigLIP vision tower."""
|
| 64 |
+
vm = model.vision_model
|
| 65 |
+
vm = getattr(vm, "vision_model", vm)
|
| 66 |
+
return list(vm.encoder.layers)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def mlp_width(layer):
|
| 70 |
+
return np.array(layer.mlp.fc1.bias).shape[0]
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
class _Tap:
|
| 74 |
+
"""Wraps an nn.Linear so the mean |output| per channel is recorded.
|
| 75 |
+
|
| 76 |
+
MLX calls `self.fc1(x)` by attribute, so swapping the attribute is enough
|
| 77 |
+
to observe it — no hooks, no forked forward pass that could drift from the
|
| 78 |
+
real one.
|
| 79 |
+
"""
|
| 80 |
+
|
| 81 |
+
def __init__(self, inner):
|
| 82 |
+
self.inner = inner
|
| 83 |
+
self.acc = None
|
| 84 |
+
self.n = 0
|
| 85 |
+
|
| 86 |
+
def __call__(self, x):
|
| 87 |
+
out = self.inner(x)
|
| 88 |
+
a = mx.mean(mx.abs(out).reshape(-1, out.shape[-1]), axis=0)
|
| 89 |
+
self.acc = a if self.acc is None else self.acc + a
|
| 90 |
+
self.n += 1
|
| 91 |
+
return out
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
# ===========================================================================
|
| 95 |
+
# 1. Saliency: activation magnitude x downstream weight norm
|
| 96 |
+
# ===========================================================================
|
| 97 |
+
def channel_saliency(model, images, model_id, batch=16):
|
| 98 |
+
"""Per-layer array of per-channel saliency, measured on REAL frames."""
|
| 99 |
+
from .embeddings import _embed_images
|
| 100 |
+
layers = vision_layers(model)
|
| 101 |
+
taps = []
|
| 102 |
+
for lyr in layers:
|
| 103 |
+
t = _Tap(lyr.mlp.fc1)
|
| 104 |
+
lyr.mlp.fc1 = t
|
| 105 |
+
taps.append(t)
|
| 106 |
+
try:
|
| 107 |
+
for i in range(0, len(images), batch):
|
| 108 |
+
_embed_images(images[i:i + batch], model_id)
|
| 109 |
+
finally:
|
| 110 |
+
for lyr, t in zip(layers, taps):
|
| 111 |
+
lyr.mlp.fc1 = t.inner
|
| 112 |
+
|
| 113 |
+
out = []
|
| 114 |
+
for lyr, t in zip(layers, taps):
|
| 115 |
+
act = np.array(t.acc) / max(t.n, 1)
|
| 116 |
+
down = np.linalg.norm(np.array(lyr.mlp.fc2.weight), axis=0)
|
| 117 |
+
out.append({"act": act.astype(np.float32),
|
| 118 |
+
"down": down.astype(np.float32),
|
| 119 |
+
"saliency": (act * down).astype(np.float32)})
|
| 120 |
+
return out
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
# ===========================================================================
|
| 124 |
+
# 2. Masking and compaction
|
| 125 |
+
# ===========================================================================
|
| 126 |
+
def _snapshot(layers):
|
| 127 |
+
return [(np.array(l.mlp.fc1.weight), np.array(l.mlp.fc1.bias))
|
| 128 |
+
for l in layers]
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def _restore(layers, snap):
|
| 132 |
+
for l, (w, b) in zip(layers, snap):
|
| 133 |
+
l.mlp.fc1.weight = mx.array(w)
|
| 134 |
+
l.mlp.fc1.bias = mx.array(b)
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def apply_masks(layers, snap, masks):
|
| 138 |
+
"""Silence channels: zeroing the fc1 row makes gelu(0)=0 downstream."""
|
| 139 |
+
for l, (w, b), m in zip(layers, snap, masks):
|
| 140 |
+
keep = m.astype(np.float32)[:, None]
|
| 141 |
+
l.mlp.fc1.weight = mx.array(w * keep)
|
| 142 |
+
l.mlp.fc1.bias = mx.array(b * keep[:, 0])
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def compact_mlps(model, masks):
|
| 146 |
+
"""Physically delete the dead channels. This is where speed comes from."""
|
| 147 |
+
layers = vision_layers(model)
|
| 148 |
+
removed = 0
|
| 149 |
+
for l, m in zip(layers, masks):
|
| 150 |
+
keep = np.where(m > 0.5)[0]
|
| 151 |
+
if len(keep) == mlp_width(l):
|
| 152 |
+
continue
|
| 153 |
+
removed += mlp_width(l) - len(keep)
|
| 154 |
+
l.mlp.fc1.weight = mx.array(np.array(l.mlp.fc1.weight)[keep])
|
| 155 |
+
l.mlp.fc1.bias = mx.array(np.array(l.mlp.fc1.bias)[keep])
|
| 156 |
+
l.mlp.fc2.weight = mx.array(np.array(l.mlp.fc2.weight)[:, keep])
|
| 157 |
+
mx.eval(model.parameters())
|
| 158 |
+
return removed
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# ===========================================================================
|
| 162 |
+
# 3. Fidelity against the unpruned teacher, on the user's own frames
|
| 163 |
+
# ===========================================================================
|
| 164 |
+
def fidelity(model_id, images, teacher, batch=16):
|
| 165 |
+
from .embeddings import _embed_images
|
| 166 |
+
out = []
|
| 167 |
+
for i in range(0, len(images), batch):
|
| 168 |
+
out.append(_embed_images(images[i:i + batch], model_id))
|
| 169 |
+
v = np.concatenate(out, axis=0)
|
| 170 |
+
return float((v * teacher).sum(axis=1).mean())
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
# ===========================================================================
|
| 174 |
+
# 4. Reverse attention (unchanged in spirit from FDNN)
|
| 175 |
+
# ===========================================================================
|
| 176 |
+
def reverse_attention(importance, temperature=1.0):
|
| 177 |
+
imp = np.asarray(importance, dtype=np.float64)
|
| 178 |
+
imp = (imp - imp.mean()) / (imp.std() + 1e-8)
|
| 179 |
+
logits = -imp / max(temperature, 1e-6)
|
| 180 |
+
logits -= logits.max()
|
| 181 |
+
e = np.exp(logits)
|
| 182 |
+
return (e / (e.sum() + 1e-12)).astype(np.float32)
|
| 183 |
+
|
| 184 |
+
|
| 185 |
+
# ===========================================================================
|
| 186 |
+
# 5. The cycle
|
| 187 |
+
# ===========================================================================
|
| 188 |
+
def prune_encoder(model_id, images, keep=0.5, layer_probe=True, iters=12,
|
| 189 |
+
lam=0.35, batch=16, verbose=True, seed=0):
|
| 190 |
+
"""Prune MLP channels of the vision tower against corpus fidelity.
|
| 191 |
+
|
| 192 |
+
`keep` is the global target fraction of MLP channels to retain. The
|
| 193 |
+
per-layer budget is not uniform: layers whose ablation barely moves
|
| 194 |
+
fidelity give up more channels than layers that matter, which is the whole
|
| 195 |
+
point of measuring instead of assuming.
|
| 196 |
+
|
| 197 |
+
Returns (model, report). The returned model is COMPACTED — smaller
|
| 198 |
+
matmuls, not masked ones.
|
| 199 |
+
"""
|
| 200 |
+
from .embeddings import _embed_images, _load_model
|
| 201 |
+
rng = np.random.default_rng(seed)
|
| 202 |
+
model, _ = _load_model(model_id)
|
| 203 |
+
layers = vision_layers(model)
|
| 204 |
+
widths = [mlp_width(l) for l in layers]
|
| 205 |
+
snap = _snapshot(layers)
|
| 206 |
+
|
| 207 |
+
teacher = np.concatenate(
|
| 208 |
+
[_embed_images(images[i:i + batch], model_id)
|
| 209 |
+
for i in range(0, len(images), batch)], axis=0)
|
| 210 |
+
|
| 211 |
+
t0 = time.time()
|
| 212 |
+
sal = channel_saliency(model, images, model_id, batch=batch)
|
| 213 |
+
if verbose:
|
| 214 |
+
print(f" saliency over {len(images)} corpus frames "
|
| 215 |
+
f"({time.time() - t0:.1f}s)", flush=True)
|
| 216 |
+
|
| 217 |
+
# ---- per-layer ablation: how much does this layer matter at all? -------
|
| 218 |
+
layer_cost = np.ones(len(layers), dtype=np.float64)
|
| 219 |
+
if layer_probe:
|
| 220 |
+
for li in range(len(layers)):
|
| 221 |
+
masks = [np.ones(w, np.float32) for w in widths]
|
| 222 |
+
masks[li][:] = 0.0
|
| 223 |
+
apply_masks(layers, snap, masks)
|
| 224 |
+
layer_cost[li] = max(1.0 - fidelity(model_id, images[:batch],
|
| 225 |
+
teacher[:batch], batch), 1e-6)
|
| 226 |
+
_restore(layers, snap)
|
| 227 |
+
if verbose:
|
| 228 |
+
order = np.argsort(layer_cost)
|
| 229 |
+
print(f" layer ablation: cheapest {order[:4].tolist()} "
|
| 230 |
+
f"costliest {order[-4:].tolist()}", flush=True)
|
| 231 |
+
|
| 232 |
+
# ---- allocate the budget across layers by measured importance ---------
|
| 233 |
+
# A layer that costs little when removed entirely can afford to lose more
|
| 234 |
+
# of its channels. Normalised so the global kept fraction hits `keep`.
|
| 235 |
+
w = layer_cost / layer_cost.sum()
|
| 236 |
+
share = w / w.mean() # 1.0 == average importance
|
| 237 |
+
per_layer_keep = np.clip(keep * share, 0.05, 1.0)
|
| 238 |
+
total = sum(widths)
|
| 239 |
+
scale = (keep * total) / sum(k * n for k, n in zip(per_layer_keep, widths))
|
| 240 |
+
per_layer_keep = np.clip(per_layer_keep * scale, 0.05, 1.0)
|
| 241 |
+
|
| 242 |
+
masks = []
|
| 243 |
+
for li, (s, wdt) in enumerate(zip(sal, widths)):
|
| 244 |
+
n_keep = max(int(round(per_layer_keep[li] * wdt)), 1)
|
| 245 |
+
idx = np.argsort(-s["saliency"])[:n_keep]
|
| 246 |
+
m = np.zeros(wdt, np.float32)
|
| 247 |
+
m[idx] = 1.0
|
| 248 |
+
masks.append(m)
|
| 249 |
+
|
| 250 |
+
apply_masks(layers, snap, masks)
|
| 251 |
+
fid = fidelity(model_id, images, teacher, batch)
|
| 252 |
+
kept = sum(m.sum() for m in masks) / total
|
| 253 |
+
if verbose:
|
| 254 |
+
print(f" saliency prune: kept {kept:.1%}, fidelity {fid:.4f}",
|
| 255 |
+
flush=True)
|
| 256 |
+
|
| 257 |
+
# ---- PPO refinement over the per-layer budget --------------------------
|
| 258 |
+
# The candidate set is the 27 per-layer keep fractions rather than 116,208
|
| 259 |
+
# independent channels: one forward pass per sampled mask makes per-channel
|
| 260 |
+
# sampling unaffordable, and the per-layer budget is where the leverage
|
| 261 |
+
# actually is (saliency already orders channels within a layer).
|
| 262 |
+
best = (fid, [m.copy() for m in masks], kept)
|
| 263 |
+
ra = reverse_attention(layer_cost)
|
| 264 |
+
logit = np.zeros(len(layers))
|
| 265 |
+
for it in range(iters):
|
| 266 |
+
cand = per_layer_keep * (1.0 + 0.25 * np.tanh(logit)
|
| 267 |
+
+ 0.15 * rng.standard_normal(len(layers)))
|
| 268 |
+
cand = np.clip(cand, 0.05, 1.0)
|
| 269 |
+
cand *= (keep * total) / sum(c * n for c, n in zip(cand, widths))
|
| 270 |
+
cand = np.clip(cand, 0.05, 1.0)
|
| 271 |
+
trial = []
|
| 272 |
+
for li, (s, wdt) in enumerate(zip(sal, widths)):
|
| 273 |
+
n_keep = max(int(round(cand[li] * wdt)), 1)
|
| 274 |
+
idx = np.argsort(-s["saliency"])[:n_keep]
|
| 275 |
+
m = np.zeros(wdt, np.float32)
|
| 276 |
+
m[idx] = 1.0
|
| 277 |
+
trial.append(m)
|
| 278 |
+
apply_masks(layers, snap, trial)
|
| 279 |
+
f = fidelity(model_id, images[:batch * 2], teacher[:batch * 2], batch)
|
| 280 |
+
k = sum(m.sum() for m in trial) / total
|
| 281 |
+
reward = -(1.0 - f) - lam * k
|
| 282 |
+
best_reward = -(1.0 - best[0]) - lam * best[2]
|
| 283 |
+
if reward > best_reward:
|
| 284 |
+
# Direction of the accepted move, computed BEFORE the budget is
|
| 285 |
+
# updated — comparing `cand` against itself would make every sign
|
| 286 |
+
# zero and the search a pure random walk.
|
| 287 |
+
step = np.sign(cand - per_layer_keep)
|
| 288 |
+
best = (f, [m.copy() for m in trial], k)
|
| 289 |
+
per_layer_keep = cand
|
| 290 |
+
logit += 0.5 * ra * step
|
| 291 |
+
if verbose and (it % 4 == 0 or it == iters - 1):
|
| 292 |
+
print(f" PPO {it:2d} | fidelity {f:.4f} | kept {k:.1%} | "
|
| 293 |
+
f"reward {reward:+.4f}", flush=True)
|
| 294 |
+
|
| 295 |
+
fid, masks, kept = best
|
| 296 |
+
_restore(layers, snap)
|
| 297 |
+
removed = compact_mlps(model, masks)
|
| 298 |
+
final = fidelity(model_id, images, teacher, batch)
|
| 299 |
+
report = {"model": model_id, "kept_fraction": float(kept),
|
| 300 |
+
"channels_removed": int(removed),
|
| 301 |
+
"channels_total": int(total),
|
| 302 |
+
"fidelity": float(final),
|
| 303 |
+
"layer_keep": [float(x) for x in per_layer_keep]}
|
| 304 |
+
if verbose:
|
| 305 |
+
print(f" compacted: removed {removed:,}/{total:,} MLP channels, "
|
| 306 |
+
f"fidelity {final:.4f}", flush=True)
|
| 307 |
+
return model, report
|
python/elidedb/fdnnstudent.py
ADDED
|
@@ -0,0 +1,231 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FDNN students shaped like their teachers.
|
| 2 |
+
|
| 3 |
+
A student is not a generic head bolted onto whatever vectors happen to
|
| 4 |
+
be lying around. That was tried and it failed for a reason worth
|
| 5 |
+
keeping: an attention-pool + ReLU MLP over FDNN-V frame vectors reached
|
| 6 |
+
0.9168 cosine to the PE teacher while a CONSTANT prediction of the
|
| 7 |
+
corpus mean scored 0.8629, and nearest-neighbour agreement was 0.005.
|
| 8 |
+
The head had learned the mean. A least-squares oracle on the same input
|
| 9 |
+
reached nn_top1 0.0155, so no head over those vectors could have worked
|
| 10 |
+
- the information was not in them.
|
| 11 |
+
|
| 12 |
+
So the student takes the TEACHER'S OWN ARCHITECTURE as its skeleton,
|
| 13 |
+
reads pixels like the teacher does, and shrinks. What changes is the
|
| 14 |
+
neuron, per the three FDNN rules:
|
| 15 |
+
|
| 16 |
+
RULE 1 each neuron is a sub-network. Every transformer block's MLP -
|
| 17 |
+
the 2/3 of a ViT's parameters that is elementwise GELU -
|
| 18 |
+
becomes a KAN-style sum over k heterogeneous bases of the
|
| 19 |
+
same pre-activation: FINER (variable-period oscillator),
|
| 20 |
+
Gabor (wavelet, fires on a burst), sine, poly-phase (chirp,
|
| 21 |
+
fires on acceleration). omega bands partition the spectrum so
|
| 22 |
+
different channels answer to different rates of change.
|
| 23 |
+
A GELU unit can only answer "how much"; these answer "how
|
| 24 |
+
much, how fast, and is it accelerating" - which is what a
|
| 25 |
+
video model's MLP is being asked for in the first place.
|
| 26 |
+
|
| 27 |
+
RULE 2 neurogenesis and apoptosis after training, each followed by a
|
| 28 |
+
re-settle fine-tune. `mask` is the aliveness vector the cycle
|
| 29 |
+
writes; it is frozen against the optimizer, because AdamW
|
| 30 |
+
weight-decays an unfrozen mask off 1.0 and every channel
|
| 31 |
+
quietly shrinks (measured on the context tower).
|
| 32 |
+
|
| 33 |
+
RULE 3 PPO + reverse attention decide who lives: utilization is the
|
| 34 |
+
drop in held-out fidelity when a channel is silenced.
|
| 35 |
+
|
| 36 |
+
WHY THE SKELETON MATTERS. V-JEPA2 is 24 blocks of width 1024 over 8,192
|
| 37 |
+
spatiotemporal tokens (256px, 64 frames, patch 16, tubelet 2). Most of
|
| 38 |
+
that cost is token count, not depth. A student that keeps the tubelet
|
| 39 |
+
embedding and the attentive pooler - the parts that decide WHAT is
|
| 40 |
+
compared - while cutting resolution, frames, width and depth, is doing
|
| 41 |
+
the teacher's computation at the teacher's shape. A student with a
|
| 42 |
+
different shape is a different model that happens to be trained on the
|
| 43 |
+
teacher's outputs, which is what failed.
|
| 44 |
+
"""
|
| 45 |
+
from __future__ import annotations
|
| 46 |
+
|
| 47 |
+
import numpy as np
|
| 48 |
+
|
| 49 |
+
import mlx.core as mx
|
| 50 |
+
import mlx.nn as nn
|
| 51 |
+
|
| 52 |
+
# Basis identifiers, matching FDNNTemporalCell so the two implementations
|
| 53 |
+
# cannot drift apart: 0 FINER, 1 Gabor, 2 sine, 3 poly-phase.
|
| 54 |
+
FINER, GABOR, SINE, POLY = 0, 1, 2, 3
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
class FDNNFeedForward(nn.Module):
|
| 58 |
+
"""Rule 1, applied to a transformer block's MLP.
|
| 59 |
+
|
| 60 |
+
Standard: y = W2 @ gelu(W1 @ x). Here each of C output channels is a
|
| 61 |
+
KAN sum over k sub-functions of the SAME pre-activation, so one
|
| 62 |
+
channel can be periodic, another can fire only on a burst, another
|
| 63 |
+
only on acceleration - heterogeneous by construction rather than by
|
| 64 |
+
hoping a homogeneous nonlinearity specialises.
|
| 65 |
+
|
| 66 |
+
omega bands are spatial-temporal here, not purely temporal: tokens
|
| 67 |
+
are tubelets, so a band is a rate of change across the token grid.
|
| 68 |
+
"""
|
| 69 |
+
|
| 70 |
+
def __init__(self, dim, channels, k_width=4,
|
| 71 |
+
omega_bands=(0.8, 2.5, 8.0),
|
| 72 |
+
band_fractions=(0.34, 0.33, 0.33), seed=0):
|
| 73 |
+
super().__init__()
|
| 74 |
+
rng = np.random.default_rng(seed)
|
| 75 |
+
C, k = channels, k_width
|
| 76 |
+
self.C, self.k = C, k
|
| 77 |
+
|
| 78 |
+
om = []
|
| 79 |
+
for b, fr in zip(omega_bands, band_fractions):
|
| 80 |
+
om.extend([b] * int(round(fr * C)))
|
| 81 |
+
om = (om + [omega_bands[-1]] * C)[:C]
|
| 82 |
+
self.omegas_per_neuron = np.array(om, np.float32)
|
| 83 |
+
self.omegas = mx.array(np.repeat(self.omegas_per_neuron, k))
|
| 84 |
+
|
| 85 |
+
half, quarter = max(k // 2, 1), max(k // 4, 1)
|
| 86 |
+
per = np.array([FINER] * half + [GABOR] * quarter
|
| 87 |
+
+ [POLY] * max(k - half - quarter, 0), np.int32)[:k]
|
| 88 |
+
self.basis_types = mx.array(np.tile(per, C).astype(np.int32))
|
| 89 |
+
# categorical, not a weight: unfrozen the optimizer drifts it off
|
| 90 |
+
# its exact values and silently reroutes every Gabor to poly
|
| 91 |
+
self.freeze(keys=["basis_types"], recurse=False)
|
| 92 |
+
|
| 93 |
+
mean_om = float(self.omegas_per_neuron.mean())
|
| 94 |
+
lim = float(np.sqrt(6.0 / dim) / mean_om)
|
| 95 |
+
self.W1 = mx.array(rng.uniform(-lim, lim, (dim, C * k)).astype(np.float32))
|
| 96 |
+
self.b1 = mx.array(rng.uniform(-2.0, 2.0, (C * k,)).astype(np.float32))
|
| 97 |
+
self.phases = mx.array(rng.uniform(0, 2 * np.pi, (C * k,)).astype(np.float32))
|
| 98 |
+
self.gabor_s = mx.array(rng.uniform(0.3, 1.5, (C * k,)).astype(np.float32))
|
| 99 |
+
log_om = np.log(np.clip(np.repeat(self.omegas_per_neuron, k), 1e-3, None))
|
| 100 |
+
self.log_alpha = mx.array(rng.uniform(np.minimum(0.0, log_om),
|
| 101 |
+
np.maximum(0.0, log_om) + 1e-6
|
| 102 |
+
).astype(np.float32))
|
| 103 |
+
w2s = float(np.sqrt(6.0 / (C * k)))
|
| 104 |
+
self.w2 = mx.array(rng.uniform(-w2s, w2s, (C, k)).astype(np.float32))
|
| 105 |
+
pl = float(np.sqrt(6.0 / C))
|
| 106 |
+
self.Wp = mx.array(rng.uniform(-pl, pl, (C, dim)).astype(np.float32))
|
| 107 |
+
|
| 108 |
+
self.mask = mx.array(np.ones((C,), np.float32)) # rule 2
|
| 109 |
+
self.freeze(keys=["mask"], recurse=False)
|
| 110 |
+
|
| 111 |
+
def channels_out(self, x):
|
| 112 |
+
"""Per-channel activation before the mask - rule 3's signal."""
|
| 113 |
+
pre = x @ self.W1 + self.b1
|
| 114 |
+
om_h = self.omegas * pre
|
| 115 |
+
sq = pre * pre
|
| 116 |
+
alpha = mx.exp(self.log_alpha)
|
| 117 |
+
finer = mx.sin(self.omegas * (mx.abs(pre) + 1.0) * pre + self.phases)
|
| 118 |
+
gab = mx.exp(-(self.gabor_s ** 2) * sq) * mx.sin(om_h + self.phases)
|
| 119 |
+
sine = mx.sin(om_h + self.phases)
|
| 120 |
+
poly = mx.sin(alpha * sq + om_h + self.phases)
|
| 121 |
+
acts = mx.where(self.basis_types == FINER, finer,
|
| 122 |
+
mx.where(self.basis_types == GABOR, gab,
|
| 123 |
+
mx.where(self.basis_types == SINE, sine, poly)))
|
| 124 |
+
acts = acts.reshape(*pre.shape[:-1], self.C, self.k)
|
| 125 |
+
return mx.sum(acts * self.w2, axis=-1)
|
| 126 |
+
|
| 127 |
+
def __call__(self, x):
|
| 128 |
+
return (self.channels_out(x) * self.mask) @ self.Wp
|
| 129 |
+
|
| 130 |
+
def set_active_mask(self, m):
|
| 131 |
+
self.mask = mx.array(np.asarray(m, np.float32))
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
class Block(nn.Module):
|
| 135 |
+
"""A ViT block with the teacher's shape and an FDNN neuron inside."""
|
| 136 |
+
|
| 137 |
+
def __init__(self, dim, heads, channels, k_width=4, seed=0):
|
| 138 |
+
super().__init__()
|
| 139 |
+
self.n1 = nn.LayerNorm(dim)
|
| 140 |
+
self.attn = nn.MultiHeadAttention(dim, heads)
|
| 141 |
+
self.n2 = nn.LayerNorm(dim)
|
| 142 |
+
self.ff = FDNNFeedForward(dim, channels, k_width, seed=seed)
|
| 143 |
+
|
| 144 |
+
def __call__(self, x):
|
| 145 |
+
h = self.n1(x)
|
| 146 |
+
x = x + self.attn(h, h, h)
|
| 147 |
+
return x + self.ff(self.n2(x))
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
class TubeletEmbed(nn.Module):
|
| 151 |
+
"""The teacher's input stage: non-overlapping spatiotemporal patches.
|
| 152 |
+
|
| 153 |
+
Kept because it is what decides what a token IS. Resolution, frame
|
| 154 |
+
count, width and depth are all shrunk; the tokenisation is not
|
| 155 |
+
changed, so the student is comparing the same kind of thing.
|
| 156 |
+
"""
|
| 157 |
+
|
| 158 |
+
def __init__(self, dim, patch=16, tubelet=2, in_ch=3):
|
| 159 |
+
super().__init__()
|
| 160 |
+
self.patch, self.tubelet = patch, tubelet
|
| 161 |
+
self.proj = nn.Linear(in_ch * patch * patch * tubelet, dim)
|
| 162 |
+
|
| 163 |
+
def __call__(self, v): # (B, T, H, W, 3) in [-1,1]
|
| 164 |
+
B, T, H, W, C = v.shape
|
| 165 |
+
p, t = self.patch, self.tubelet
|
| 166 |
+
gt, gh, gw = T // t, H // p, W // p
|
| 167 |
+
v = v[:, :gt * t, :gh * p, :gw * p]
|
| 168 |
+
v = v.reshape(B, gt, t, gh, p, gw, p, C)
|
| 169 |
+
v = v.transpose(0, 1, 3, 5, 2, 4, 6, 7).reshape(B, gt * gh * gw, -1)
|
| 170 |
+
return self.proj(v), (gt, gh, gw)
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
class AttentivePool(nn.Module):
|
| 174 |
+
"""The teacher pools with a learned query, so the student does too.
|
| 175 |
+
Mean pooling buries the few tokens that carry the event under the
|
| 176 |
+
many that carry the unchanged room."""
|
| 177 |
+
|
| 178 |
+
def __init__(self, dim, out_dim):
|
| 179 |
+
super().__init__()
|
| 180 |
+
self.q = mx.array((np.random.default_rng(0).normal(size=(dim,))
|
| 181 |
+
/ np.sqrt(dim)).astype(np.float32))
|
| 182 |
+
self.norm = nn.LayerNorm(dim)
|
| 183 |
+
self.proj = nn.Linear(dim, out_dim)
|
| 184 |
+
|
| 185 |
+
def __call__(self, x):
|
| 186 |
+
a = mx.softmax(x @ self.q, axis=1)[..., None]
|
| 187 |
+
y = self.proj(self.norm(mx.sum(a * x, axis=1)))
|
| 188 |
+
return y / (mx.linalg.norm(y, axis=-1, keepdims=True) + 1e-8)
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
class VideoStudent(nn.Module):
|
| 192 |
+
"""V-JEPA2's skeleton at 1/N the size, with FDNN neurons.
|
| 193 |
+
|
| 194 |
+
teacher 256px, 64 frames, patch 16, tubelet 2 -> 8192 tokens,
|
| 195 |
+
24 blocks, width 1024, 16 heads
|
| 196 |
+
student configurable; the default cuts tokens 16x and width 5x,
|
| 197 |
+
which is where a ViT's cost actually lives.
|
| 198 |
+
"""
|
| 199 |
+
|
| 200 |
+
def __init__(self, out_dim=1024, dim=192, depth=4, heads=3,
|
| 201 |
+
channels=256, k_width=4, patch=16, tubelet=2,
|
| 202 |
+
size=128, frames=16, seed=0):
|
| 203 |
+
super().__init__()
|
| 204 |
+
self.size, self.frames = size, frames
|
| 205 |
+
self.embed = TubeletEmbed(dim, patch, tubelet)
|
| 206 |
+
n_tok = (frames // tubelet) * (size // patch) ** 2
|
| 207 |
+
self.pos = mx.array((np.random.default_rng(seed).normal(
|
| 208 |
+
size=(1, n_tok, dim)) * 0.02).astype(np.float32))
|
| 209 |
+
self.blocks = [Block(dim, heads, channels, k_width, seed=seed + i)
|
| 210 |
+
for i in range(depth)]
|
| 211 |
+
self.pool = AttentivePool(dim, out_dim)
|
| 212 |
+
|
| 213 |
+
def __call__(self, v):
|
| 214 |
+
x, _ = self.embed(v)
|
| 215 |
+
x = x + self.pos[:, :x.shape[1]]
|
| 216 |
+
for b in self.blocks:
|
| 217 |
+
x = b(x)
|
| 218 |
+
return self.pool(x)
|
| 219 |
+
|
| 220 |
+
def utilization(self, v):
|
| 221 |
+
"""Rule 3's raw signal: per-channel mean |activation| per block."""
|
| 222 |
+
x, _ = self.embed(v)
|
| 223 |
+
x = x + self.pos[:, :x.shape[1]]
|
| 224 |
+
out = []
|
| 225 |
+
for b in self.blocks:
|
| 226 |
+
h = b.n1(x)
|
| 227 |
+
x = x + b.attn(h, h, h)
|
| 228 |
+
c = b.ff.channels_out(b.n2(x))
|
| 229 |
+
out.append(np.abs(np.array(c)).mean((0, 1)))
|
| 230 |
+
x = x + b.ff(b.n2(x))
|
| 231 |
+
return out
|
python/elidedb/fdnnv2.py
ADDED
|
@@ -0,0 +1,249 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FDNN-V2: the context head. Same speed class, different training signal.
|
| 2 |
+
|
| 3 |
+
V1's measured dead end: distilled to an appearance teacher whose own space
|
| 4 |
+
separates close/open at 0.60, no student can learn verbs — the target lacks
|
| 5 |
+
them. V2 keeps V1's skeleton (stem -> FDNN recurrent cell, 0.27 ms/frame)
|
| 6 |
+
and warm-starts from its weights; what changes is WHERE the gradient comes
|
| 7 |
+
from (docs/FDNNV2_PLAN.md, all research-grounded):
|
| 8 |
+
|
| 9 |
+
L1 PREDICTIVE (V-JEPA-style, arXiv 2506.09985): from state h_t, predict
|
| 10 |
+
the embedding of frame t+k at two horizons. The future is free
|
| 11 |
+
supervision on every frame, and dynamics must enter h to predict it.
|
| 12 |
+
This is the biological claim made mechanical: understanding = a scene
|
| 13 |
+
model good enough to predict what happens next.
|
| 14 |
+
L2 ARROW OF TIME (Wei et al., CVPR 2018): classify forward vs reversed
|
| 15 |
+
clips from the final state. Open and close are time-reversals; one
|
| 16 |
+
binary head separates them by construction.
|
| 17 |
+
L3 VERB-FOCUSED CONTRASTIVE (arXiv 2304.06708): sigmoid contrastive
|
| 18 |
+
between clip context embeddings and 7B captions, with hard negatives
|
| 19 |
+
built by verb/direction swaps — the joint space queries actually live
|
| 20 |
+
in. A 2-layer text adapter maps SigLIP text embeddings into the
|
| 21 |
+
256-d context space (~0.1 ms per query).
|
| 22 |
+
|
| 23 |
+
The appearance head keeps its V1 distillation target so every existing
|
| 24 |
+
index and the SigLIP text tower continue to work unchanged.
|
| 25 |
+
|
| 26 |
+
THE GATE IS AN EVENT DETECTOR. The cell's update gate z_t measures how much
|
| 27 |
+
of the scene model this frame is allowed to overwrite. Sustained high gate
|
| 28 |
+
activity marks an event boundary — so the encoder emits event segmentation
|
| 29 |
+
for free, and event embeddings are pooled with NOVELTY WEIGHTS (gate
|
| 30 |
+
activity), not uniformly: the moment the drawer closes outweighs the seconds
|
| 31 |
+
it sat still (the TempMe lesson, arXiv 2409.01156, applied at pooling time).
|
| 32 |
+
"""
|
| 33 |
+
from __future__ import annotations
|
| 34 |
+
|
| 35 |
+
import json
|
| 36 |
+
import re
|
| 37 |
+
from pathlib import Path
|
| 38 |
+
|
| 39 |
+
import mlx.core as mx
|
| 40 |
+
import mlx.nn as nn
|
| 41 |
+
import numpy as np
|
| 42 |
+
|
| 43 |
+
from .fdnnvideo import EMBED_DIM, FDNNVideoEncoder
|
| 44 |
+
|
| 45 |
+
CTX_DIM = 256
|
| 46 |
+
|
| 47 |
+
# Verb/direction swaps for L3 hard negatives: shared lexicon (moved to
|
| 48 |
+
# lexicon.py so the QUERY path can import it without this module's mlx)
|
| 49 |
+
from .lexicon import derived_swaps # noqa: E402,F401
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def swap_verbs(text: str, rng, store=None) -> str | None:
|
| 53 |
+
"""One randomly chosen applicable swap -> a hard negative. None if no
|
| 54 |
+
swap applies (caption has no directional content to invert)."""
|
| 55 |
+
# Training negatives obey the no-hardwire rule too: a model taught
|
| 56 |
+
# from hand-written oppositions has the hand-writing baked into its
|
| 57 |
+
# weights, which is worse than a lookup table because it cannot be
|
| 58 |
+
# grepped out afterwards.
|
| 59 |
+
pairs = list(derived_swaps(store)) if store is not None else []
|
| 60 |
+
if not pairs:
|
| 61 |
+
return None
|
| 62 |
+
t = " " + text.lower() + " "
|
| 63 |
+
hits = []
|
| 64 |
+
for a, b in pairs:
|
| 65 |
+
if f" {a} " in t:
|
| 66 |
+
hits.append((a, b))
|
| 67 |
+
if f" {b} " in t:
|
| 68 |
+
hits.append((b, a))
|
| 69 |
+
if not hits:
|
| 70 |
+
# clause order flip is the fallback inversion for compounds
|
| 71 |
+
parts = re.split(r"\band then\b|\bthen\b|,", text)
|
| 72 |
+
if len(parts) >= 2:
|
| 73 |
+
return " then ".join(p.strip() for p in reversed(parts)
|
| 74 |
+
if p.strip())
|
| 75 |
+
return None
|
| 76 |
+
a, b = hits[rng.integers(len(hits))]
|
| 77 |
+
return re.sub(rf"\b{re.escape(a)}\b", b, text.lower(), count=1)
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
class TextAdapter(nn.Module):
|
| 81 |
+
"""SigLIP text embedding (1152) -> context space (256). Two layers,
|
| 82 |
+
~0.7M params, ~0.1 ms — the entire query-time cost of verb awareness."""
|
| 83 |
+
|
| 84 |
+
def __init__(self, in_dim=EMBED_DIM, dim=CTX_DIM):
|
| 85 |
+
super().__init__()
|
| 86 |
+
self.fc1 = nn.Linear(in_dim, 512)
|
| 87 |
+
self.fc2 = nn.Linear(512, dim)
|
| 88 |
+
|
| 89 |
+
def __call__(self, x):
|
| 90 |
+
e = self.fc2(nn.silu(self.fc1(x)))
|
| 91 |
+
return e * mx.rsqrt(mx.sum(e * e, axis=-1, keepdims=True) + 1e-8)
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
class FDNNv2(nn.Module):
|
| 95 |
+
"""V1 skeleton + context head + predictors + arrow-of-time head."""
|
| 96 |
+
|
| 97 |
+
def __init__(self, base: FDNNVideoEncoder, ctx_dim=CTX_DIM,
|
| 98 |
+
horizons=(2, 10), motion_dim=64):
|
| 99 |
+
super().__init__()
|
| 100 |
+
self.base = base
|
| 101 |
+
C = base.channels
|
| 102 |
+
G = base.cfg["glimpse"]
|
| 103 |
+
self.horizons = list(horizons)
|
| 104 |
+
self.motion_dim = motion_dim
|
| 105 |
+
# MOTION PATHWAY — the reversal-symmetry breaker. The gated cell is a
|
| 106 |
+
# leaky integrator; for smooth inputs an EMA is nearly order-invariant,
|
| 107 |
+
# and stage-A measured exactly that: AoT stuck at chance (0.47) while
|
| 108 |
+
# dominating the loss. Frame-difference features flip SIGN under time
|
| 109 |
+
# reversal, so direction becomes linearly readable. Velocity is not a
|
| 110 |
+
# nicety here; it is the only anti-symmetric signal in the model.
|
| 111 |
+
self.motion = nn.Linear(G, motion_dim)
|
| 112 |
+
# context head reads (glimpse, state, motion)
|
| 113 |
+
self.head_ctx = nn.Linear(G + C + motion_dim, ctx_dim)
|
| 114 |
+
# Predictors target the model's OWN future glimpse-change (stop-grad).
|
| 115 |
+
# A2 measured why not student-embedding changes: adjacent-frame true
|
| 116 |
+
# change is ~0.2 in norm while the student's own error is ~0.45 — the
|
| 117 |
+
# difference of noisy embeddings is noise, and the predictors sat at
|
| 118 |
+
# cosine 0 for 8 epochs. Own-latent prediction is JEPA's actual
|
| 119 |
+
# recipe; the appearance anchor stops the stem collapsing to constants.
|
| 120 |
+
self.pred = [nn.Linear(C + motion_dim, G) for _ in horizons]
|
| 121 |
+
# AoT reads the velocity SEQUENCE through a temporal conv, not the
|
| 122 |
+
# mean: this corpus is reciprocal motion (arm out, arm back), so mean
|
| 123 |
+
# signed velocity cancels — measured at chance twice. Order within
|
| 124 |
+
# the window is the signal; a conv kernel can be asymmetric in time.
|
| 125 |
+
self.aot_conv = nn.Conv1d(motion_dim, 32, 5, padding=2)
|
| 126 |
+
self.head_aot = nn.Linear(32, 1)
|
| 127 |
+
self.cfg = {"ctx_dim": ctx_dim, "horizons": self.horizons,
|
| 128 |
+
"motion_dim": motion_dim, "base": base.cfg}
|
| 129 |
+
|
| 130 |
+
# ---- forward over a sequence, exposing everything the losses and the
|
| 131 |
+
# event segmenter need ---------------------------------------------------
|
| 132 |
+
def run(self, seq, h0=None):
|
| 133 |
+
"""(B,T,H,W,3) -> dict of app (B,T,1152), ctx (B,T,ctx), preds,
|
| 134 |
+
gates (B,T), h_last."""
|
| 135 |
+
b = self.base
|
| 136 |
+
B, T = seq.shape[0], seq.shape[1]
|
| 137 |
+
g = b.stem(seq.reshape(B * T, *seq.shape[2:])).reshape(B, T, -1)
|
| 138 |
+
h = b.init_state(B) if h0 is None else h0
|
| 139 |
+
hs, gates = [], []
|
| 140 |
+
cell = b.cell
|
| 141 |
+
for t in range(T):
|
| 142 |
+
z = mx.sigmoid(g[:, t] @ cell.Wz + h @ cell.Uz + cell.bz)
|
| 143 |
+
h = cell(g[:, t], h)
|
| 144 |
+
hs.append(h)
|
| 145 |
+
gates.append(mx.mean(z, axis=-1))
|
| 146 |
+
hs = mx.stack(hs, axis=1) # (B,T,C)
|
| 147 |
+
gates = mx.stack(gates, axis=1) # (B,T)
|
| 148 |
+
app = b.head_g(g) + b.head_h(hs)
|
| 149 |
+
app = app * mx.rsqrt(mx.sum(app * app, axis=-1, keepdims=True) + 1e-8)
|
| 150 |
+
# signed velocity of the glimpse; first frame gets zero motion.
|
| 151 |
+
# NORMALISED: measured, ||dg|| is real motion (corr 0.63 with pixel
|
| 152 |
+
# motion, 5x on moving frames) but only ~2% of the feature norm — fed
|
| 153 |
+
# raw, it drowned next to signals 50x larger and every motion head
|
| 154 |
+
# starved. Direction is unit-normalised; magnitude re-enters as a
|
| 155 |
+
# bounded gain, so both the WHAT and the HOW-MUCH of motion survive.
|
| 156 |
+
dg = mx.concatenate([mx.zeros_like(g[:, :1]),
|
| 157 |
+
g[:, 1:] - g[:, :-1]], axis=1)
|
| 158 |
+
mag = mx.sqrt(mx.sum(dg * dg, axis=-1, keepdims=True) + 1e-8)
|
| 159 |
+
mfeat = nn.silu(self.motion(dg / mag)) * mx.tanh(mag)
|
| 160 |
+
cat = mx.concatenate([g, hs, mfeat], axis=-1)
|
| 161 |
+
ctx = self.head_ctx(cat)
|
| 162 |
+
ctx = ctx * mx.rsqrt(mx.sum(ctx * ctx, axis=-1, keepdims=True) + 1e-8)
|
| 163 |
+
hm = mx.concatenate([hs, mfeat], axis=-1)
|
| 164 |
+
preds = [p(hm) for p in self.pred] # each (B,T,G)
|
| 165 |
+
a = nn.silu(self.aot_conv(mfeat)) # (B,T,32) time-conv
|
| 166 |
+
aot = self.head_aot(mx.mean(a, axis=1))[:, 0]
|
| 167 |
+
return {"app": app, "ctx": ctx, "preds": preds, "gates": gates,
|
| 168 |
+
"aot": aot, "h": h, "g": g}
|
| 169 |
+
|
| 170 |
+
# ---- streaming embed (the write path): app + ctx + gate per frame -----
|
| 171 |
+
def embed_stream_np(self, frames_u8, batch=128):
|
| 172 |
+
h = self.base.init_state(1)
|
| 173 |
+
apps, ctxs, gates = [], [], []
|
| 174 |
+
for i in range(0, len(frames_u8), batch):
|
| 175 |
+
x = mx.array(frames_u8[i:i + batch].astype(np.float32)
|
| 176 |
+
/ 127.5 - 1.0)[None]
|
| 177 |
+
out = self.run(x, h0=h)
|
| 178 |
+
h = out["h"]
|
| 179 |
+
apps.append(np.array(out["app"][0], np.float32))
|
| 180 |
+
ctxs.append(np.array(out["ctx"][0], np.float32))
|
| 181 |
+
gates.append(np.array(out["gates"][0], np.float32))
|
| 182 |
+
return (np.concatenate(apps), np.concatenate(ctxs),
|
| 183 |
+
np.concatenate(gates))
|
| 184 |
+
|
| 185 |
+
|
| 186 |
+
# ===========================================================================
|
| 187 |
+
# Event segmentation from the gate signal + novelty-weighted pooling
|
| 188 |
+
# ===========================================================================
|
| 189 |
+
def segment_events(gates, ts, min_len=6, smooth=5, thresh_pct=75.0):
|
| 190 |
+
"""Gate-activity peaks -> event boundaries.
|
| 191 |
+
|
| 192 |
+
A boundary is where smoothed gate activity crosses above its own
|
| 193 |
+
percentile threshold after having been below — the scene model is being
|
| 194 |
+
rewritten. Percentile (not absolute) because gate scale is a trained
|
| 195 |
+
quantity; per-stream calibration is free.
|
| 196 |
+
Returns list of (start_idx, end_idx) covering the stream.
|
| 197 |
+
"""
|
| 198 |
+
g = np.convolve(gates, np.ones(smooth) / smooth, mode="same")
|
| 199 |
+
thr = np.percentile(g, thresh_pct)
|
| 200 |
+
above = g > thr
|
| 201 |
+
bounds = [0]
|
| 202 |
+
for i in range(1, len(g)):
|
| 203 |
+
if above[i] and not above[i - 1] and i - bounds[-1] >= min_len:
|
| 204 |
+
bounds.append(i)
|
| 205 |
+
bounds.append(len(g))
|
| 206 |
+
return [(a, b) for a, b in zip(bounds[:-1], bounds[1:]) if b - a >= 2]
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def pool_event(vecs, gates, lo, hi):
|
| 210 |
+
"""Novelty-weighted pool: frames weighted by gate activity, so change
|
| 211 |
+
dominates stillness. Uniform mean is the verb-eraser; this is not."""
|
| 212 |
+
w = gates[lo:hi] + 1e-3
|
| 213 |
+
w = w / w.sum()
|
| 214 |
+
v = (vecs[lo:hi] * w[:, None]).sum(0)
|
| 215 |
+
return v / (np.linalg.norm(v) + 1e-8)
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
# ===========================================================================
|
| 219 |
+
# persistence
|
| 220 |
+
# ===========================================================================
|
| 221 |
+
def save_v2(model, adapter, meta, path):
|
| 222 |
+
from mlx.utils import tree_flatten
|
| 223 |
+
path = Path(path)
|
| 224 |
+
path.mkdir(parents=True, exist_ok=True)
|
| 225 |
+
np.savez(path / "v2.npz",
|
| 226 |
+
**{k: np.array(v) for k, v in tree_flatten(model.parameters())})
|
| 227 |
+
np.savez(path / "adapter.npz",
|
| 228 |
+
**{k: np.array(v)
|
| 229 |
+
for k, v in tree_flatten(adapter.parameters())})
|
| 230 |
+
(path / "v2.json").write_text(json.dumps({**meta, "cfg": model.cfg},
|
| 231 |
+
indent=2))
|
| 232 |
+
|
| 233 |
+
|
| 234 |
+
def load_v2(path):
|
| 235 |
+
from mlx.utils import tree_unflatten
|
| 236 |
+
path = Path(path)
|
| 237 |
+
meta = json.loads((path / "v2.json").read_text())
|
| 238 |
+
base = FDNNVideoEncoder(**meta["cfg"]["base"])
|
| 239 |
+
model = FDNNv2(base, ctx_dim=meta["cfg"]["ctx_dim"],
|
| 240 |
+
horizons=tuple(meta["cfg"]["horizons"]),
|
| 241 |
+
motion_dim=meta["cfg"].get("motion_dim", 64))
|
| 242 |
+
z = np.load(path / "v2.npz")
|
| 243 |
+
model.update(tree_unflatten([(k, mx.array(z[k])) for k in z.files]))
|
| 244 |
+
model.base.cell.freeze(keys=["basis_types", "mask"], recurse=False)
|
| 245 |
+
adapter = TextAdapter()
|
| 246 |
+
z = np.load(path / "adapter.npz")
|
| 247 |
+
adapter.update(tree_unflatten([(k, mx.array(z[k])) for k in z.files]))
|
| 248 |
+
mx.eval(model.parameters(), adapter.parameters())
|
| 249 |
+
return model, adapter, meta
|
python/elidedb/fdnnvideo.py
ADDED
|
@@ -0,0 +1,500 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FDNN-V: a video-native embedding encoder, distilled from SigLIP.
|
| 2 |
+
|
| 3 |
+
WHY THIS MODEL EXISTS
|
| 4 |
+
---------------------
|
| 5 |
+
The database must embed EVERY frame at write time. SigLIP cannot do that at
|
| 6 |
+
scale because it is a 428M-parameter model that treats each frame as an
|
| 7 |
+
unrelated photograph — measured 13-90 ms/frame on this machine. But video is
|
| 8 |
+
not a pile of photographs: frame t is almost entirely explained by frame t-1
|
| 9 |
+
(Deep Feature Flow, arXiv 1611.07715, built a whole recognition system on that
|
| 10 |
+
observation). A human watching a video does not re-parse the scene 30 times a
|
| 11 |
+
second either; they maintain a scene model and update it with what changed.
|
| 12 |
+
|
| 13 |
+
FDNN-V is that shape, made of the three FDNN rules:
|
| 14 |
+
|
| 15 |
+
spatial glimpse (cheap, per frame) "what is in front of me right now"
|
| 16 |
+
Gabor-initialised first conv — V1 simple cells ARE Gabor filters
|
| 17 |
+
small conv pyramid — ventral stream
|
| 18 |
+
|
|
| 19 |
+
FDNN temporal core (rule 1) "what has been going on"
|
| 20 |
+
recurrent state h_t; each channel is a KAN-style sum over k
|
| 21 |
+
heterogeneous sub-functions of (glimpse, state):
|
| 22 |
+
FINER variable-period oscillator — periodic motion (gait, wipers)
|
| 23 |
+
Gabor temporal wavelet — bursts (a grasp, a brake light)
|
| 24 |
+
poly chirp — acceleration (pulling away)
|
| 25 |
+
omega bands partition the temporal spectrum: slow = scene identity,
|
| 26 |
+
mid = object motion, fast = transitions
|
| 27 |
+
|
|
| 28 |
+
head -> SigLIP space (1152-d) so every existing index, text query,
|
| 29 |
+
and centroid keeps working unchanged
|
| 30 |
+
|
| 31 |
+
rules 2+3 (apoptosis -> fine-tune -> neurogenesis -> fine-tune, decided by
|
| 32 |
+
PPO + reverse attention) run post-training in fdnnv_prune — and unlike the
|
| 33 |
+
attempt to prune SigLIP itself, the fine-tune step EXISTS here, because
|
| 34 |
+
distillation pairs are free: every frame of the corpus already has a teacher
|
| 35 |
+
embedding in `frame_vectors`.
|
| 36 |
+
|
| 37 |
+
THE STREAMING CONTRACT
|
| 38 |
+
----------------------
|
| 39 |
+
`step(frame, h) -> (embedding, h')` is causal and O(1) per frame: no
|
| 40 |
+
lookahead, no window buffer. That is what makes embed-on-write real — the
|
| 41 |
+
encoder can sit inside the ingest loop and emit an embedding as each frame
|
| 42 |
+
arrives, like any other index maintenance.
|
| 43 |
+
|
| 44 |
+
WHAT DISTILLATION CAN AND CANNOT GIVE
|
| 45 |
+
-------------------------------------
|
| 46 |
+
The student lands in the teacher's embedding space, so text queries (embedded
|
| 47 |
+
by the frozen SigLIP text tower) keep working. It can match the teacher ON
|
| 48 |
+
THIS CORPUS's manifold; it is not a zero-shot model for arbitrary imagery.
|
| 49 |
+
That is the correct trade for a database: specialise the index to the data it
|
| 50 |
+
serves, keep the teacher for what it is — an offline labeller.
|
| 51 |
+
"""
|
| 52 |
+
from __future__ import annotations
|
| 53 |
+
|
| 54 |
+
import json
|
| 55 |
+
from pathlib import Path
|
| 56 |
+
|
| 57 |
+
import mlx.core as mx
|
| 58 |
+
import mlx.nn as nn
|
| 59 |
+
import numpy as np
|
| 60 |
+
|
| 61 |
+
EMBED_DIM = 1152 # SigLIP so400m space — compatibility is the point
|
| 62 |
+
INPUT_HW = (144, 192) # decode width 192 -> 192x144 (h, w)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
# ===========================================================================
|
| 66 |
+
# V1: Gabor-initialised first convolution
|
| 67 |
+
# ===========================================================================
|
| 68 |
+
def gabor_bank(n_filters, ksize, rng):
|
| 69 |
+
"""Oriented Gabor filters spanning orientation x frequency x phase.
|
| 70 |
+
|
| 71 |
+
Not decoration: the first layer of every competent visual system —
|
| 72 |
+
biological or learned — converges to oriented band-pass filters. Starting
|
| 73 |
+
there instead of at noise removes the epochs a small model would spend
|
| 74 |
+
rediscovering V1.
|
| 75 |
+
"""
|
| 76 |
+
k = np.zeros((n_filters, ksize, ksize, 3), dtype=np.float32)
|
| 77 |
+
half = ksize // 2
|
| 78 |
+
ys, xs = np.mgrid[-half:half + 1, -half:half + 1]
|
| 79 |
+
for i in range(n_filters):
|
| 80 |
+
theta = np.pi * (i % 8) / 8.0
|
| 81 |
+
lam = ksize / (1.5 + (i // 8) % 3)
|
| 82 |
+
psi = 0.0 if (i // 24) % 2 == 0 else np.pi / 2
|
| 83 |
+
sigma = 0.4 * lam
|
| 84 |
+
xr = xs * np.cos(theta) + ys * np.sin(theta)
|
| 85 |
+
yr = -xs * np.sin(theta) + ys * np.cos(theta)
|
| 86 |
+
g = np.exp(-(xr**2 + 0.8 * yr**2) / (2 * sigma**2)) \
|
| 87 |
+
* np.cos(2 * np.pi * xr / lam + psi)
|
| 88 |
+
g -= g.mean()
|
| 89 |
+
g /= (np.abs(g).sum() + 1e-6)
|
| 90 |
+
# colour-opponent weighting: some filters luminance, some R-G, B-Y
|
| 91 |
+
cw = [(1, 1, 1), (1, -1, 0), (0.5, 0.5, -1)][i % 3]
|
| 92 |
+
for c in range(3):
|
| 93 |
+
k[i, :, :, c] = g * cw[c]
|
| 94 |
+
k += rng.normal(0, 0.01, k.shape).astype(np.float32)
|
| 95 |
+
return k
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
# ===========================================================================
|
| 99 |
+
# Ventral stream: small depthwise-separable pyramid
|
| 100 |
+
# ===========================================================================
|
| 101 |
+
class DWBlock(nn.Module):
|
| 102 |
+
def __init__(self, c_in, c_out, stride):
|
| 103 |
+
super().__init__()
|
| 104 |
+
self.dw = nn.Conv2d(c_in, c_in, 3, stride=stride, padding=1,
|
| 105 |
+
groups=c_in)
|
| 106 |
+
self.pw = nn.Conv2d(c_in, c_out, 1)
|
| 107 |
+
self.norm = nn.LayerNorm(c_out)
|
| 108 |
+
self.res = (c_in == c_out and stride == 1)
|
| 109 |
+
|
| 110 |
+
def __call__(self, x):
|
| 111 |
+
y = self.norm(self.pw(nn.silu(self.dw(x))))
|
| 112 |
+
return x + y if self.res else y
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
class ConvStem(nn.Module):
|
| 116 |
+
"""Per-frame spatial encoder -> one glimpse vector. Small on purpose:
|
| 117 |
+
fine discrimination is amortised into the temporal state, and capacity
|
| 118 |
+
here is paid for EVERY frame FOREVER."""
|
| 119 |
+
|
| 120 |
+
def __init__(self, width=48, out_dim=256, seed=0):
|
| 121 |
+
super().__init__()
|
| 122 |
+
rng = np.random.default_rng(seed)
|
| 123 |
+
w = width
|
| 124 |
+
self.v1 = nn.Conv2d(3, w, 7, stride=2, padding=3)
|
| 125 |
+
self.v1.weight = mx.array(gabor_bank(w, 7, rng))
|
| 126 |
+
self.s1 = DWBlock(w, w * 2, 2)
|
| 127 |
+
self.s2 = DWBlock(w * 2, w * 4, 2)
|
| 128 |
+
self.s3 = DWBlock(w * 4, w * 8, 2)
|
| 129 |
+
self.r1 = DWBlock(w * 8, w * 8, 1)
|
| 130 |
+
self.r2 = DWBlock(w * 8, w * 8, 1)
|
| 131 |
+
self.proj = nn.Linear(w * 16, out_dim)
|
| 132 |
+
self.norm = nn.LayerNorm(out_dim)
|
| 133 |
+
|
| 134 |
+
def __call__(self, x): # (B, H, W, 3) in [-1, 1]
|
| 135 |
+
h = nn.silu(self.v1(x))
|
| 136 |
+
h = self.s1(h)
|
| 137 |
+
h = self.s2(h)
|
| 138 |
+
h = self.s3(h)
|
| 139 |
+
h = self.r2(self.r1(h))
|
| 140 |
+
mean = mx.mean(h, axis=(1, 2))
|
| 141 |
+
peak = mx.max(h, axis=(1, 2)) # mean = layout, max = salient
|
| 142 |
+
return self.norm(self.proj(mx.concatenate([mean, peak], axis=-1)))
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
class ViTStem(nn.Module):
|
| 146 |
+
"""The alternative stem: patchify + tiny transformer. Pure matmuls, which
|
| 147 |
+
Apple-silicon GEMM kernels love; raced against ConvStem by measurement,
|
| 148 |
+
never by taste."""
|
| 149 |
+
|
| 150 |
+
def __init__(self, dim=192, depth=4, heads=3, out_dim=256,
|
| 151 |
+
patch=16, hw=INPUT_HW):
|
| 152 |
+
super().__init__()
|
| 153 |
+
self.patch = patch
|
| 154 |
+
self.nh, self.nw = hw[0] // patch, hw[1] // patch
|
| 155 |
+
self.embed = nn.Linear(patch * patch * 3, dim)
|
| 156 |
+
self.pos = mx.zeros((1, self.nh * self.nw, dim))
|
| 157 |
+
self.blocks = [_TinyBlock(dim, heads) for _ in range(depth)]
|
| 158 |
+
self.proj = nn.Linear(dim * 2, out_dim)
|
| 159 |
+
self.norm = nn.LayerNorm(out_dim)
|
| 160 |
+
|
| 161 |
+
def __call__(self, x): # (B, H, W, 3)
|
| 162 |
+
B, H, W, _ = x.shape
|
| 163 |
+
p = self.patch
|
| 164 |
+
t = x.reshape(B, self.nh, p, self.nw, p, 3).transpose(0, 1, 3, 2, 4, 5)
|
| 165 |
+
t = t.reshape(B, self.nh * self.nw, p * p * 3)
|
| 166 |
+
t = self.embed(t) + self.pos
|
| 167 |
+
for blk in self.blocks:
|
| 168 |
+
t = blk(t)
|
| 169 |
+
pooled = mx.concatenate([mx.mean(t, axis=1), mx.max(t, axis=1)],
|
| 170 |
+
axis=-1)
|
| 171 |
+
return self.norm(self.proj(pooled))
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
class _TinyBlock(nn.Module):
|
| 175 |
+
def __init__(self, dim, heads):
|
| 176 |
+
super().__init__()
|
| 177 |
+
self.n1 = nn.LayerNorm(dim)
|
| 178 |
+
self.att = nn.MultiHeadAttention(dim, heads)
|
| 179 |
+
self.n2 = nn.LayerNorm(dim)
|
| 180 |
+
self.fc1 = nn.Linear(dim, dim * 2)
|
| 181 |
+
self.fc2 = nn.Linear(dim * 2, dim)
|
| 182 |
+
|
| 183 |
+
def __call__(self, x):
|
| 184 |
+
y = self.n1(x)
|
| 185 |
+
x = x + self.att(y, y, y)
|
| 186 |
+
return x + self.fc2(nn.silu(self.fc1(self.n2(x))))
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
# ===========================================================================
|
| 190 |
+
# Rule 1: the temporal core — every channel is a sub-network
|
| 191 |
+
# ===========================================================================
|
| 192 |
+
class FDNNTemporalCell(nn.Module):
|
| 193 |
+
"""Recurrent state whose channels are KAN-style sums over k heterogeneous
|
| 194 |
+
temporal sub-functions, gated GRU-style so state persists by default.
|
| 195 |
+
|
| 196 |
+
The bases read the JOINT signal (current glimpse, previous state), so a
|
| 197 |
+
FINER channel can oscillate with repeated motion, a Gabor channel can fire
|
| 198 |
+
on a burst of change, and a poly-phase channel can track acceleration —
|
| 199 |
+
while the gate decides how much of the old scene model each step is
|
| 200 |
+
allowed to overwrite. This is FDNN's HybridBiomimeticLayer with time as
|
| 201 |
+
the signal axis instead of a coordinate.
|
| 202 |
+
"""
|
| 203 |
+
|
| 204 |
+
def __init__(self, in_dim=256, channels=256, k_width=4,
|
| 205 |
+
omega_bands=(0.8, 2.5, 8.0),
|
| 206 |
+
band_fractions=(0.34, 0.33, 0.33), seed=0):
|
| 207 |
+
# omega bands are LOWER than the context tower's (2, 6, 18) on
|
| 208 |
+
# purpose: that tower convolved over time feed-forward, this cell
|
| 209 |
+
# FEEDS BACK. sin bases at omega 15 inside a 32-step recurrence give
|
| 210 |
+
# chaotic gradients — measured: stage-2 loss climbed from 0.086 to
|
| 211 |
+
# 0.112 and val fidelity fell 0.019. Same bases, calmer spectrum.
|
| 212 |
+
super().__init__()
|
| 213 |
+
rng = np.random.default_rng(seed)
|
| 214 |
+
C, k = channels, k_width
|
| 215 |
+
self.C, self.k = C, k
|
| 216 |
+
|
| 217 |
+
omegas = []
|
| 218 |
+
for om, fr in zip(omega_bands, band_fractions):
|
| 219 |
+
omegas.extend([om] * int(round(fr * C)))
|
| 220 |
+
omegas = (omegas + [omega_bands[-1]] * C)[:C]
|
| 221 |
+
self.omegas_per_neuron = np.array(omegas, dtype=np.float32)
|
| 222 |
+
self.omegas = mx.array(np.repeat(self.omegas_per_neuron, k))
|
| 223 |
+
|
| 224 |
+
half, quarter = max(k // 2, 1), max(k // 4, 1)
|
| 225 |
+
per = np.array([0] * half + [1] * quarter
|
| 226 |
+
+ [3] * max(k - half - quarter, 0), np.int32)[:k]
|
| 227 |
+
self.basis_types = mx.array(np.tile(per, C).astype(np.int32))
|
| 228 |
+
# A categorical selector, not a weight: unfrozen, the optimizer
|
| 229 |
+
# promotes it to float and drifts it off its exact values, silently
|
| 230 |
+
# rerouting every Gabor neuron to the poly branch (measured on the
|
| 231 |
+
# context tower: 1.0 -> 0.9992 and `== 1` matched nothing).
|
| 232 |
+
self.freeze(keys=["basis_types"], recurse=False)
|
| 233 |
+
|
| 234 |
+
mean_om = float(self.omegas_per_neuron.mean())
|
| 235 |
+
lim_g = float(np.sqrt(6.0 / in_dim) / mean_om)
|
| 236 |
+
lim_h = float(np.sqrt(6.0 / C) / mean_om)
|
| 237 |
+
self.Wg = mx.array(rng.uniform(-lim_g, lim_g,
|
| 238 |
+
(in_dim, C * k)).astype(np.float32))
|
| 239 |
+
self.Wh = mx.array(rng.uniform(-lim_h, lim_h,
|
| 240 |
+
(C, C * k)).astype(np.float32))
|
| 241 |
+
self.b1 = mx.array(rng.uniform(-2.0, 2.0, (C * k,)).astype(np.float32))
|
| 242 |
+
self.phases = mx.array(rng.uniform(0, 2 * np.pi,
|
| 243 |
+
(C * k,)).astype(np.float32))
|
| 244 |
+
self.gabor_s = mx.array(rng.uniform(0.3, 1.5,
|
| 245 |
+
(C * k,)).astype(np.float32))
|
| 246 |
+
log_om = np.log(np.clip(np.repeat(self.omegas_per_neuron, k),
|
| 247 |
+
1e-3, None))
|
| 248 |
+
# alpha starts log-uniform between 1 and omega. With sub-unit omegas
|
| 249 |
+
# (the calmed recurrent bands) log-omega is negative, so the interval
|
| 250 |
+
# must be ordered explicitly — uniform(0, negative) is an error.
|
| 251 |
+
self.log_alpha = mx.array(rng.uniform(np.minimum(0.0, log_om),
|
| 252 |
+
np.maximum(0.0, log_om) + 1e-6
|
| 253 |
+
).astype(np.float32))
|
| 254 |
+
w2s = float(np.sqrt(6.0 / (C * k)))
|
| 255 |
+
self.w2 = mx.array(rng.uniform(-w2s, w2s, (C, k)).astype(np.float32))
|
| 256 |
+
|
| 257 |
+
zlim = float(np.sqrt(6.0 / (in_dim + C)))
|
| 258 |
+
self.Wz = mx.array(rng.uniform(-zlim, zlim,
|
| 259 |
+
(in_dim, C)).astype(np.float32))
|
| 260 |
+
self.Uz = mx.array(rng.uniform(-zlim, zlim,
|
| 261 |
+
(C, C)).astype(np.float32))
|
| 262 |
+
# Gate bias starts NEGATIVE: sigmoid(-1) ~ 0.27, so at init the state
|
| 263 |
+
# persists — a scene model that forgets everything every frame is just
|
| 264 |
+
# a per-frame model with extra steps.
|
| 265 |
+
self.bz = mx.array(np.full((C,), -1.0, np.float32))
|
| 266 |
+
|
| 267 |
+
self.mask = mx.array(np.ones((C,), np.float32)) # aliveness (rule 2)
|
| 268 |
+
# Aliveness is set by the pruning cycle, never by the optimizer —
|
| 269 |
+
# unfrozen, AdamW weight-decays it off 1.0 and every channel quietly
|
| 270 |
+
# shrinks (the measured context-tower failure mode).
|
| 271 |
+
self.freeze(keys=["mask"], recurse=False)
|
| 272 |
+
|
| 273 |
+
def set_active_mask(self, m):
|
| 274 |
+
self.mask = mx.array(np.asarray(m, dtype=np.float32))
|
| 275 |
+
|
| 276 |
+
def _cand(self, g, h):
|
| 277 |
+
pre = g @ self.Wg + h @ self.Wh + self.b1
|
| 278 |
+
om_h = self.omegas * pre
|
| 279 |
+
sq = pre * pre
|
| 280 |
+
alpha = mx.exp(self.log_alpha)
|
| 281 |
+
finer = mx.sin(self.omegas * (mx.abs(pre) + 1.0) * pre + self.phases)
|
| 282 |
+
gab = mx.exp(-(self.gabor_s ** 2) * sq) * mx.sin(om_h + self.phases)
|
| 283 |
+
sine = mx.sin(om_h + self.phases)
|
| 284 |
+
poly = mx.sin(alpha * sq + om_h + self.phases)
|
| 285 |
+
acts = mx.where(self.basis_types == 0, finer,
|
| 286 |
+
mx.where(self.basis_types == 1, gab,
|
| 287 |
+
mx.where(self.basis_types == 2, sine, poly)))
|
| 288 |
+
acts = acts.reshape(-1, self.C, self.k)
|
| 289 |
+
return mx.sum(acts * self.w2, axis=-1)
|
| 290 |
+
|
| 291 |
+
def neuron_outputs(self, g, h):
|
| 292 |
+
"""Per-neuron candidate BEFORE gate and mask — the pruning signal."""
|
| 293 |
+
return self._cand(g, h)
|
| 294 |
+
|
| 295 |
+
def __call__(self, g, h):
|
| 296 |
+
cand = self._cand(g, h) * self.mask
|
| 297 |
+
z = mx.sigmoid(g @ self.Wz + h @ self.Uz + self.bz) * self.mask
|
| 298 |
+
return (1.0 - z) * h + z * cand
|
| 299 |
+
|
| 300 |
+
|
| 301 |
+
# ===========================================================================
|
| 302 |
+
# The encoder
|
| 303 |
+
# ===========================================================================
|
| 304 |
+
class FDNNVideoEncoder(nn.Module):
|
| 305 |
+
def __init__(self, stem="conv", stem_width=48, glimpse=256, channels=256,
|
| 306 |
+
k_width=4, embed_dim=EMBED_DIM, vit_depth=4, seed=0):
|
| 307 |
+
super().__init__()
|
| 308 |
+
self.cfg = dict(stem=stem, stem_width=stem_width, glimpse=glimpse,
|
| 309 |
+
channels=channels, k_width=k_width,
|
| 310 |
+
embed_dim=embed_dim, vit_depth=vit_depth, seed=seed)
|
| 311 |
+
if stem == "conv":
|
| 312 |
+
self.stem = ConvStem(width=stem_width, out_dim=glimpse, seed=seed)
|
| 313 |
+
else:
|
| 314 |
+
self.stem = ViTStem(dim=stem_width * 4, depth=vit_depth,
|
| 315 |
+
out_dim=glimpse)
|
| 316 |
+
self.cell = FDNNTemporalCell(in_dim=glimpse, channels=channels,
|
| 317 |
+
k_width=k_width, seed=seed)
|
| 318 |
+
self.head_g = nn.Linear(glimpse, embed_dim)
|
| 319 |
+
# Temporal head starts at zero: at init the model IS the per-frame
|
| 320 |
+
# model (stage 1), and training can only add information from state.
|
| 321 |
+
# Same identity-safe discipline as every other init in this repo.
|
| 322 |
+
self.head_h = nn.Linear(channels, embed_dim)
|
| 323 |
+
self.head_h.weight = mx.zeros(self.head_h.weight.shape)
|
| 324 |
+
self.head_h.bias = mx.zeros((embed_dim,))
|
| 325 |
+
self.channels = channels
|
| 326 |
+
|
| 327 |
+
# ---- streaming: this is the embed-on-write contract -------------------
|
| 328 |
+
def init_state(self, batch=1):
|
| 329 |
+
return mx.zeros((batch, self.channels))
|
| 330 |
+
|
| 331 |
+
def step(self, frame, h):
|
| 332 |
+
"""One frame in, one embedding out, O(1) state carried. Causal."""
|
| 333 |
+
g = self.stem(frame)
|
| 334 |
+
h = self.cell(g, h)
|
| 335 |
+
e = self.head_g(g) + self.head_h(h)
|
| 336 |
+
return e * mx.rsqrt(mx.sum(e * e, axis=-1, keepdims=True) + 1e-8), h
|
| 337 |
+
|
| 338 |
+
# ---- batched sequences (training / bulk ingest) -----------------------
|
| 339 |
+
def __call__(self, seq, h0=None):
|
| 340 |
+
"""(B, T, H, W, 3) -> (B, T, D). Stem runs on all frames as one big
|
| 341 |
+
batch (the GEMM-friendly part); only the tiny cell recurs."""
|
| 342 |
+
B, T = seq.shape[0], seq.shape[1]
|
| 343 |
+
g = self.stem(seq.reshape(B * T, *seq.shape[2:])).reshape(B, T, -1)
|
| 344 |
+
h = self.init_state(B) if h0 is None else h0
|
| 345 |
+
outs = []
|
| 346 |
+
for t in range(T):
|
| 347 |
+
h = self.cell(g[:, t], h)
|
| 348 |
+
outs.append(h)
|
| 349 |
+
hs = mx.stack(outs, axis=1)
|
| 350 |
+
e = self.head_g(g) + self.head_h(hs)
|
| 351 |
+
return e * mx.rsqrt(mx.sum(e * e, axis=-1, keepdims=True) + 1e-8), h
|
| 352 |
+
|
| 353 |
+
def embed_frames_np(self, frames_u8, batch=64, chunk=None):
|
| 354 |
+
"""uint8 (N, H, W, 3) of ONE stream, in time order -> (N, 1152).
|
| 355 |
+
Stateful across batches — one continuous pass over the stream."""
|
| 356 |
+
h = self.init_state(1)
|
| 357 |
+
out = []
|
| 358 |
+
for i in range(0, len(frames_u8), batch):
|
| 359 |
+
x = mx.array(frames_u8[i:i + batch].astype(np.float32)
|
| 360 |
+
/ 127.5 - 1.0)[None]
|
| 361 |
+
e, h = self(x, h0=h)
|
| 362 |
+
out.append(np.array(e[0], dtype=np.float32))
|
| 363 |
+
return np.concatenate(out, axis=0)
|
| 364 |
+
|
| 365 |
+
|
| 366 |
+
# ===========================================================================
|
| 367 |
+
# Distillation loss: pointwise + affinity mimicking
|
| 368 |
+
# ===========================================================================
|
| 369 |
+
def distill_loss(student, teacher, affinity_w=0.25, mu=None, centered_w=1.0,
|
| 370 |
+
anchors=None, anchor_w=50.0):
|
| 371 |
+
"""Distillation aimed at RETRIEVAL, not at raw closeness.
|
| 372 |
+
|
| 373 |
+
Plain pointwise cosine is a trap on a homogeneous corpus: every teacher
|
| 374 |
+
vector shares a huge common mode, so matching that alone buys ~0.9 cosine
|
| 375 |
+
while scrambling the thin discriminative residual that ranking runs on.
|
| 376 |
+
Measured: a student at fidelity 0.907 kept only 4.4% of the teacher's
|
| 377 |
+
top-10 neighbours, while the teacher AGAINST ITSELF at a different input
|
| 378 |
+
resolution — fidelity 0.921 — keeps 50.7%. Same closeness, 10x the
|
| 379 |
+
retrieval agreement: the difference is WHERE the error lives.
|
| 380 |
+
|
| 381 |
+
So three additional terms put the error where it does no harm:
|
| 382 |
+
centered cosine on (v - mu): the mean-free residual is exactly what
|
| 383 |
+
ranking compares, so it gets its own gradient.
|
| 384 |
+
affinity within-batch similarity matching (TinyCLIP, arXiv
|
| 385 |
+
2309.12314): preserve the teacher's ordering structure.
|
| 386 |
+
anchors similarity profile against real caption-text embeddings from
|
| 387 |
+
this store: text queries live in those directions, and
|
| 388 |
+
image-text sims occupy a band ~50x narrower than image-image
|
| 389 |
+
sims — hence the weight.
|
| 390 |
+
"""
|
| 391 |
+
t = teacher * mx.rsqrt(mx.sum(teacher * teacher, axis=-1,
|
| 392 |
+
keepdims=True) + 1e-8)
|
| 393 |
+
loss = mx.mean(1.0 - mx.sum(student * t, axis=-1))
|
| 394 |
+
if mu is not None and centered_w > 0:
|
| 395 |
+
sc = student - mu
|
| 396 |
+
tc = t - mu
|
| 397 |
+
sc = sc * mx.rsqrt(mx.sum(sc * sc, axis=-1, keepdims=True) + 1e-8)
|
| 398 |
+
tc = tc * mx.rsqrt(mx.sum(tc * tc, axis=-1, keepdims=True) + 1e-8)
|
| 399 |
+
loss = loss + centered_w * mx.mean(1.0 - mx.sum(sc * tc, axis=-1))
|
| 400 |
+
if affinity_w > 0:
|
| 401 |
+
s2 = student.reshape(-1, student.shape[-1])
|
| 402 |
+
t2 = t.reshape(-1, t.shape[-1])
|
| 403 |
+
loss = loss + affinity_w * mx.mean(mx.square(s2 @ s2.T - t2 @ t2.T))
|
| 404 |
+
if anchors is not None and anchor_w > 0:
|
| 405 |
+
sa = student.reshape(-1, student.shape[-1]) @ anchors.T
|
| 406 |
+
ta = t.reshape(-1, t.shape[-1]) @ anchors.T
|
| 407 |
+
loss = loss + anchor_w * mx.mean(mx.square(sa - ta))
|
| 408 |
+
return loss
|
| 409 |
+
|
| 410 |
+
|
| 411 |
+
# ===========================================================================
|
| 412 |
+
# persistence
|
| 413 |
+
# ===========================================================================
|
| 414 |
+
def save_encoder(model, meta, path):
|
| 415 |
+
from mlx.utils import tree_flatten
|
| 416 |
+
path = Path(path)
|
| 417 |
+
path.mkdir(parents=True, exist_ok=True)
|
| 418 |
+
np.savez(path / "weights.npz",
|
| 419 |
+
**{k: np.array(v) for k, v in tree_flatten(model.parameters())})
|
| 420 |
+
(path / "encoder.json").write_text(json.dumps(
|
| 421 |
+
{**meta, "cfg": model.cfg}, indent=2))
|
| 422 |
+
|
| 423 |
+
|
| 424 |
+
def fdnnv_dir() -> Path:
|
| 425 |
+
"""Where the FDNN-V encoder lives.
|
| 426 |
+
|
| 427 |
+
Repo-level `models/fdnnv`, NOT inside any store. The encoder used to
|
| 428 |
+
sit at `lake/bridge/models/fdnnv`, so clearing the stores deleted a
|
| 429 |
+
trained model along with the data (2026-07-28). A model is not store
|
| 430 |
+
data. The legacy path is still accepted for stores that predate the
|
| 431 |
+
move."""
|
| 432 |
+
here = Path("models/fdnnv")
|
| 433 |
+
if (here / "encoder.json").exists():
|
| 434 |
+
return here
|
| 435 |
+
legacy = Path("lake/bridge/models/fdnnv")
|
| 436 |
+
if (legacy / "encoder.json").exists():
|
| 437 |
+
return legacy
|
| 438 |
+
return here
|
| 439 |
+
|
| 440 |
+
|
| 441 |
+
def load_encoder(path):
|
| 442 |
+
from mlx.utils import tree_unflatten
|
| 443 |
+
path = Path(path)
|
| 444 |
+
meta = json.loads((path / "encoder.json").read_text())
|
| 445 |
+
model = FDNNVideoEncoder(**meta["cfg"])
|
| 446 |
+
z = np.load(path / "weights.npz")
|
| 447 |
+
model.update(tree_unflatten([(k, mx.array(z[k])) for k in z.files]))
|
| 448 |
+
model.cell.freeze(keys=["basis_types", "mask"], recurse=False)
|
| 449 |
+
mx.eval(model.parameters())
|
| 450 |
+
return model, meta
|
| 451 |
+
|
| 452 |
+
|
| 453 |
+
# ===========================================================================
|
| 454 |
+
# The write path: chunked byte-range decode feeding the streaming encoder
|
| 455 |
+
# ===========================================================================
|
| 456 |
+
def embed_stream(store, model, rows, width=192, chunk=512, batch_cb=None):
|
| 457 |
+
"""Embed one stream's frames in time order, state carried across chunks.
|
| 458 |
+
|
| 459 |
+
`chunk` bounds the decoder subprocess's rawvideo buffer (~1 GB at 512
|
| 460 |
+
frames of 640x480); the encoder state flows straight through, so the
|
| 461 |
+
result is identical to one infinite pass. This loop is the write path:
|
| 462 |
+
ingest can call it as frames land.
|
| 463 |
+
Returns (ts int64 array, vectors float32 (N, D), decode_s, embed_s).
|
| 464 |
+
"""
|
| 465 |
+
import time as _time
|
| 466 |
+
|
| 467 |
+
from .video import FrameSet
|
| 468 |
+
ts_out, vecs = [], []
|
| 469 |
+
h = model.init_state(1)
|
| 470 |
+
dec_s = emb_s = 0.0
|
| 471 |
+
for i in range(0, len(rows), chunk):
|
| 472 |
+
t0 = _time.perf_counter()
|
| 473 |
+
dec = FrameSet(store, "frames", rows.slice(i, chunk)).decode(
|
| 474 |
+
width=width)
|
| 475 |
+
dec_s += _time.perf_counter() - t0
|
| 476 |
+
if not dec:
|
| 477 |
+
continue
|
| 478 |
+
frames = np.stack([d[1] for d in dec])
|
| 479 |
+
# the stem takes EXACTLY (in_h, in_w); sources with a different
|
| 480 |
+
# aspect ratio decode to other shapes (lab video came back square
|
| 481 |
+
# and crashed the reshape). Stretch — the encoder was distilled on
|
| 482 |
+
# stretched frames, so aspect distortion is in-distribution.
|
| 483 |
+
ih, iw = model.cfg["in_hw"] if "in_hw" in model.cfg else (144, 192)
|
| 484 |
+
if frames.shape[1] != ih or frames.shape[2] != iw:
|
| 485 |
+
xr = np.linspace(0, frames.shape[2] - 1, iw).round().astype(int)
|
| 486 |
+
yr = np.linspace(0, frames.shape[1] - 1, ih).round().astype(int)
|
| 487 |
+
frames = frames[:, yr][:, :, xr]
|
| 488 |
+
t0 = _time.perf_counter()
|
| 489 |
+
x = mx.array(frames.astype(np.float32) / 127.5 - 1.0)[None]
|
| 490 |
+
e, h = model(x, h0=h)
|
| 491 |
+
e = np.array(e[0], dtype=np.float32)
|
| 492 |
+
emb_s += _time.perf_counter() - t0
|
| 493 |
+
ts_out.extend(d[0] for d in dec)
|
| 494 |
+
vecs.append(e)
|
| 495 |
+
if batch_cb:
|
| 496 |
+
batch_cb(len(ts_out))
|
| 497 |
+
if not vecs:
|
| 498 |
+
return np.array([], np.int64), np.zeros((0, EMBED_DIM), np.float32), \
|
| 499 |
+
dec_s, emb_s
|
| 500 |
+
return (np.array(ts_out, np.int64), np.concatenate(vecs), dec_s, emb_s)
|
python/elidedb/fftools.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Locate the ffmpeg/ffprobe binaries robustly.
|
| 2 |
+
|
| 3 |
+
Why this exists: a macOS .app (or anything launched from Finder/launchd) runs
|
| 4 |
+
with a minimal PATH — typically `/usr/bin:/bin:/usr/sbin:/sbin` — which does
|
| 5 |
+
NOT include Homebrew's `/opt/homebrew/bin`. A bare `["ffmpeg", ...]` then
|
| 6 |
+
fails with `FileNotFoundError` even though ffmpeg is installed. We resolve the
|
| 7 |
+
absolute path once, checking (in order): an explicit env override, the current
|
| 8 |
+
PATH, and the common install locations on macOS/Linux.
|
| 9 |
+
"""
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import functools
|
| 13 |
+
import os
|
| 14 |
+
import shutil
|
| 15 |
+
|
| 16 |
+
_COMMON_DIRS = [
|
| 17 |
+
"/opt/homebrew/bin", # Apple Silicon Homebrew
|
| 18 |
+
"/usr/local/bin", # Intel Homebrew / manual installs
|
| 19 |
+
"/opt/local/bin", # MacPorts
|
| 20 |
+
"/usr/bin",
|
| 21 |
+
]
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
@functools.lru_cache(maxsize=None)
|
| 25 |
+
def find(tool: str) -> str:
|
| 26 |
+
"""Absolute path to `ffmpeg`/`ffprobe`, or raise with a fixable message."""
|
| 27 |
+
env = os.environ.get(f"ELIDEDB_{tool.upper()}")
|
| 28 |
+
if env and os.path.exists(env):
|
| 29 |
+
return env
|
| 30 |
+
found = shutil.which(tool)
|
| 31 |
+
if found:
|
| 32 |
+
return found
|
| 33 |
+
for d in _COMMON_DIRS:
|
| 34 |
+
cand = os.path.join(d, tool)
|
| 35 |
+
if os.path.exists(cand):
|
| 36 |
+
return cand
|
| 37 |
+
raise FileNotFoundError(
|
| 38 |
+
f"{tool} not found. Install it with `brew install ffmpeg`, or set "
|
| 39 |
+
f"ELIDEDB_{tool.upper()}=/path/to/{tool}. (Apps launched from Finder "
|
| 40 |
+
f"get a minimal PATH, so Homebrew's bin dir may not be visible.)")
|
python/elidedb/fusion.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Reciprocal rank fusion.
|
| 2 |
+
|
| 3 |
+
WHY NOT JUST ADD THE SCORES
|
| 4 |
+
---------------------------
|
| 5 |
+
Because they are not on the same scale, and no amount of tuning fixes that.
|
| 6 |
+
SigLIP image-text cosines on a homogeneous corpus occupy roughly 0.01-0.15
|
| 7 |
+
(the modality gap squashes them); caption-LSA cosines spread over most of
|
| 8 |
+
[-1, 1]; TF-IDF cosines are mostly 0 with a thin tail near 1. A weighted sum
|
| 9 |
+
of those is governed by whichever signal happens to have the widest spread on
|
| 10 |
+
that particular query, so the weight you set is not the weight you get.
|
| 11 |
+
|
| 12 |
+
Standardising per query (z-scores) helps but still lets one ranker's outlier
|
| 13 |
+
drag the result: a single candidate three sigma out on appearance outranks a
|
| 14 |
+
candidate that every ranker agrees is second-best.
|
| 15 |
+
|
| 16 |
+
RRF (Cormack, Clarke & Buettcher, SIGIR 2009) throws the magnitudes away and
|
| 17 |
+
keeps only the ORDER:
|
| 18 |
+
|
| 19 |
+
score(d) = sum over rankers r of w_r / (K + rank_r(d))
|
| 20 |
+
|
| 21 |
+
That makes fusion scale-free by construction, and it rewards CONSENSUS: a
|
| 22 |
+
document ranked 2nd by three rankers beats one ranked 1st by a single ranker
|
| 23 |
+
and 200th by the rest. Which is exactly the failure being fixed — "crossing
|
| 24 |
+
red car" returned any clip of a person crossing, because one strong signal on
|
| 25 |
+
"crossing" was allowed to win alone. Under RRF the answer has to look right to
|
| 26 |
+
the appearance index AND the caption index AND the lexical index.
|
| 27 |
+
|
| 28 |
+
K (default 60, the constant from the paper) damps the top of the curve: it is
|
| 29 |
+
the number of rank positions over which differences stop mattering much, so a
|
| 30 |
+
ranker cannot dominate purely by being extremely confident about its #1.
|
| 31 |
+
"""
|
| 32 |
+
from __future__ import annotations
|
| 33 |
+
|
| 34 |
+
import numpy as np
|
| 35 |
+
|
| 36 |
+
DEFAULT_K = 60.0
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def ranks_from_scores(scores: np.ndarray) -> np.ndarray:
|
| 40 |
+
"""0-based ranks; highest score gets rank 0.
|
| 41 |
+
|
| 42 |
+
NaN means ABSTAIN, and an abstention is given the median rank rather than
|
| 43 |
+
the last. This matters: the lexical ranker can only score windows that
|
| 44 |
+
actually have a caption, and windows whose context vector was estimated by
|
| 45 |
+
the tower have none. Sending those to the bottom would let a ranker veto
|
| 46 |
+
every item it has no opinion about, which is the opposite of what "no
|
| 47 |
+
evidence" should mean. Median rank is the neutral position — it neither
|
| 48 |
+
helps nor hurts.
|
| 49 |
+
"""
|
| 50 |
+
s = np.asarray(scores, dtype=float)
|
| 51 |
+
known = ~np.isnan(s)
|
| 52 |
+
r = np.empty(len(s), dtype=float)
|
| 53 |
+
if not known.any():
|
| 54 |
+
r[:] = 0.0
|
| 55 |
+
return r
|
| 56 |
+
sub = s[known]
|
| 57 |
+
order = np.argsort(-sub, kind="stable")
|
| 58 |
+
rr = np.empty(len(sub), dtype=float)
|
| 59 |
+
rr[order] = np.arange(len(sub), dtype=float)
|
| 60 |
+
r[known] = rr
|
| 61 |
+
r[~known] = float(np.median(rr))
|
| 62 |
+
return r
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def rrf(rankings: dict[str, np.ndarray], weights: dict[str, float] | None = None,
|
| 66 |
+
k: float = DEFAULT_K) -> np.ndarray:
|
| 67 |
+
"""Fuse score vectors (all over the SAME candidate set, same order).
|
| 68 |
+
|
| 69 |
+
Pass raw scores, not ranks — this converts. Returns the fused score, where
|
| 70 |
+
higher is better, on an arbitrary but consistent scale.
|
| 71 |
+
"""
|
| 72 |
+
if not rankings:
|
| 73 |
+
raise ValueError("rrf() needs at least one ranking")
|
| 74 |
+
n = len(next(iter(rankings.values())))
|
| 75 |
+
out = np.zeros(n, dtype=float)
|
| 76 |
+
for name, sc in rankings.items():
|
| 77 |
+
sc = np.asarray(sc, dtype=float)
|
| 78 |
+
if len(sc) != n:
|
| 79 |
+
raise ValueError(
|
| 80 |
+
f"ranker '{name}' has {len(sc)} scores, expected {n} — every "
|
| 81 |
+
"ranker must score the same candidate set")
|
| 82 |
+
w = 1.0 if weights is None else float(weights.get(name, 1.0))
|
| 83 |
+
if w == 0.0:
|
| 84 |
+
continue
|
| 85 |
+
out += w / (k + ranks_from_scores(sc) + 1.0)
|
| 86 |
+
return out
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def variant_max(vs):
|
| 90 |
+
"""Row-wise max over query-variant score arrays. All-NaN rows
|
| 91 |
+
(episodes absent from a channel's table — they abstain, they are
|
| 92 |
+
not errors) stay NaN without numpy's All-NaN-slice RuntimeWarning,
|
| 93 |
+
which was leaking source lines into bench stdout."""
|
| 94 |
+
m = np.stack(vs)
|
| 95 |
+
out = np.full(m.shape[1], np.nan)
|
| 96 |
+
fin = np.isfinite(m).any(0)
|
| 97 |
+
if fin.any():
|
| 98 |
+
out[fin] = np.nanmax(m[:, fin], 0)
|
| 99 |
+
return out
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def explain_fusion(rankings, weights=None, k=DEFAULT_K, idx=None, top=5):
|
| 103 |
+
"""Per-ranker rank of each fused winner. Answers 'why is this here?' —
|
| 104 |
+
which is the question a fused score alone can never answer."""
|
| 105 |
+
fused = rrf(rankings, weights, k)
|
| 106 |
+
order = np.argsort(-fused)[:top] if idx is None else idx
|
| 107 |
+
rk = {n: ranks_from_scores(s) for n, s in rankings.items()}
|
| 108 |
+
return [{"i": int(i), "fused": float(fused[i]),
|
| 109 |
+
"ranks": {n: int(r[i]) + 1 for n, r in rk.items()}}
|
| 110 |
+
for i in order]
|
python/elidedb/grounding.py
ADDED
|
@@ -0,0 +1,218 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""GEOMETRIC RELATIONAL VERIFIER — boxes over time instead of a VLM.
|
| 2 |
+
|
| 3 |
+
Calibration verdict that forced this: VLM judges see the right scene
|
| 4 |
+
and say yes (7B est 1.0 on visually ~10%-pure sets); they cannot bind
|
| 5 |
+
"X ends up inside/on Y". A detector can: ground the query's noun
|
| 6 |
+
phrases to boxes on sampled frames, and the relation IS the geometry —
|
| 7 |
+
containment (intersection-over-area of X in Y) low at the start and
|
| 8 |
+
high at the end verifies "put X in/on Y"; the reverse verifies "take X
|
| 9 |
+
out of Y". Deterministic, explainable, video-evidence-only.
|
| 10 |
+
|
| 11 |
+
Interim detector: Grounding DINO (open-vocabulary, ungated). Measured
|
| 12 |
+
on our frames (2026-07-24): green-object→drawer and lid→pot both
|
| 13 |
+
verified by end-state containment; one mid-flight false positive under
|
| 14 |
+
occlusion (sink read as lid) — which start/end predicates never see.
|
| 15 |
+
SAM 3 (pending license approval) swaps in for masks + tracking later.
|
| 16 |
+
|
| 17 |
+
Latency: ~1.2 s/frame on MPS — a VERIFICATION-tier cost (top-K
|
| 18 |
+
candidates, 6 frames each), never a scan cost.
|
| 19 |
+
"""
|
| 20 |
+
from __future__ import annotations
|
| 21 |
+
|
| 22 |
+
import numpy as np
|
| 23 |
+
|
| 24 |
+
_G = {}
|
| 25 |
+
|
| 26 |
+
MODEL_ID = "IDEA-Research/grounding-dino-base"
|
| 27 |
+
|
| 28 |
+
# "put X in Y" family vs "take X out of Y" family: sign of the expected
|
| 29 |
+
# containment CHANGE. Generic English, nothing per-dataset.
|
| 30 |
+
_INWARD = ("in", "into", "inside", "on", "onto", "on top of", "over")
|
| 31 |
+
_OUTWARD = ("out", "out of", "from")
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def _load():
|
| 35 |
+
if "model" in _G:
|
| 36 |
+
return _G
|
| 37 |
+
import torch
|
| 38 |
+
from transformers import (AutoModelForZeroShotObjectDetection,
|
| 39 |
+
AutoProcessor)
|
| 40 |
+
dev = "mps" if torch.backends.mps.is_available() else "cpu"
|
| 41 |
+
_G["proc"] = AutoProcessor.from_pretrained(MODEL_ID)
|
| 42 |
+
_G["model"] = AutoModelForZeroShotObjectDetection.from_pretrained(
|
| 43 |
+
MODEL_ID).to(dev).eval()
|
| 44 |
+
_G["dev"] = dev
|
| 45 |
+
return _G
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def detect_phrases(images, phrases, threshold=0.3):
|
| 49 |
+
"""PIL images x phrases -> per image {phrase: (box, score)|None},
|
| 50 |
+
best box per phrase."""
|
| 51 |
+
import torch
|
| 52 |
+
g = _load()
|
| 53 |
+
prompt = " . ".join(p.lower().strip(". ") for p in phrases) + " ."
|
| 54 |
+
inputs = g["proc"](images=images, text=[prompt] * len(images),
|
| 55 |
+
return_tensors="pt").to(g["dev"])
|
| 56 |
+
with torch.no_grad():
|
| 57 |
+
out = g["model"](**inputs)
|
| 58 |
+
res = g["proc"].post_process_grounded_object_detection(
|
| 59 |
+
out, inputs.input_ids, threshold=threshold, text_threshold=0.25,
|
| 60 |
+
target_sizes=[im.size[::-1] for im in images])
|
| 61 |
+
per = []
|
| 62 |
+
for r in res:
|
| 63 |
+
best = {}
|
| 64 |
+
for box, sc, lb in zip(r["boxes"], r["scores"],
|
| 65 |
+
r["text_labels"]):
|
| 66 |
+
for p in phrases:
|
| 67 |
+
if p.lower().strip(". ") in lb or lb in p.lower():
|
| 68 |
+
if p not in best or sc > best[p][1]:
|
| 69 |
+
best[p] = (np.array([float(v) for v in box]),
|
| 70 |
+
float(sc))
|
| 71 |
+
per.append({p: best.get(p) for p in phrases})
|
| 72 |
+
return per
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def detect_regions(images, prompt="object", threshold=0.25):
|
| 76 |
+
"""ALL boxes above threshold, per image — a class-agnostic REGION
|
| 77 |
+
PROPOSER rather than a phrase verifier.
|
| 78 |
+
|
| 79 |
+
detect_phrases keeps the single best box per phrase, which is right
|
| 80 |
+
when you are asking "where is the lid" and wrong when you are asking
|
| 81 |
+
"where are the objects": it returns at most one region per frame and
|
| 82 |
+
an object store needs every region in the frame.
|
| 83 |
+
"""
|
| 84 |
+
import torch
|
| 85 |
+
g = _load()
|
| 86 |
+
text = prompt.lower().strip(". ") + " ."
|
| 87 |
+
inputs = g["proc"](images=images, text=[text] * len(images),
|
| 88 |
+
return_tensors="pt").to(g["dev"])
|
| 89 |
+
with torch.no_grad():
|
| 90 |
+
out = g["model"](**inputs)
|
| 91 |
+
res = g["proc"].post_process_grounded_object_detection(
|
| 92 |
+
out, inputs.input_ids, threshold=threshold, text_threshold=0.2,
|
| 93 |
+
target_sizes=[im.size[::-1] for im in images])
|
| 94 |
+
per = []
|
| 95 |
+
for r in res:
|
| 96 |
+
boxes = [([float(v) for v in b], float(sc))
|
| 97 |
+
for b, sc in zip(r["boxes"], r["scores"])]
|
| 98 |
+
boxes.sort(key=lambda x: -x[1])
|
| 99 |
+
per.append(boxes)
|
| 100 |
+
return per
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def _ioa(a, b):
|
| 104 |
+
"""Intersection over area of A — how much of X sits inside Y."""
|
| 105 |
+
x0 = max(a[0], b[0]); y0 = max(a[1], b[1])
|
| 106 |
+
x1 = min(a[2], b[2]); y1 = min(a[3], b[3])
|
| 107 |
+
inter = max(0.0, x1 - x0) * max(0.0, y1 - y0)
|
| 108 |
+
area = max(1.0, (a[2] - a[0]) * (a[3] - a[1]))
|
| 109 |
+
return inter / area
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def parse_relation(text):
|
| 113 |
+
"""'put the green object in the drawer' ->
|
| 114 |
+
('a green object', 'in', 'a drawer') | None. Mechanical determiner-
|
| 115 |
+
phrase split around a spatial preposition; generic English only."""
|
| 116 |
+
import re
|
| 117 |
+
t = text.lower()
|
| 118 |
+
prep = None
|
| 119 |
+
for p in sorted(_INWARD + _OUTWARD, key=len, reverse=True):
|
| 120 |
+
m = re.search(rf"\b{p}\b", t)
|
| 121 |
+
if m:
|
| 122 |
+
prep = (p, m.start(), m.end())
|
| 123 |
+
break
|
| 124 |
+
if not prep:
|
| 125 |
+
return None
|
| 126 |
+
left, right = t[:prep[1]], t[prep[2]:]
|
| 127 |
+
# the landmark ends at the first conjunction — "out of the drawer
|
| 128 |
+
# and put it..." must not leak "and put" into the phrase
|
| 129 |
+
right = right.split(" and ")[0].split(" then ")[0]
|
| 130 |
+
|
| 131 |
+
def np_of(seg, last):
|
| 132 |
+
# GREEDY middle words bounded at closed-class function words
|
| 133 |
+
# (same _STOP mechanism as sig2.atoms_of): without the
|
| 134 |
+
# boundary this produced "a vessel and put" as the audit's X
|
| 135 |
+
# phrase — it grounded nowhere and the audit executed a 9/10
|
| 136 |
+
# true set (audit-bench-caught)
|
| 137 |
+
from .sig2 import _STOP
|
| 138 |
+
ms = list(re.finditer(
|
| 139 |
+
r"\b(?:a|an|the)\s+(?:(?!(?:%s)\b)\w+\s+){0,2}\w+"
|
| 140 |
+
r"(?=\s|$|\.)" % "|".join(_STOP), seg))
|
| 141 |
+
if not ms:
|
| 142 |
+
return None
|
| 143 |
+
m = ms[-1] if last else ms[0]
|
| 144 |
+
w = m.group(0).split()
|
| 145 |
+
while len(w) > 1 and w[-1] in _STOP:
|
| 146 |
+
w.pop()
|
| 147 |
+
head = " ".join(w[1:])
|
| 148 |
+
art = "an" if head[:1] in "aeiou" else "a"
|
| 149 |
+
return f"{art} {head}"
|
| 150 |
+
x, y = np_of(left, last=True), np_of(right, last=False)
|
| 151 |
+
if not x or not y:
|
| 152 |
+
return None
|
| 153 |
+
return x, prep[0], y
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def relation_margin(store, stream, t0, t1, x_phrase, y_phrase,
|
| 157 |
+
inward, n_frames=8):
|
| 158 |
+
"""Containment-change margin for one episode, computed on the SAM
|
| 159 |
+
3.1 VIDEO tracker's masklets (user rule: never per-frame images
|
| 160 |
+
when a video API exists — tracked identity through time is the
|
| 161 |
+
point). Positive = geometry agrees with the query direction; NaN =
|
| 162 |
+
abstain (landmark never seen, or no signal defined).
|
| 163 |
+
|
| 164 |
+
Two complementary signals, because containers OCCLUDE (bench-
|
| 165 |
+
caught: every put-in-drawer IoA delta was 0 — the object disappears
|
| 166 |
+
inside):
|
| 167 |
+
IoA change — X's mask/box overlap with Y rises (put ON, or
|
| 168 |
+
open container where X stays visible)
|
| 169 |
+
presence — the tracker loses X while Y persists (put IN), or
|
| 170 |
+
transition acquires X late (take OUT); tracker probabilities,
|
| 171 |
+
not thresholded detections
|
| 172 |
+
"""
|
| 173 |
+
from .sam3x import track_concepts
|
| 174 |
+
tr = track_concepts(store, stream, t0, t1, [x_phrase, y_phrase],
|
| 175 |
+
n_frames=n_frames)
|
| 176 |
+
if tr is None:
|
| 177 |
+
return float("nan")
|
| 178 |
+
X, Y = tr[x_phrase], tr[y_phrase]
|
| 179 |
+
if max(Y["presence"]) <= 0:
|
| 180 |
+
return float("nan") # scene lacks the landmark: abstain
|
| 181 |
+
n = len(X["presence"])
|
| 182 |
+
k = max(1, n // 3)
|
| 183 |
+
|
| 184 |
+
sigs = []
|
| 185 |
+
traj = []
|
| 186 |
+
for i in range(n):
|
| 187 |
+
if X["masks"][i] is not None and Y["masks"][i] is not None:
|
| 188 |
+
inter = float((X["masks"][i] & Y["masks"][i]).sum())
|
| 189 |
+
traj.append((i, inter / max(1.0,
|
| 190 |
+
float(X["masks"][i].sum()))))
|
| 191 |
+
elif X["boxes"][i] is not None and Y["boxes"][i] is not None:
|
| 192 |
+
traj.append((i, _ioa(X["boxes"][i], Y["boxes"][i])))
|
| 193 |
+
if len(traj) >= 2:
|
| 194 |
+
kk = max(1, len(traj) // 3)
|
| 195 |
+
sigs.append(float(np.mean([v for _, v in traj[-kk:]])
|
| 196 |
+
- np.mean([v for _, v in traj[:kk]])))
|
| 197 |
+
pres = X["presence"]
|
| 198 |
+
early, late = float(np.mean(pres[:k])), float(np.mean(pres[-k:]))
|
| 199 |
+
if abs(early - late) > 0.2:
|
| 200 |
+
# disappearing INTO the container is inward-positive
|
| 201 |
+
sigs.append(early - late)
|
| 202 |
+
if not sigs:
|
| 203 |
+
return float("nan")
|
| 204 |
+
delta = float(np.mean(sigs))
|
| 205 |
+
return delta if inward else -delta
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def verify_relation(store, text, episodes, n_frames=6):
|
| 209 |
+
"""Batch: query text + [(stream,t0,t1)] -> margins array (NaN =
|
| 210 |
+
abstain). Only called on verification-tier candidates."""
|
| 211 |
+
rel = parse_relation(text)
|
| 212 |
+
if rel is None:
|
| 213 |
+
return None
|
| 214 |
+
x, prep, y = rel
|
| 215 |
+
inward = prep in _INWARD
|
| 216 |
+
return np.array([relation_margin(store, s, a, b, x, y, inward,
|
| 217 |
+
n_frames=n_frames)
|
| 218 |
+
for s, a, b in episodes])
|
python/elidedb/identity.py
ADDED
|
@@ -0,0 +1,915 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""OBJECT IDENTITY on the write path, text-free, riding on tracks.
|
| 2 |
+
|
| 3 |
+
Identity is a property of a physical object, not of a name and not of a
|
| 4 |
+
category. The same spoon across every demo gets one id. A spatula gets
|
| 5 |
+
a different id. A red pepper and a green pepper get different ids, even
|
| 6 |
+
though every text-aligned encoder calls them both "pepper".
|
| 7 |
+
|
| 8 |
+
Nothing here turns pixels into language. Regions come from a
|
| 9 |
+
class-agnostic segmenter, identity from a ReID encoder; both vision
|
| 10 |
+
only, both real time. This is separate from objects.py, which proposes
|
| 11 |
+
regions for SEMANTIC retrieval - that channel answers "does this look
|
| 12 |
+
like a cloth", this one answers "is this the same object again".
|
| 13 |
+
|
| 14 |
+
CONTINUITY IS THE ARCHITECTURE
|
| 15 |
+
------------------------------
|
| 16 |
+
The earlier version sampled 3 frames per episode and asked the gallery
|
| 17 |
+
about every crop independently. That is the wrong question asked ~500
|
| 18 |
+
times: it made identity a threshold problem, and the threshold had no
|
| 19 |
+
good setting - 0.55 built attractor ids that swallowed a green box, a
|
| 20 |
+
broccoli, a toy mouse and a carrot; 0.85 fragmented to 1.15 sightings
|
| 21 |
+
per object. Neither end works because a single crop of an object from
|
| 22 |
+
one angle simply is not enough evidence.
|
| 23 |
+
|
| 24 |
+
A tracker answers it for free instead. While an object stays visible,
|
| 25 |
+
IoU + Kalman association says "same object" from GEOMETRY, at 0.25
|
| 26 |
+
ms/frame, with no appearance model and no threshold. The object store
|
| 27 |
+
is consulted only when continuity BREAKS - a new track appears - and
|
| 28 |
+
then it is asked once, about a descriptor pooled over the whole track,
|
| 29 |
+
not once per frame about one view.
|
| 30 |
+
|
| 31 |
+
Measured over 24 episodes / 883 frames of Bridge:
|
| 32 |
+
|
| 33 |
+
median track length 24-48 frames (episodes are ~37)
|
| 34 |
+
singleton tracks 5%
|
| 35 |
+
object-store uploads 19, against 415 tracked detections
|
| 36 |
+
= 95% fewer questions asked
|
| 37 |
+
|
| 38 |
+
That 95% is the design. Fewer questions, each with far more evidence
|
| 39 |
+
behind it.
|
| 40 |
+
|
| 41 |
+
COST, and where it goes
|
| 42 |
+
-----------------------
|
| 43 |
+
Detection is STATELESS, so it batches across episode boundaries;
|
| 44 |
+
association is STATEFUL, so it runs per episode, sequentially. Splitting
|
| 45 |
+
them is what makes every-frame tracking affordable - ultralytics' own
|
| 46 |
+
track() ties them together and pays batch-size-1 prices for both.
|
| 47 |
+
|
| 48 |
+
per-frame track() calls 15.8 ms/frame
|
| 49 |
+
batched track() (list per ep) 9.3
|
| 50 |
+
detect batched + associate 4.8 <- this module
|
| 51 |
+
of which detect (fp16, 448) 4.16
|
| 52 |
+
of which associate 0.13
|
| 53 |
+
of which reid 0.50 (was 4.1 per FRAME)
|
| 54 |
+
|
| 55 |
+
At 18,000 frames per hour of video that is 1.43 min/hour marginal,
|
| 56 |
+
against a 1 min/hour write budget of which 0.23 is already spent by the
|
| 57 |
+
existing pass. STATED STRAIGHT: identity does not fit the budget. What
|
| 58 |
+
riding on tracks bought is the ReID half, which used to dominate and is
|
| 59 |
+
now 0.1 min/hour; the detector is the entire remaining cost, so closing
|
| 60 |
+
the gap is a detector question, not a design one.
|
| 61 |
+
|
| 62 |
+
fp16 is worth 14% of detect (4.94 -> 4.16 ms/frame) with the fitted cut
|
| 63 |
+
and track count unchanged. That took interleaved runs to establish - a
|
| 64 |
+
first attempt compared an fp16 run against fp32 runs taken earlier in a
|
| 65 |
+
long session and read the thermal drift as fp16 being SLOWER.
|
| 66 |
+
|
| 67 |
+
DEVICE. Defaults to the GPU, and this is now measured rather than
|
| 68 |
+
assumed. The old 3-frames-per-episode path put CPU and MPS within 5%
|
| 69 |
+
(19.3 vs 20.3 ms/frame) because per-call overhead dominated a batch of
|
| 70 |
+
three. On real batches the gap is 3.7x - mps 6.7 ms/frame vs cpu 24.8
|
| 71 |
+
at imgsz 640 - and the CPU gets WORSE with batching (93.9 ms/frame at
|
| 72 |
+
batch 16) as threads contend. Frames are already resident on the GPU
|
| 73 |
+
for the FDNN-V pass, and on CUDA a host round trip per frame costs more
|
| 74 |
+
than the inference. CPU is the fallback, never the default.
|
| 75 |
+
"""
|
| 76 |
+
from __future__ import annotations
|
| 77 |
+
|
| 78 |
+
import os
|
| 79 |
+
from types import SimpleNamespace
|
| 80 |
+
|
| 81 |
+
import numpy as np
|
| 82 |
+
|
| 83 |
+
_M: dict = {}
|
| 84 |
+
|
| 85 |
+
# WEIGHTS LIVE IN models/weights, not the repo root. ultralytics
|
| 86 |
+
# downloads into the CWD by default, which is how 86 MB of checkpoints
|
| 87 |
+
# ended up beside the README with 62 MB of it committed to git. _weights
|
| 88 |
+
# resolves a bare name against models/weights and falls back to the bare
|
| 89 |
+
# name so a missing file still auto-downloads - into models/weights,
|
| 90 |
+
# because that is where we then look.
|
| 91 |
+
def _weights(name):
|
| 92 |
+
from pathlib import Path as _P
|
| 93 |
+
d = _P(__file__).resolve().parents[2] / "models" / "weights"
|
| 94 |
+
d.mkdir(parents=True, exist_ok=True)
|
| 95 |
+
p = d / name
|
| 96 |
+
return str(p) if p.exists() else str(d / name)
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
SEG_MODEL = os.environ.get("ELIDEDB_SEG") or _weights("yolo11n-seg.pt")
|
| 100 |
+
REID_MODEL = os.environ.get("ELIDEDB_REID") or _weights("yolo26n-reid.onnx")
|
| 101 |
+
# imgsz 448 rather than 384, deliberately paying 0.5 ms/frame for it.
|
| 102 |
+
# 384 is cheaper (3.76 vs 4.23 ms/frame with fp16) and finds MORE boxes
|
| 103 |
+
# (2.63 vs 2.44 per frame), but the extra boxes are marginal ones: its
|
| 104 |
+
# proven-different pairs reach 0.757 against 0.555 at 448, i.e. noisier
|
| 105 |
+
# negatives, and it recovered 72 multi-episode objects against 81. A
|
| 106 |
+
# cheaper size that degrades the identity is not cheaper.
|
| 107 |
+
DET_SZ = int(os.environ.get("ELIDEDB_DET_SZ", "448"))
|
| 108 |
+
# free: same 2.44 boxes/frame, 4.74 -> 4.23 ms/frame
|
| 109 |
+
DET_FP16 = os.environ.get("ELIDEDB_DET_FP16", "1") not in ("0", "")
|
| 110 |
+
# detection is stateless, so the batch may cross episode boundaries;
|
| 111 |
+
# 3.87 ms/frame at ~37 (one episode) against 3.32 at 64
|
| 112 |
+
DET_BATCH = int(os.environ.get("ELIDEDB_DET_BATCH", "64"))
|
| 113 |
+
# Global motion compensation estimates camera movement before
|
| 114 |
+
# association. Default OFF from MEASUREMENT, not from an assumption
|
| 115 |
+
# about this corpus: with sparseOptFlow, median track length was 47 at
|
| 116 |
+
# imgsz 448 and 38 at 640; with none, 48 and 38 - identical - while the
|
| 117 |
+
# optical flow cost 1.8 ms/frame, 7x the entire association step. A
|
| 118 |
+
# corpus with a moving camera must re-measure and will likely want it
|
| 119 |
+
# back; that is what the knob is for.
|
| 120 |
+
GMC = os.environ.get("ELIDEDB_GMC", "none")
|
| 121 |
+
# how many views of a track to embed. The pooled descriptor is what the
|
| 122 |
+
# object store sees, so this trades evidence against ReID calls.
|
| 123 |
+
EXEMPLARS = int(os.environ.get("ELIDEDB_OBJ_VIEWS", "4"))
|
| 124 |
+
|
| 125 |
+
# "same physical object" cut. Not a constant to be tuned by hand - see
|
| 126 |
+
# calibrate(), which fits it from the corpus being ingested. This value
|
| 127 |
+
# is only the fallback for a corpus too small to fit on.
|
| 128 |
+
MATCH = float(os.environ.get("ELIDEDB_OBJ_MATCH", "0.85"))
|
| 129 |
+
# quantile of the free-negative distribution used as the cut
|
| 130 |
+
CAL_Q = float(os.environ.get("ELIDEDB_OBJ_CAL_Q", "99.5"))
|
| 131 |
+
# Adding an exemplar WIDENS the id's acceptance region, because matching
|
| 132 |
+
# takes a max over the gallery. Unchecked that is a runaway: absorb ->
|
| 133 |
+
# widen -> absorb more. A new view is admitted only if it is close to
|
| 134 |
+
# EVERY exemplar already held, which keeps an id tight instead of
|
| 135 |
+
# letting it sprawl across a chain of intermediate appearances.
|
| 136 |
+
MAX_EXEMPLARS = 5
|
| 137 |
+
ADMIT = float(os.environ.get("ELIDEDB_OBJ_ADMIT", "0.70"))
|
| 138 |
+
MIN_SIDE = 16
|
| 139 |
+
MAX_AREA = 0.5
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
def device():
|
| 143 |
+
d = os.environ.get("ELIDEDB_DEVICE")
|
| 144 |
+
if d:
|
| 145 |
+
return d
|
| 146 |
+
try:
|
| 147 |
+
import torch
|
| 148 |
+
if torch.cuda.is_available():
|
| 149 |
+
return "cuda"
|
| 150 |
+
if torch.backends.mps.is_available():
|
| 151 |
+
return "mps"
|
| 152 |
+
except Exception:
|
| 153 |
+
pass
|
| 154 |
+
return "cpu"
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def _load():
|
| 158 |
+
# check for THIS function's own keys, not for _M being non-empty:
|
| 159 |
+
# propose() also caches into _M, so `if _M` short-circuited here and
|
| 160 |
+
# left reid unloaded the moment a proposer had run first.
|
| 161 |
+
if "reid" in _M:
|
| 162 |
+
return _M
|
| 163 |
+
from ultralytics import YOLO
|
| 164 |
+
from ultralytics.trackers.utils.reid import ReID
|
| 165 |
+
dev = device()
|
| 166 |
+
_M["seg"] = YOLO(SEG_MODEL)
|
| 167 |
+
# the ReID encoder runs under onnxruntime and picks its own
|
| 168 |
+
# execution provider; it does not share the torch device
|
| 169 |
+
_M["reid"] = ReID(REID_MODEL, device="cpu" if dev == "mps" else dev)
|
| 170 |
+
_M["dev"] = dev
|
| 171 |
+
return _M
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
PROPOSER = os.environ.get("ELIDEDB_PROPOSER") or _weights("FastSAM-s.pt")
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
def propose(frames, imgsz=DET_SZ, batch=DET_BATCH, conf=0.25):
|
| 178 |
+
"""CLASS-AGNOSTIC region proposal. Not YOLO-with-single_cls.
|
| 179 |
+
|
| 180 |
+
MEASURED, same frames, same model: single_cls=True and
|
| 181 |
+
single_cls=False give byte-identical output - 3.10 det/frame, all of
|
| 182 |
+
it {oven, sink, bowl, person, spoon, wine glass}. single_cls is a
|
| 183 |
+
TRAINING flag; at inference the detector still fires only on
|
| 184 |
+
COCO-shaped things. Everything built on it inherited an 80-class
|
| 185 |
+
prior that was never intended, which is why looking for a container
|
| 186 |
+
found 0.7 regions per episode and they were ovens.
|
| 187 |
+
|
| 188 |
+
FastSAM is trained on SA-1B with no class list at all, and on this
|
| 189 |
+
corpus returns 45.3 regions/frame against COCO YOLO's 2.6, with 3.5
|
| 190 |
+
large structures per frame against 1.1.
|
| 191 |
+
|
| 192 |
+
It costs 24.93 ms/frame batched fp16 = 7.48 min per hour of video,
|
| 193 |
+
which is 7x the write budget - so this is a TEACHER. It is paid once
|
| 194 |
+
to make training pairs and then distilled, exactly like every other
|
| 195 |
+
teacher here. Never put it on the write path directly.
|
| 196 |
+
"""
|
| 197 |
+
from ultralytics import FastSAM
|
| 198 |
+
if "prop" not in _M:
|
| 199 |
+
_M["prop"] = FastSAM(PROPOSER)
|
| 200 |
+
out = []
|
| 201 |
+
kw = {"quantize": "fp16"} if DET_FP16 else {}
|
| 202 |
+
for i in range(0, len(frames), batch):
|
| 203 |
+
chunk = frames[i:i + batch]
|
| 204 |
+
res = _M["prop"].predict(chunk, device=device(), verbose=False,
|
| 205 |
+
imgsz=imgsz, conf=conf, **kw)
|
| 206 |
+
for im, r in zip(chunk, res):
|
| 207 |
+
H, W = im.shape[:2]
|
| 208 |
+
if r.boxes is None or not len(r.boxes):
|
| 209 |
+
out.append(np.zeros((0, 4), np.int32))
|
| 210 |
+
continue
|
| 211 |
+
b = r.boxes.xyxy.cpu().numpy()
|
| 212 |
+
b = np.stack([b[:, 0].clip(0, W), b[:, 1].clip(0, H),
|
| 213 |
+
b[:, 2].clip(0, W), b[:, 3].clip(0, H)],
|
| 214 |
+
1).astype(np.int32)
|
| 215 |
+
wh = (b[:, 2] - b[:, 0]), (b[:, 3] - b[:, 1])
|
| 216 |
+
out.append(b[(wh[0] >= MIN_SIDE) & (wh[1] >= MIN_SIDE)])
|
| 217 |
+
return out
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
def detect(frames, imgsz=DET_SZ, batch=DET_BATCH, conf=0.25):
|
| 221 |
+
"""Class-agnostic regions per frame: [(boxes Nx4, conf N, area N)].
|
| 222 |
+
|
| 223 |
+
Stateless and therefore batched - the caller may hand in frames from
|
| 224 |
+
several episodes at once. single_cls collapses the 80 COCO classes
|
| 225 |
+
into one anonymous class, so the model reports "a thing is here" and
|
| 226 |
+
never a label.
|
| 227 |
+
"""
|
| 228 |
+
m = _load()
|
| 229 |
+
out = []
|
| 230 |
+
kw = {"quantize": "fp16"} if DET_FP16 else {}
|
| 231 |
+
for i in range(0, len(frames), batch):
|
| 232 |
+
chunk = frames[i:i + batch]
|
| 233 |
+
res = m["seg"].predict(chunk, device=m["dev"], single_cls=True,
|
| 234 |
+
verbose=False, imgsz=imgsz, conf=conf, **kw)
|
| 235 |
+
for im, r in zip(chunk, res):
|
| 236 |
+
H, W = im.shape[:2]
|
| 237 |
+
if r.boxes is None or not len(r.boxes):
|
| 238 |
+
out.append((np.zeros((0, 4), np.int32), np.zeros(0, np.float32),
|
| 239 |
+
np.zeros(0, np.float32)))
|
| 240 |
+
continue
|
| 241 |
+
xy = r.boxes.xyxy.cpu().numpy()
|
| 242 |
+
cf = r.boxes.conf.cpu().numpy()
|
| 243 |
+
ar = (r.masks.data.sum((1, 2)).cpu().numpy()
|
| 244 |
+
if r.masks is not None else np.zeros(len(xy)))
|
| 245 |
+
b = np.stack([xy[:, 0].clip(0, W), xy[:, 1].clip(0, H),
|
| 246 |
+
xy[:, 2].clip(0, W), xy[:, 3].clip(0, H)],
|
| 247 |
+
1).astype(np.int32)
|
| 248 |
+
wh = (b[:, 2] - b[:, 0]), (b[:, 3] - b[:, 1])
|
| 249 |
+
ok = ((wh[0] >= MIN_SIDE) & (wh[1] >= MIN_SIDE)
|
| 250 |
+
& (wh[0] * wh[1] <= MAX_AREA * W * H))
|
| 251 |
+
out.append((b[ok], cf[ok].astype(np.float32),
|
| 252 |
+
ar[ok].astype(np.float32)))
|
| 253 |
+
return out
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
class _Dets:
|
| 257 |
+
"""The minimal shape ultralytics' trackers consume: xywh/conf/cls,
|
| 258 |
+
a length, and boolean-mask indexing (they split detections into
|
| 259 |
+
high- and low-confidence subsets). Handing them this instead of a
|
| 260 |
+
Results object is what lets detection run batched somewhere else."""
|
| 261 |
+
|
| 262 |
+
__slots__ = ("xywh", "conf", "cls")
|
| 263 |
+
|
| 264 |
+
def __init__(self, xywh, conf, cls):
|
| 265 |
+
self.xywh, self.conf, self.cls = xywh, conf, cls
|
| 266 |
+
|
| 267 |
+
def __len__(self):
|
| 268 |
+
return len(self.conf)
|
| 269 |
+
|
| 270 |
+
def __getitem__(self, m):
|
| 271 |
+
return _Dets(self.xywh[m], self.conf[m], self.cls[m])
|
| 272 |
+
|
| 273 |
+
@property
|
| 274 |
+
def xyxy(self):
|
| 275 |
+
"""Corner boxes, which global motion compensation asks for.
|
| 276 |
+
|
| 277 |
+
byte_tracker.py calls `self.gmc.apply(img, results_high.xyxy)` to
|
| 278 |
+
mask moving objects out before estimating camera motion, inside a
|
| 279 |
+
try/except that WARNS and falls back to an identity warp. Without
|
| 280 |
+
this property every frame took that fallback: 1,434 warnings in
|
| 281 |
+
the first two minutes of a write, and camera motion silently not
|
| 282 |
+
compensated.
|
| 283 |
+
|
| 284 |
+
It was invisible on this corpus because Bridge's camera is fixed,
|
| 285 |
+
where identity is the correct warp anyway - so the bug cost
|
| 286 |
+
nothing here and everything on a corpus that moves. A robot that
|
| 287 |
+
drives or a vehicle camera would have had its tracks fragmented
|
| 288 |
+
by exactly the motion GMC exists to remove.
|
| 289 |
+
|
| 290 |
+
xywh is centre-based (built that way in Stream.push, and what
|
| 291 |
+
ultralytics' xywh2ltwh assumes), so the corners are centre +/-
|
| 292 |
+
half-extent.
|
| 293 |
+
"""
|
| 294 |
+
cx, cy, w, h = (self.xywh[:, 0], self.xywh[:, 1],
|
| 295 |
+
self.xywh[:, 2], self.xywh[:, 3])
|
| 296 |
+
return np.stack([cx - w / 2, cy - h / 2,
|
| 297 |
+
cx + w / 2, cy + h / 2], 1)
|
| 298 |
+
|
| 299 |
+
|
| 300 |
+
def _tracker(gmc=None):
|
| 301 |
+
from ultralytics.trackers.bot_sort import BOTSORT
|
| 302 |
+
from ultralytics.utils import YAML
|
| 303 |
+
import ultralytics
|
| 304 |
+
from pathlib import Path
|
| 305 |
+
cfg = YAML.load(Path(ultralytics.__file__).parent
|
| 306 |
+
/ "cfg/trackers/botsort.yaml")
|
| 307 |
+
cfg["gmc_method"] = gmc or GMC
|
| 308 |
+
# BoT-SORT's own ReID would embed every detection on every frame,
|
| 309 |
+
# which is the cost this design exists to avoid. Association here is
|
| 310 |
+
# geometry only; appearance is consulted once per track, later.
|
| 311 |
+
cfg["with_reid"] = False
|
| 312 |
+
return BOTSORT(args=SimpleNamespace(**cfg))
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
class Stream:
|
| 316 |
+
"""CONTINUOUS tracking over a stream. No episode boundaries.
|
| 317 |
+
|
| 318 |
+
link() below takes one episode's frames and builds a fresh tracker
|
| 319 |
+
for each - correct when a corpus ships discrete demos, and an
|
| 320 |
+
assumption the engine has no right to make. Raw capture is a
|
| 321 |
+
continuous recording; episodes are something the engine must
|
| 322 |
+
PRODUCE. A driving log, a surveillance feed and a surgical recording
|
| 323 |
+
have no cuts to reset on.
|
| 324 |
+
|
| 325 |
+
So the tracker runs for the life of the stream and a track ends when
|
| 326 |
+
the OBJECT does - it leaves frame, is occluded past the buffer, or
|
| 327 |
+
is carried away. That track's span is exactly a PRESENCE INTERVAL:
|
| 328 |
+
"this object was here, from t0 to t1", one row however long it
|
| 329 |
+
lasted. An object sitting still for four hours is one row, not
|
| 330 |
+
72,000, and it is finally recorded at all - the motion-triggered
|
| 331 |
+
element path never produced a row for anything that did not move.
|
| 332 |
+
|
| 333 |
+
Online by construction: detections go in frame by frame, closed
|
| 334 |
+
tracks come out as they close, and only the open tracks are held.
|
| 335 |
+
A stream that does not fit in memory is the normal case, not an
|
| 336 |
+
edge case.
|
| 337 |
+
|
| 338 |
+
Exemplar crops are retained per OPEN track and embedded once, at
|
| 339 |
+
close - the same "ask the object store one question per track, with
|
| 340 |
+
a whole sighting behind it" that made identity work, now without
|
| 341 |
+
needing the episode to know when to ask.
|
| 342 |
+
"""
|
| 343 |
+
|
| 344 |
+
def __init__(self, gmc=None, views=EXEMPLARS, buffer_s=6.0):
|
| 345 |
+
self.tr = _tracker(gmc)
|
| 346 |
+
self.views = views
|
| 347 |
+
# how long a track survives with no detection before it is
|
| 348 |
+
# declared ended, in SECONDS not frames: a frame count means
|
| 349 |
+
# different things at 5 fps and 30 fps, and the corpus chooses
|
| 350 |
+
# the frame rate.
|
| 351 |
+
self.buffer_s = buffer_s
|
| 352 |
+
self.open: dict[int, dict] = {}
|
| 353 |
+
self.last_ts = None
|
| 354 |
+
|
| 355 |
+
def update(self, ts, det, frame=None):
|
| 356 |
+
"""One frame in; the tracks that CLOSED at this frame out.
|
| 357 |
+
|
| 358 |
+
`det` is (boxes Nx4, conf N, area N) from detect()/propose().
|
| 359 |
+
"""
|
| 360 |
+
b, c, a = det
|
| 361 |
+
self.last_ts = int(ts)
|
| 362 |
+
seen = set()
|
| 363 |
+
if len(b):
|
| 364 |
+
xywh = np.stack([(b[:, 0] + b[:, 2]) / 2,
|
| 365 |
+
(b[:, 1] + b[:, 3]) / 2,
|
| 366 |
+
b[:, 2] - b[:, 0], b[:, 3] - b[:, 1]], 1)
|
| 367 |
+
rows = self.tr.update(
|
| 368 |
+
_Dets(xywh.astype(np.float32), c,
|
| 369 |
+
np.zeros(len(b), np.float32)), frame)
|
| 370 |
+
for row in rows:
|
| 371 |
+
tid, di = int(row[-4]), int(row[-1])
|
| 372 |
+
if di >= len(b):
|
| 373 |
+
continue
|
| 374 |
+
seen.add(tid)
|
| 375 |
+
t = self.open.setdefault(tid, {
|
| 376 |
+
"ts": int(ts), "t1": int(ts), "n": 0, "t": [],
|
| 377 |
+
"box": [], "conf": [], "area": [], "crops": []})
|
| 378 |
+
t["t1"] = int(ts)
|
| 379 |
+
t["n"] += 1
|
| 380 |
+
# the box list without its timestamps is a SHAPE, not a
|
| 381 |
+
# trajectory; every consumer that wanted motion had to
|
| 382 |
+
# guess the time axis back from (ts, t1, n)
|
| 383 |
+
t["t"].append(int(ts))
|
| 384 |
+
t["box"].append(b[di])
|
| 385 |
+
t["conf"].append(float(c[di]))
|
| 386 |
+
t["area"].append(float(a[di]))
|
| 387 |
+
# retain a bounded, spread set of views for the one
|
| 388 |
+
# ReID call this track will ever cost
|
| 389 |
+
if frame is not None and len(t["crops"]) < self.views:
|
| 390 |
+
x0, y0, x1, y1 = (int(v) for v in b[di])
|
| 391 |
+
if x1 > x0 and y1 > y0:
|
| 392 |
+
t["crops"].append((b[di], frame[y0:y1, x0:x1]))
|
| 393 |
+
return self._reap(seen)
|
| 394 |
+
|
| 395 |
+
def _reap(self, seen):
|
| 396 |
+
gap = self.buffer_s * 1e9
|
| 397 |
+
done = []
|
| 398 |
+
for tid in list(self.open):
|
| 399 |
+
if tid in seen:
|
| 400 |
+
continue
|
| 401 |
+
if self.last_ts - self.open[tid]["t1"] > gap:
|
| 402 |
+
done.append((tid, self.open.pop(tid)))
|
| 403 |
+
return done
|
| 404 |
+
|
| 405 |
+
def flush(self):
|
| 406 |
+
"""End of stream: everything still open is still real."""
|
| 407 |
+
out = list(self.open.items())
|
| 408 |
+
self.open = {}
|
| 409 |
+
return out
|
| 410 |
+
|
| 411 |
+
|
| 412 |
+
def link(dets, frames=None, gmc=None):
|
| 413 |
+
"""Associate one episode's consecutive detections into tracks.
|
| 414 |
+
|
| 415 |
+
Returns {track_id: {"f": [frame idx], "box": [...], "conf": [...],
|
| 416 |
+
"area": [...]}}. A fresh tracker per call, which is right only when
|
| 417 |
+
the caller really does have a scene cut. Prefer Stream for
|
| 418 |
+
continuous capture; this remains for corpora that ship discrete
|
| 419 |
+
clips, where resetting at a genuine cut avoids carrying Kalman state
|
| 420 |
+
across it.
|
| 421 |
+
"""
|
| 422 |
+
tr = _tracker(gmc)
|
| 423 |
+
out: dict[int, dict] = {}
|
| 424 |
+
for j, (b, c, a) in enumerate(dets):
|
| 425 |
+
if not len(b):
|
| 426 |
+
continue
|
| 427 |
+
xywh = np.stack([(b[:, 0] + b[:, 2]) / 2, (b[:, 1] + b[:, 3]) / 2,
|
| 428 |
+
b[:, 2] - b[:, 0], b[:, 3] - b[:, 1]], 1)
|
| 429 |
+
det = _Dets(xywh.astype(np.float32), c, np.zeros(len(b), np.float32))
|
| 430 |
+
rows = tr.update(det, None if frames is None else frames[j])
|
| 431 |
+
for row in rows:
|
| 432 |
+
# ultralytics packs [*xyxy, track_id, conf, cls, det_idx]
|
| 433 |
+
tid, di = int(row[-4]), int(row[-1])
|
| 434 |
+
if di >= len(b):
|
| 435 |
+
continue
|
| 436 |
+
t = out.setdefault(tid, {"f": [], "box": [], "conf": [],
|
| 437 |
+
"area": []})
|
| 438 |
+
t["f"].append(j)
|
| 439 |
+
t["box"].append(b[di])
|
| 440 |
+
t["conf"].append(float(c[di]))
|
| 441 |
+
t["area"].append(float(a[di]))
|
| 442 |
+
return out
|
| 443 |
+
|
| 444 |
+
|
| 445 |
+
def _views(track, k=EXEMPLARS):
|
| 446 |
+
"""Pick up to k frames of a track: the best-scoring view in each of
|
| 447 |
+
k equal slices of the track's life. Spread matters more than score -
|
| 448 |
+
k views of the same instant carry one view's worth of evidence."""
|
| 449 |
+
n = len(track["f"])
|
| 450 |
+
if n <= k:
|
| 451 |
+
return list(range(n))
|
| 452 |
+
edges = np.linspace(0, n, k + 1).round().astype(int)
|
| 453 |
+
conf = np.asarray(track["conf"])
|
| 454 |
+
return [int(a + conf[a:b].argmax())
|
| 455 |
+
for a, b in zip(edges[:-1], edges[1:]) if b > a]
|
| 456 |
+
|
| 457 |
+
|
| 458 |
+
def features(frame, boxes):
|
| 459 |
+
"""Appearance embedding per region. Vision only, no text."""
|
| 460 |
+
if not len(boxes):
|
| 461 |
+
return np.zeros((0, 512), np.float32)
|
| 462 |
+
m = _load()
|
| 463 |
+
b = np.asarray(boxes, np.float32)
|
| 464 |
+
xywh = np.stack([(b[:, 0] + b[:, 2]) / 2, (b[:, 1] + b[:, 3]) / 2,
|
| 465 |
+
b[:, 2] - b[:, 0], b[:, 3] - b[:, 1]], 1)
|
| 466 |
+
f = np.asarray(m["reid"](frame, xywh))
|
| 467 |
+
f = f.reshape(len(boxes), -1).astype(np.float32)
|
| 468 |
+
return f / (np.linalg.norm(f, axis=1, keepdims=True) + 1e-8)
|
| 469 |
+
|
| 470 |
+
|
| 471 |
+
def descriptors(frames, tracks, k=EXEMPLARS):
|
| 472 |
+
"""One pooled appearance vector per track - the unit of upload.
|
| 473 |
+
|
| 474 |
+
This is the "break in continuity" moment: while a track holds, the
|
| 475 |
+
object store hears nothing about it. When the track is done, it is
|
| 476 |
+
described ONCE, by a mean over k views spread across its life, and
|
| 477 |
+
that single descriptor is what the store is asked about.
|
| 478 |
+
"""
|
| 479 |
+
ids, out = [], []
|
| 480 |
+
for tid, t in tracks.items():
|
| 481 |
+
V = []
|
| 482 |
+
for i in _views(t, k):
|
| 483 |
+
V.append(features(frames[t["f"][i]], [t["box"][i]])[0])
|
| 484 |
+
v = np.mean(V, 0)
|
| 485 |
+
ids.append(tid)
|
| 486 |
+
out.append(v / (np.linalg.norm(v) + 1e-8))
|
| 487 |
+
return ids, (np.stack(out) if out else np.zeros((0, 512), np.float32))
|
| 488 |
+
|
| 489 |
+
|
| 490 |
+
def _iou(a, b):
|
| 491 |
+
x0, y0 = max(a[0], b[0]), max(a[1], b[1])
|
| 492 |
+
x1, y1 = min(a[2], b[2]), min(a[3], b[3])
|
| 493 |
+
inter = max(x1 - x0, 0) * max(y1 - y0, 0)
|
| 494 |
+
ua = ((a[2] - a[0]) * (a[3] - a[1])
|
| 495 |
+
+ (b[2] - b[0]) * (b[3] - b[1]) - inter)
|
| 496 |
+
return inter / ua if ua > 0 else 0.0
|
| 497 |
+
|
| 498 |
+
|
| 499 |
+
def free_negatives(tracks, iou_max=0.1):
|
| 500 |
+
"""Track-id pairs that are CERTAINLY different physical objects.
|
| 501 |
+
|
| 502 |
+
Two tracks occupying DISJOINT REGIONS OF THE SAME FRAME are two
|
| 503 |
+
different objects: one object cannot be in two places at one
|
| 504 |
+
instant, and the tracker has already linked each object to itself,
|
| 505 |
+
so a second track is a second thing. No annotation, no dataset
|
| 506 |
+
knowledge, no prior - geometry alone.
|
| 507 |
+
|
| 508 |
+
The disjointness test is not decoration. Co-existence ALONE is not
|
| 509 |
+
enough, because a detector will sometimes put two boxes on one
|
| 510 |
+
object; those two tracks co-exist and look identical, and they land
|
| 511 |
+
exactly in the high tail that a calibration quantile reads. Fitting
|
| 512 |
+
on co-existence alone drove the cut to its 0.95 ceiling on a 300-
|
| 513 |
+
episode run and fragmented the store to 1.21 sightings per object.
|
| 514 |
+
Requiring the boxes to be apart removes the double detections, which
|
| 515 |
+
are the only way two tracks of the same object can co-exist.
|
| 516 |
+
"""
|
| 517 |
+
out, ids = [], list(tracks)
|
| 518 |
+
for x in range(len(ids)):
|
| 519 |
+
a = tracks[ids[x]]
|
| 520 |
+
fa = {f: i for i, f in enumerate(a["f"])}
|
| 521 |
+
for y in range(x + 1, len(ids)):
|
| 522 |
+
b = tracks[ids[y]]
|
| 523 |
+
shared = [f for f in b["f"] if f in fa]
|
| 524 |
+
if not shared:
|
| 525 |
+
continue
|
| 526 |
+
bf = {f: i for i, f in enumerate(b["f"])}
|
| 527 |
+
if all(_iou(a["box"][fa[f]], b["box"][bf[f]]) <= iou_max
|
| 528 |
+
for f in shared):
|
| 529 |
+
out.append((ids[x], ids[y]))
|
| 530 |
+
return out
|
| 531 |
+
|
| 532 |
+
|
| 533 |
+
def free_positives(tracks, iou_min=0.8):
|
| 534 |
+
"""Track-id pairs that are CERTAINLY the SAME physical object.
|
| 535 |
+
|
| 536 |
+
The mirror of free_negatives, from the mirror of its argument. That
|
| 537 |
+
one says two tracks in DISJOINT regions of one frame are two objects,
|
| 538 |
+
because one object cannot be in two places at once. This one says two
|
| 539 |
+
tracks in the SAME region of one frame are ONE object, because two
|
| 540 |
+
objects cannot occupy one place at once. Both are geometry; neither
|
| 541 |
+
needs an annotation, a label or a dataset prior.
|
| 542 |
+
|
| 543 |
+
These pairs are DOUBLE DETECTIONS - the detector put two boxes on one
|
| 544 |
+
thing - which is why free_negatives works to exclude them. They were
|
| 545 |
+
treated as waste. They are not waste: they are the only proven-same
|
| 546 |
+
evidence this corpus can produce for free, and without them the cut
|
| 547 |
+
was fitted from one side of a two-sided decision.
|
| 548 |
+
|
| 549 |
+
iou_min is high on purpose. Overlap alone does not prove identity:
|
| 550 |
+
nested and contacting things - a lid on a jar, a hand on a pot -
|
| 551 |
+
overlap heavily and are two objects. Measured on fresh_bench, the
|
| 552 |
+
positive sample gets monotonically cleaner as the bar rises, and the
|
| 553 |
+
AUC it implies rises with it (0.860 at IoU>0.5 to 0.902 at IoU>0.95),
|
| 554 |
+
i.e. the loose bands are contaminated by exactly those nested pairs.
|
| 555 |
+
0.8 keeps thousands of pairs while paying most of that gap.
|
| 556 |
+
|
| 557 |
+
BIAS, stated, because it bounds what this can conclude: a same-frame
|
| 558 |
+
positive is the EASIEST positive there is - one instant, one
|
| 559 |
+
viewpoint, one exposure - while a same-frame negative is the HARDEST
|
| 560 |
+
negative. So the fitted cut is bracketed by two optimistic samples
|
| 561 |
+
pulling in opposite directions, and a genuine cross-episode
|
| 562 |
+
re-identification is harder than anything measured here.
|
| 563 |
+
"""
|
| 564 |
+
out, ids = [], list(tracks)
|
| 565 |
+
for x in range(len(ids)):
|
| 566 |
+
a = tracks[ids[x]]
|
| 567 |
+
fa = {f: i for i, f in enumerate(a["f"])}
|
| 568 |
+
for y in range(x + 1, len(ids)):
|
| 569 |
+
b = tracks[ids[y]]
|
| 570 |
+
shared = [f for f in b["f"] if f in fa]
|
| 571 |
+
if not shared:
|
| 572 |
+
continue
|
| 573 |
+
bf = {f: i for i, f in enumerate(b["f"])}
|
| 574 |
+
if all(_iou(a["box"][fa[f]], b["box"][bf[f]]) >= iou_min
|
| 575 |
+
for f in shared):
|
| 576 |
+
out.append((ids[x], ids[y]))
|
| 577 |
+
return out
|
| 578 |
+
|
| 579 |
+
|
| 580 |
+
def interval_pairs(rows, iou_diff=0.1, iou_same=0.8, start_frac=0.25):
|
| 581 |
+
"""free_negatives and free_positives over CLOSED intervals.
|
| 582 |
+
|
| 583 |
+
The write path does not hold tracks - it holds closed presence
|
| 584 |
+
intervals, each reduced to one box. Same two geometric arguments,
|
| 585 |
+
applied to what the writer actually has.
|
| 586 |
+
|
| 587 |
+
APPROXIMATION, stated: the retained box is the track's FIRST box, so
|
| 588 |
+
two boxes are only comparable when the tracks START together. That is
|
| 589 |
+
always true of the double detections the positive test is looking
|
| 590 |
+
for, and start_frac enforces it rather than assuming it. The negative
|
| 591 |
+
test does not need the guard - a pair wrongly called disjoint is a
|
| 592 |
+
pair dropped, not a pair mislabelled.
|
| 593 |
+
|
| 594 |
+
Args:
|
| 595 |
+
rows: (stream, meta) or (stream, meta, vec) as close_track emits;
|
| 596 |
+
meta carries ts, t1 and box.
|
| 597 |
+
Returns:
|
| 598 |
+
(neg, pos) index-pair lists into `rows`.
|
| 599 |
+
"""
|
| 600 |
+
n = len(rows)
|
| 601 |
+
key = sorted(range(n), key=lambda i: (rows[i][0], rows[i][1]["ts"]))
|
| 602 |
+
neg, pos = [], []
|
| 603 |
+
for a in range(len(key)):
|
| 604 |
+
i = key[a]
|
| 605 |
+
si, mi = rows[i][0], rows[i][1]
|
| 606 |
+
for b in range(a + 1, len(key)):
|
| 607 |
+
j = key[b]
|
| 608 |
+
sj, mj = rows[j][0], rows[j][1]
|
| 609 |
+
if sj != si or mj["ts"] > mi["t1"]:
|
| 610 |
+
break # sorted: no later one overlaps
|
| 611 |
+
ov = _iou(mi["box"], mj["box"])
|
| 612 |
+
# the sweep runs in time order, so emit (min, max) rather
|
| 613 |
+
# than (earlier, later) - a pair is unordered and callers
|
| 614 |
+
# should not have to know which way round it came out.
|
| 615 |
+
e = (i, j) if i < j else (j, i)
|
| 616 |
+
if ov <= iou_diff:
|
| 617 |
+
neg.append(e)
|
| 618 |
+
elif ov >= iou_same:
|
| 619 |
+
span = min(mi["t1"] - mi["ts"], mj["t1"] - mj["ts"])
|
| 620 |
+
if abs(mj["ts"] - mi["ts"]) <= start_frac * max(span, 1):
|
| 621 |
+
pos.append(e)
|
| 622 |
+
return neg, pos
|
| 623 |
+
|
| 624 |
+
|
| 625 |
+
def handoff_pairs(rows, gap_s=1.0, iou_min=0.5):
|
| 626 |
+
"""Proven-same evidence every corpus produces in BULK: track handoffs.
|
| 627 |
+
|
| 628 |
+
An interval that ends and another that begins at the same place
|
| 629 |
+
within `gap_s` is one object twice - things do not teleport, and
|
| 630 |
+
nothing else can arrive at that exact spot that fast. This is the
|
| 631 |
+
within-group continuity that double-detection positives cannot
|
| 632 |
+
supply (a clean proposer produces almost none: 77 pairs on the sim
|
| 633 |
+
corpus vs 927k negatives), and it is exactly the decision event
|
| 634 |
+
binding depends on: "is the block after the occlusion the block
|
| 635 |
+
from before it".
|
| 636 |
+
|
| 637 |
+
The retained box is the track's FIRST box, so a pair is emitted only
|
| 638 |
+
when the dying track's origin still overlaps the resuming track's
|
| 639 |
+
origin - true for the static objects that dominate any corpus, a
|
| 640 |
+
miss (never a mislabel) for objects whose track broke mid-flight.
|
| 641 |
+
"""
|
| 642 |
+
from bisect import bisect_right
|
| 643 |
+
by = {}
|
| 644 |
+
for i, r in enumerate(rows):
|
| 645 |
+
by.setdefault(r[0], []).append(i)
|
| 646 |
+
out = []
|
| 647 |
+
gap = int(gap_s * 1e9)
|
| 648 |
+
for idxs in by.values():
|
| 649 |
+
idxs.sort(key=lambda i: rows[i][1]["ts"])
|
| 650 |
+
ts = [rows[i][1]["ts"] for i in idxs]
|
| 651 |
+
for i in idxs:
|
| 652 |
+
e = rows[i][1]["t1"]
|
| 653 |
+
for k in range(bisect_right(ts, e), bisect_right(ts, e + gap)):
|
| 654 |
+
j = idxs[k]
|
| 655 |
+
if _iou(rows[i][1]["box"], rows[j][1]["box"]) >= iou_min:
|
| 656 |
+
out.append((i, j) if i < j else (j, i))
|
| 657 |
+
return out
|
| 658 |
+
|
| 659 |
+
|
| 660 |
+
def calibrate(V, pairs, pos=None, q=CAL_Q, floor=0.5, ceil=0.95):
|
| 661 |
+
"""Fit the match cut from proven pairs, on this corpus.
|
| 662 |
+
|
| 663 |
+
TWO-SIDED when proven-same pairs are supplied, and that is the
|
| 664 |
+
correction. The one-sided form below reads the q-th percentile of the
|
| 665 |
+
proven-DIFFERENT similarities - "tighter than all but q% of what I
|
| 666 |
+
can show is not the same object" - which is sound only if that sample
|
| 667 |
+
is clean. On fresh_bench it was not: double detections are 0.5% of
|
| 668 |
+
co-existing pairs and q=99.5 reads the top 0.5%, so the cut was very
|
| 669 |
+
nearly a readout of the contamination. It sat at 0.739, where only
|
| 670 |
+
26% of proven-same pairs are accepted, and the store fragmented to
|
| 671 |
+
78% singletons.
|
| 672 |
+
|
| 673 |
+
Fixing the contamination alone moves it to 0.692 - still 34% recall.
|
| 674 |
+
A one-sided fit cannot do better, because the tail of the negatives
|
| 675 |
+
says nothing about where the positives are. Hence Youden's J over
|
| 676 |
+
both: the cut that maximises (true accept rate - false merge rate),
|
| 677 |
+
the standard criterion when neither error has a stated price. It is
|
| 678 |
+
fitted, not chosen; no number here is hand-picked.
|
| 679 |
+
|
| 680 |
+
Fails safe in stages: no positives, use the one-sided percentile; too
|
| 681 |
+
few negatives to estimate at all, keep MATCH.
|
| 682 |
+
|
| 683 |
+
Args:
|
| 684 |
+
V: (N, D) unit-norm track descriptors.
|
| 685 |
+
pairs: proven-DIFFERENT index pairs, from free_negatives.
|
| 686 |
+
pos: proven-SAME index pairs, from free_positives. Optional.
|
| 687 |
+
"""
|
| 688 |
+
if len(pairs) < 20:
|
| 689 |
+
return MATCH, len(pairs)
|
| 690 |
+
neg = np.array([float(V[i] @ V[j]) for i, j in pairs])
|
| 691 |
+
if pos is None or len(pos) < 20:
|
| 692 |
+
return float(np.clip(np.percentile(neg, q), floor, ceil)), len(neg)
|
| 693 |
+
p = np.array([float(V[i] @ V[j]) for i, j in pos])
|
| 694 |
+
grid = np.linspace(0.0, 1.0, 201)
|
| 695 |
+
j = ((p[None, :] >= grid[:, None]).mean(1)
|
| 696 |
+
- (neg[None, :] >= grid[:, None]).mean(1))
|
| 697 |
+
return float(np.clip(grid[int(j.argmax())], floor, ceil)), len(neg)
|
| 698 |
+
|
| 699 |
+
|
| 700 |
+
def fit_cut(V, groups, neg, pos=None, grid=None, handoff=None):
|
| 701 |
+
"""Fit the cut by MAXIMISING what identity is actually for.
|
| 702 |
+
|
| 703 |
+
Both free samples - proven-different and proven-same - are SAME-FRAME
|
| 704 |
+
pairs. The decision the gallery actually makes is not that one. It is
|
| 705 |
+
"is this new track the object I saw in a different episode?", and no
|
| 706 |
+
same-frame pair is evidence about it. So a statistic of those pairs
|
| 707 |
+
can bound the descriptor's quality but cannot choose the operating
|
| 708 |
+
point, and choosing it from one anyway is how the cut ended up where
|
| 709 |
+
it did.
|
| 710 |
+
|
| 711 |
+
The target quantity is measurable directly and needs no annotation:
|
| 712 |
+
HOW MANY OBJECTS RECUR ACROSS GROUPS. That is the whole point of a
|
| 713 |
+
persistent id, and it is self-limiting as an objective - too tight
|
| 714 |
+
and every sighting is a new id that recurs nowhere; too loose and
|
| 715 |
+
distinct objects collapse into one attractor, which also recurs
|
| 716 |
+
nowhere because there are no longer distinct objects to recur. It
|
| 717 |
+
therefore has an interior maximum, and that maximum is the fit.
|
| 718 |
+
|
| 719 |
+
Measured on fresh_bench (2,097 episodes, 73,521 intervals):
|
| 720 |
+
|
| 721 |
+
cut objects singletons RECUR worst false-merge
|
| 722 |
+
0.75 34,063 80.2% 6,447 0.18%
|
| 723 |
+
0.70 26,159 74.0% 6,475 0.22% <- peak
|
| 724 |
+
0.60 14,011 57.3% 5,732 0.54%
|
| 725 |
+
0.50 6,320 36.8% 3,876 1.05%
|
| 726 |
+
0.42 2,618 19.4% 2,075 1.72%
|
| 727 |
+
|
| 728 |
+
which also disposes of the singleton rate as a target: driving it
|
| 729 |
+
from 80% to 19% costs two thirds of the recurrence and multiplies
|
| 730 |
+
proven-wrong merges by eight. Fewer, bigger, wronger objects.
|
| 731 |
+
|
| 732 |
+
Sweeping a cut used to be impossible - Gallery.assign was ~20 minutes
|
| 733 |
+
per pass - so the fit had to be a closed-form statistic. It is now
|
| 734 |
+
seconds, and an objective beats a proxy.
|
| 735 |
+
|
| 736 |
+
Args:
|
| 737 |
+
V: (N, D) unit-norm descriptors.
|
| 738 |
+
groups: (N,) group label per descriptor - the episode it sits in.
|
| 739 |
+
neg: proven-different pairs, for the reported false-merge rate.
|
| 740 |
+
pos: proven-same pairs, for the reported recovery rate.
|
| 741 |
+
grid: candidate cuts. Defaults to percentiles of the proven-
|
| 742 |
+
different similarities, so the range adapts to the corpus
|
| 743 |
+
rather than being a hand-written span.
|
| 744 |
+
handoff: track-handoff proven-same pairs (handoff_pairs). When
|
| 745 |
+
supplied in bulk they CONSTRAIN the objective - see below.
|
| 746 |
+
Returns:
|
| 747 |
+
(cut, report) - report lists every candidate that was tried.
|
| 748 |
+
"""
|
| 749 |
+
g = np.asarray(groups)
|
| 750 |
+
g = g - g.min() + 1
|
| 751 |
+
span = int(g.max()) + 2
|
| 752 |
+
N = np.asarray(neg, np.int64).reshape(-1, 2)
|
| 753 |
+
Pp = np.asarray(pos if pos is not None else [], np.int64).reshape(-1, 2)
|
| 754 |
+
H = np.asarray(handoff if handoff is not None else [],
|
| 755 |
+
np.int64).reshape(-1, 2)
|
| 756 |
+
if grid is None:
|
| 757 |
+
s = (np.einsum("ij,ij->i", V[N[:, 0]], V[N[:, 1]]) if len(N)
|
| 758 |
+
else np.array([MATCH]))
|
| 759 |
+
grid = np.unique(np.round(np.percentile(
|
| 760 |
+
s, [90, 95, 97.5, 99, 99.3, 99.5, 99.7, 99.9]), 3))
|
| 761 |
+
|
| 762 |
+
def _row(cut):
|
| 763 |
+
ids = Gallery(match=float(cut)).assign(V)
|
| 764 |
+
n = int(ids.max()) + 1
|
| 765 |
+
# objects present in more than one group
|
| 766 |
+
u = np.unique(ids.astype(np.int64) * span + g)
|
| 767 |
+
per = np.bincount((u // span).astype(np.int64), minlength=n)
|
| 768 |
+
row = {"cut": float(cut), "objects": n,
|
| 769 |
+
"recur": int((per > 1).sum()),
|
| 770 |
+
"singleton_pct": round(100.0 * float(
|
| 771 |
+
(np.bincount(ids) == 1).mean()), 1)}
|
| 772 |
+
if len(N):
|
| 773 |
+
row["false_merge_pct"] = round(100.0 * float(
|
| 774 |
+
(ids[N[:, 0]] == ids[N[:, 1]]).mean()), 3)
|
| 775 |
+
if len(Pp):
|
| 776 |
+
row["recovered_pct"] = round(100.0 * float(
|
| 777 |
+
(ids[Pp[:, 0]] == ids[Pp[:, 1]]).mean()), 1)
|
| 778 |
+
if len(H):
|
| 779 |
+
row["handoff_pct"] = round(100.0 * float(
|
| 780 |
+
(ids[H[:, 0]] == ids[H[:, 1]]).mean()), 1)
|
| 781 |
+
return row
|
| 782 |
+
|
| 783 |
+
def _pick(rep):
|
| 784 |
+
# Recurrence alone is GAMEABLE BY FRAGMENTATION on a lookalike
|
| 785 |
+
# corpus: identical twins across episodes mean noise-clusters of
|
| 786 |
+
# one kind still span groups, so the count keeps rising as the
|
| 787 |
+
# cut climbs while real objects shatter (measured on the sim
|
| 788 |
+
# store: 21,492 objects for ~450 true blocks, 2.67 ids per
|
| 789 |
+
# block, 8% same-block event agreement). A cut that cannot
|
| 790 |
+
# re-find an object across a one-second track break identifies
|
| 791 |
+
# nothing - so when handoff pairs exist in bulk, continuity
|
| 792 |
+
# recovery is a CONSTRAINT on the recurrence objective, not a
|
| 793 |
+
# tiebreaker: maximise recurrence among cuts that recover >=70%
|
| 794 |
+
# of handoffs, and if none does, take the best recovery there is.
|
| 795 |
+
if len(H) >= 30:
|
| 796 |
+
ok = [r for r in rep if r.get("handoff_pct", 0.0) >= 70.0]
|
| 797 |
+
if ok:
|
| 798 |
+
return max(ok, key=lambda r: r["recur"])
|
| 799 |
+
return max(rep, key=lambda r: (r.get("handoff_pct", 0.0),
|
| 800 |
+
r["recur"]))
|
| 801 |
+
return max(rep, key=lambda r: r["recur"])
|
| 802 |
+
|
| 803 |
+
report = [_row(cut) for cut in grid]
|
| 804 |
+
best = _pick(report)
|
| 805 |
+
# The grid is percentiles of the proven-DIFFERENT similarities, and
|
| 806 |
+
# that anchor breaks when the descriptor is much better than the
|
| 807 |
+
# negatives are hard: recurrence keeps rising past the negatives'
|
| 808 |
+
# entire tail, the argmax lands on the grid's top edge, and the
|
| 809 |
+
# "interior maximum" was never actually bracketed. Measured on the
|
| 810 |
+
# DINOv3 rebuild: recur still climbing 7,590 -> 10,541 at the last
|
| 811 |
+
# point. So while the best cut IS the top edge, keep extending
|
| 812 |
+
# upward (midpoint steps toward 0.99) until the maximum is interior
|
| 813 |
+
# or the ceiling is reached - the fit must end bracketed, not
|
| 814 |
+
# truncated by an artifact of where the negatives happened to end.
|
| 815 |
+
while best["cut"] == report[-1]["cut"] and best["cut"] < 0.985:
|
| 816 |
+
report.append(_row(round(best["cut"]
|
| 817 |
+
+ (0.99 - best["cut"]) / 2, 3)))
|
| 818 |
+
nxt = _pick(report)
|
| 819 |
+
if nxt["cut"] == best["cut"]:
|
| 820 |
+
break
|
| 821 |
+
best = nxt
|
| 822 |
+
return best["cut"], report
|
| 823 |
+
|
| 824 |
+
|
| 825 |
+
class Gallery:
|
| 826 |
+
"""Persistent identities, matched by appearance.
|
| 827 |
+
|
| 828 |
+
A GALLERY of exemplars per object, not a running mean. The mean is
|
| 829 |
+
what broke the earlier attempt: a centroid that updates as it
|
| 830 |
+
absorbs members lets a group CHAIN - green, to a slightly different
|
| 831 |
+
green, to red - so two objects merge through a path of intermediate
|
| 832 |
+
views. Several fixed exemplars, matched against the best of them,
|
| 833 |
+
has no such path. This is ordinary ReID gallery practice.
|
| 834 |
+
|
| 835 |
+
What it is asked about changed: entries are now TRACK descriptors
|
| 836 |
+
pooled over many frames, not single crops, so each question carries
|
| 837 |
+
a whole sighting's worth of evidence.
|
| 838 |
+
"""
|
| 839 |
+
|
| 840 |
+
def __init__(self, match=MATCH, max_ex=MAX_EXEMPLARS, admit=ADMIT):
|
| 841 |
+
self.match, self.max_ex, self.admit = match, max_ex, admit
|
| 842 |
+
self.ex: list[np.ndarray] = []
|
| 843 |
+
self.n: list[int] = []
|
| 844 |
+
# every exemplar of every object in ONE matrix, plus which object
|
| 845 |
+
# owns each row. self.ex stays the public view of the same thing.
|
| 846 |
+
self._M = np.zeros((0, 0), np.float32)
|
| 847 |
+
self._own = np.zeros(0, np.int32)
|
| 848 |
+
self._m = 0
|
| 849 |
+
|
| 850 |
+
def _grow(self, dim):
|
| 851 |
+
if self._M.shape[0] > self._m:
|
| 852 |
+
return
|
| 853 |
+
cap = max(1024, self._M.shape[0] * 2)
|
| 854 |
+
M = np.zeros((cap, dim), np.float32)
|
| 855 |
+
if self._m:
|
| 856 |
+
M[:self._m] = self._M[:self._m]
|
| 857 |
+
own = np.zeros(cap, np.int32)
|
| 858 |
+
own[:self._m] = self._own[:self._m]
|
| 859 |
+
self._M, self._own = M, own
|
| 860 |
+
|
| 861 |
+
def assign(self, feats):
|
| 862 |
+
"""Which object is each descriptor, in arrival order.
|
| 863 |
+
|
| 864 |
+
ONE matrix-vector product per descriptor, not one per object.
|
| 865 |
+
The old form looped over self.ex in Python, so a 73,521-track
|
| 866 |
+
corpus that had grown 32,756 objects did 2.4 billion iterations
|
| 867 |
+
of a two-line body - about 20 of the write's 54 minutes, and slow
|
| 868 |
+
enough that the cut could not be swept to find out it was wrong.
|
| 869 |
+
|
| 870 |
+
The reduction is exact, not an approximation: the best object is
|
| 871 |
+
the one owning the best EXEMPLAR, because max over objects of
|
| 872 |
+
(max over that object's exemplars) IS the max over all exemplars.
|
| 873 |
+
So a single argmax replaces the per-object max-then-compare.
|
| 874 |
+
"""
|
| 875 |
+
ids = []
|
| 876 |
+
for v in feats:
|
| 877 |
+
v = np.ascontiguousarray(v, np.float32)
|
| 878 |
+
best, bs = -1, -1.0
|
| 879 |
+
if self._m:
|
| 880 |
+
s = self._M[:self._m] @ v
|
| 881 |
+
k = int(s.argmax())
|
| 882 |
+
best, bs = int(self._own[k]), float(s[k])
|
| 883 |
+
if best >= 0 and bs >= self.match:
|
| 884 |
+
ids.append(best)
|
| 885 |
+
self.n[best] += 1
|
| 886 |
+
E = self.ex[best]
|
| 887 |
+
# keep a view only if it ADDS one; near-duplicates of a
|
| 888 |
+
# stored exemplar teach the gallery nothing. Admit it
|
| 889 |
+
# only if it agrees with EVERY held exemplar - min(),
|
| 890 |
+
# not max() - so the acceptance region cannot sprawl.
|
| 891 |
+
if (len(E) < self.max_ex and bs < 0.95
|
| 892 |
+
and float((E @ v).min()) >= self.admit):
|
| 893 |
+
self.ex[best] = np.vstack([E, v])
|
| 894 |
+
self._grow(len(v))
|
| 895 |
+
self._M[self._m] = v
|
| 896 |
+
self._own[self._m] = best
|
| 897 |
+
self._m += 1
|
| 898 |
+
else:
|
| 899 |
+
self.ex.append(v[None, :])
|
| 900 |
+
self.n.append(1)
|
| 901 |
+
best = len(self.ex) - 1
|
| 902 |
+
ids.append(best)
|
| 903 |
+
self._grow(len(v))
|
| 904 |
+
self._M[self._m] = v
|
| 905 |
+
self._own[self._m] = best
|
| 906 |
+
self._m += 1
|
| 907 |
+
return np.asarray(ids, np.int32)
|
| 908 |
+
|
| 909 |
+
def centroids(self):
|
| 910 |
+
return np.stack([E.mean(0) / (np.linalg.norm(E.mean(0)) + 1e-8)
|
| 911 |
+
for E in self.ex]) if self.ex else np.zeros(
|
| 912 |
+
(0, 512), np.float32)
|
| 913 |
+
|
| 914 |
+
def __len__(self):
|
| 915 |
+
return len(self.ex)
|
python/elidedb/ingest.py
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Ingest primitives: the shared surface the write scripts build on.
|
| 2 |
+
|
| 3 |
+
WHY THIS MODULE EXISTS
|
| 4 |
+
----------------------
|
| 5 |
+
`scripts/full_write.py` imported nine names from `scripts/write_once.py`,
|
| 6 |
+
which is a SCRIPT with its own `main()`. That is the wrong dependency
|
| 7 |
+
direction: a runnable entry point had become a library, so importing it
|
| 8 |
+
pulled in argv parsing, module-level side effects and a second CLI, and
|
| 9 |
+
neither file could be changed without checking the other.
|
| 10 |
+
|
| 11 |
+
The same shape appears seven more times across `scripts/` - files
|
| 12 |
+
importing `bench_product` for a constant, `extract_events` for two
|
| 13 |
+
functions. Entry points should depend on the package; the package should
|
| 14 |
+
never depend on an entry point.
|
| 15 |
+
|
| 16 |
+
So the durable pieces live here, in `elidedb`, and the scripts become
|
| 17 |
+
thin: parse arguments, call in, report. Anything with a `main()` is a
|
| 18 |
+
program. Anything imported is a module. Not both.
|
| 19 |
+
|
| 20 |
+
WHAT IS HERE
|
| 21 |
+
------------
|
| 22 |
+
Corpus geometry (where the bytes are, what the clock is) and the two
|
| 23 |
+
loops every writer needs: streaming a file once, and cutting a demo's
|
| 24 |
+
frames out of a continuous timeline.
|
| 25 |
+
|
| 26 |
+
Dataset LAYOUT belongs in `corpus.json` (see corpus.py) rather than in
|
| 27 |
+
constants here; the constants below are the Bridge defaults kept so the
|
| 28 |
+
existing scripts keep working, and each is overridable.
|
| 29 |
+
"""
|
| 30 |
+
from __future__ import annotations
|
| 31 |
+
|
| 32 |
+
import subprocess
|
| 33 |
+
|
| 34 |
+
import numpy as np
|
| 35 |
+
|
| 36 |
+
from .fftools import find
|
| 37 |
+
|
| 38 |
+
# Bridge defaults. A different corpus supplies these through corpus.json
|
| 39 |
+
# rather than by editing a constant - see corpus.Corpus.
|
| 40 |
+
CAM = "observation.images.image_0"
|
| 41 |
+
EPOCH_NS = 1_704_067_200_000_000_000
|
| 42 |
+
FILE_STRIDE_NS = 20_000_000_000_000
|
| 43 |
+
FPS = 5.0
|
| 44 |
+
NGEOM = 12 # frames per episode handed to geometry
|
| 45 |
+
GEOM_W, GEOM_H = 256, 192 # geometry runs on its own small raster
|
| 46 |
+
EMBED_W, EMBED_H = 192, 144 # FDNN-V's raster
|
| 47 |
+
GAP_S = 60.0 # silence inserted between demos
|
| 48 |
+
CRF = 26 # per-demo H.264 quality
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def gapped(spans, gap_s=GAP_S):
|
| 52 |
+
"""Uniform gap between consecutive demos of a stream.
|
| 53 |
+
|
| 54 |
+
Without it two adjacent demos are contiguous in time and a window
|
| 55 |
+
query cannot express "this demo and not the next one". Returns
|
| 56 |
+
{(stream, t0): shift_ns}.
|
| 57 |
+
"""
|
| 58 |
+
gap = int(gap_s * 1e9)
|
| 59 |
+
prev, shift = {}, {}
|
| 60 |
+
for s in sorted(spans, key=lambda r: (r["stream"], r["t0"])):
|
| 61 |
+
st = s["stream"]
|
| 62 |
+
new = s["t0"] if st not in prev else prev[st] + gap
|
| 63 |
+
shift[(st, s["t0"])] = new - s["t0"]
|
| 64 |
+
prev[st] = new + (s["t1"] - s["t0"])
|
| 65 |
+
return shift
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def frame_owner(spans, base, fps=FPS):
|
| 69 |
+
"""{frame index -> episode id} for one file.
|
| 70 |
+
|
| 71 |
+
An episode owns EXACTLY the n frames the corpus declares. Deriving
|
| 72 |
+
the end by rounding t1 claims one extra frame at an inclusive bound -
|
| 73 |
+
verified against the seek-based writer, which produced 30 frames for
|
| 74 |
+
an episode where rounding produced 31, every other frame identical.
|
| 75 |
+
Trust the declared length, not a rounded timestamp.
|
| 76 |
+
"""
|
| 77 |
+
owner = {}
|
| 78 |
+
for e in spans:
|
| 79 |
+
i0 = int(round((e["t0"] - base) / 1e9 * fps))
|
| 80 |
+
for i in range(i0, i0 + int(e["n"])):
|
| 81 |
+
owner[i] = e["episode"]
|
| 82 |
+
return owner
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def decode_stream(path, width=None, height=None, fps=FPS, chunk_frames=16):
|
| 86 |
+
"""Decode a video ONCE, yielding raw RGB frames in batches.
|
| 87 |
+
|
| 88 |
+
The write path decoded the same pixels three times - once to cut
|
| 89 |
+
per-demo segments, once for the embedding raster, once again from
|
| 90 |
+
the store's own segments for region proposal. Measured, the third
|
| 91 |
+
was 41% of the presence stage and the second was the entire embed
|
| 92 |
+
stage. Everything that needs pixels should ride ONE decode, which is
|
| 93 |
+
what this yields.
|
| 94 |
+
|
| 95 |
+
`width`/`height` None keeps the source resolution: consumers that
|
| 96 |
+
want a smaller raster resize the frames they are given, rather than
|
| 97 |
+
each opening its own decoder.
|
| 98 |
+
"""
|
| 99 |
+
vf = [f"fps={fps}"]
|
| 100 |
+
if width and height:
|
| 101 |
+
vf.append(f"scale={width}:{height}")
|
| 102 |
+
if width is None or height is None:
|
| 103 |
+
w, h = probe_size(path)
|
| 104 |
+
width, height = width or w, height or h
|
| 105 |
+
fb = width * height * 3
|
| 106 |
+
proc = subprocess.Popen(
|
| 107 |
+
[find("ffmpeg"), "-v", "error", "-i", str(path),
|
| 108 |
+
"-vf", ",".join(vf), "-f", "rawvideo", "-pix_fmt", "rgb24",
|
| 109 |
+
"pipe:1"], stdout=subprocess.PIPE, bufsize=fb * chunk_frames)
|
| 110 |
+
buf = b""
|
| 111 |
+
try:
|
| 112 |
+
while True:
|
| 113 |
+
data = proc.stdout.read(fb * chunk_frames - len(buf))
|
| 114 |
+
if data:
|
| 115 |
+
buf += data
|
| 116 |
+
n = len(buf) // fb
|
| 117 |
+
if n == 0 and not data:
|
| 118 |
+
break
|
| 119 |
+
if n == 0:
|
| 120 |
+
continue
|
| 121 |
+
yield np.frombuffer(buf[:n * fb], np.uint8).reshape(
|
| 122 |
+
n, height, width, 3)
|
| 123 |
+
buf = buf[n * fb:]
|
| 124 |
+
finally:
|
| 125 |
+
proc.stdout.close()
|
| 126 |
+
proc.wait()
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def probe_size(path):
|
| 130 |
+
"""(width, height) of a video, without decoding it."""
|
| 131 |
+
r = subprocess.run(
|
| 132 |
+
[find("ffprobe"), "-v", "error", "-select_streams", "v:0",
|
| 133 |
+
"-show_entries", "stream=width,height", "-of", "csv=p=0",
|
| 134 |
+
str(path)], capture_output=True, text=True, check=True)
|
| 135 |
+
w, h = (int(x) for x in r.stdout.strip().split(",")[:2])
|
| 136 |
+
return w, h
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def encoder(path, width, height, fps=FPS, crf=CRF):
|
| 140 |
+
"""An ffmpeg process that takes RAW FRAMES on stdin.
|
| 141 |
+
|
| 142 |
+
The store needs per-demo H.264 with exactly one IDR so a 2 s read is
|
| 143 |
+
a byte range. That encode is unavoidable; the DECODE inside it is
|
| 144 |
+
not - feeding already-decoded frames removes a whole pass over the
|
| 145 |
+
source.
|
| 146 |
+
|
| 147 |
+
-bf 0 keeps packet order equal to presentation order, which the
|
| 148 |
+
frame index depends on; -g/-keyint_min large with -sc_threshold 0
|
| 149 |
+
guarantees the single IDR.
|
| 150 |
+
"""
|
| 151 |
+
return subprocess.Popen(
|
| 152 |
+
[find("ffmpeg"), "-v", "error", "-y",
|
| 153 |
+
"-f", "rawvideo", "-pix_fmt", "rgb24",
|
| 154 |
+
"-s", f"{width}x{height}", "-r", str(fps), "-i", "pipe:0",
|
| 155 |
+
"-an", "-c:v", "libx264", "-preset", "medium", "-crf", str(crf),
|
| 156 |
+
"-bf", "0", "-g", "10000", "-keyint_min", "10000",
|
| 157 |
+
"-sc_threshold", "0", "-f", "h264", str(path)],
|
| 158 |
+
stdin=subprocess.PIPE)
|
python/elidedb/itm.py
ADDED
|
@@ -0,0 +1,175 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""ITM cross-encoder rerank — the discarded head, wired for live queries.
|
| 2 |
+
|
| 3 |
+
InternVideo2's checkpoint ships an `itm_head` that every load discarded
|
| 4 |
+
("Some weights ... were not used"). It is the cross-encoder half of the
|
| 5 |
+
model's own retrieval recipe (BLIP-2 arXiv 2301.12597): text tokens
|
| 6 |
+
cross-attend to video tokens, a 2-class head reads the fused CLS. It
|
| 7 |
+
out-ranks everything else in this store — measured at k=1.5xsupport,
|
| 8 |
+
ITM alone 0.38/0.25 against the cosine ensemble's 0.32/0.21 and the
|
| 9 |
+
shipped RRF path's 0.29/0.23.
|
| 10 |
+
|
| 11 |
+
Two measurements decide the shape of this module:
|
| 12 |
+
|
| 13 |
+
NO POOLING. Vision tokens are 1,025 x 1,408 per episode = 3.24 GB
|
| 14 |
+
corpus-wide, four times the raw source, so storing them in the store
|
| 15 |
+
is out. Pooling them is worse than out: 16x16 -> 8x8 per frame drops
|
| 16 |
+
Spearman against the full-token score to 0.18, and to NEGATIVE on two
|
| 17 |
+
of three probe queries. Cross-attention does not tolerate a reduced
|
| 18 |
+
key set. So the tokens are a disposable CACHE outside the store's
|
| 19 |
+
byte budget, never a table.
|
| 20 |
+
|
| 21 |
+
CASCADE, NOT SCAN. Reranking the top-N of the cheap fused ranking
|
| 22 |
+
reaches the full-scan number exactly: N=500 gives 0.40/0.27, N=150
|
| 23 |
+
already gives 0.38/0.25, and every query with support under ~200 is
|
| 24 |
+
saturated at N=150. So N scales with the ceiling k rather than the
|
| 25 |
+
corpus, which is L7 ("ITM is evidence inside a candidate set, not
|
| 26 |
+
authority over the corpus") arrived at from the cost side.
|
| 27 |
+
|
| 28 |
+
ITM IS A STAGE, NOT A CHANNEL — and that distinction was measured, not
|
| 29 |
+
chosen. Making it a weighted RRF voter and refitting cost 0.38 -> 0.27
|
| 30 |
+
for two compounding reasons: RRF converts every channel to RANKS, which
|
| 31 |
+
throws away the logit margin's scale, and that scale is where the
|
| 32 |
+
cross-encoder's separation lives (the 0.40 offline number came from
|
| 33 |
+
adding z-scores, not ranks); and the k-ladder refit then optimized the
|
| 34 |
+
ladder mean away from the 1.5x-support operating point the product
|
| 35 |
+
metric uses. So ITM stays a rerank STAGE alongside NMS and the
|
| 36 |
+
confidence cut - stages the fitter switches on and off rather than
|
| 37 |
+
weights - and the fit models it as such.
|
| 38 |
+
|
| 39 |
+
The cache is content-addressed by (episode, model) and capped; deleting
|
| 40 |
+
it costs time, never correctness.
|
| 41 |
+
"""
|
| 42 |
+
from __future__ import annotations
|
| 43 |
+
|
| 44 |
+
import hashlib
|
| 45 |
+
from pathlib import Path
|
| 46 |
+
|
| 47 |
+
import numpy as np
|
| 48 |
+
|
| 49 |
+
_S: dict = {}
|
| 50 |
+
# 3.24 GB holds the whole 1,122-episode corpus. This is a DISPOSABLE
|
| 51 |
+
# cache under the store's _cache/, not a table: it never enters the
|
| 52 |
+
# store's byte ledger, and deleting it costs recompute time, never
|
| 53 |
+
# correctness. Override with ELIDEDB_ITM_CACHE_GB.
|
| 54 |
+
CACHE_CAP_GB = float(__import__("os").environ.get(
|
| 55 |
+
"ELIDEDB_ITM_CACHE_GB", "4.0"))
|
| 56 |
+
NF = 4
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def _head():
|
| 60 |
+
"""The checkpoint's own 2-class matcher — AutoModel drops it because
|
| 61 |
+
the vendored class never declares the attribute."""
|
| 62 |
+
if "head" in _S:
|
| 63 |
+
return _S["head"], _S["m"], _S["dev"]
|
| 64 |
+
import torch
|
| 65 |
+
|
| 66 |
+
from .iv2 import MDIR, load_model
|
| 67 |
+
sd = torch.load(f"{MDIR}/pytorch_model.bin", map_location="cpu",
|
| 68 |
+
weights_only=True, mmap=True)
|
| 69 |
+
m, dev = load_model()
|
| 70 |
+
h = torch.nn.Linear(sd["itm_head.weight"].shape[1], 2)
|
| 71 |
+
h.load_state_dict({"weight": sd["itm_head.weight"],
|
| 72 |
+
"bias": sd["itm_head.bias"]})
|
| 73 |
+
h = h.to(dev, m.dtype).eval()
|
| 74 |
+
_S["head"], _S["m"], _S["dev"] = h, m, dev
|
| 75 |
+
return h, m, dev
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def _cache_dir(store):
|
| 79 |
+
d = Path(store.dir) / "_cache" / "itm_tokens"
|
| 80 |
+
d.mkdir(parents=True, exist_ok=True)
|
| 81 |
+
return d
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def _cache_path(store, stream, ts):
|
| 85 |
+
key = hashlib.sha1(f"{stream}|{ts}|iv2-1b|{NF}".encode()).hexdigest()
|
| 86 |
+
return _cache_dir(store) / f"{key}.npy"
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def _evict(store):
|
| 90 |
+
"""Keep the disposable cache under CACHE_CAP_GB, oldest first."""
|
| 91 |
+
d = _cache_dir(store)
|
| 92 |
+
files = sorted(d.glob("*.npy"), key=lambda p: p.stat().st_mtime)
|
| 93 |
+
total = sum(p.stat().st_size for p in files)
|
| 94 |
+
cap = CACHE_CAP_GB * 1e9
|
| 95 |
+
while total > cap and files:
|
| 96 |
+
p = files.pop(0)
|
| 97 |
+
total -= p.stat().st_size
|
| 98 |
+
try:
|
| 99 |
+
p.unlink()
|
| 100 |
+
except OSError:
|
| 101 |
+
pass
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
def _vision_tokens(store, frames_tbl, key):
|
| 105 |
+
"""Cached (1, T, C) vision tokens for one episode."""
|
| 106 |
+
import cv2
|
| 107 |
+
import torch
|
| 108 |
+
|
| 109 |
+
from .iv2 import V_MEAN, V_STD
|
| 110 |
+
from .video import FrameSet
|
| 111 |
+
s, a, b = key
|
| 112 |
+
p = _cache_path(store, s, a)
|
| 113 |
+
_, m, dev = _head()
|
| 114 |
+
if p.exists():
|
| 115 |
+
arr = np.load(p)
|
| 116 |
+
return torch.from_numpy(arr).to(dev, m.dtype)
|
| 117 |
+
import pyarrow.compute as pc
|
| 118 |
+
sel = frames_tbl.filter(pc.and_(
|
| 119 |
+
pc.equal(frames_tbl.column("stream"), s),
|
| 120 |
+
pc.and_(pc.greater_equal(frames_tbl.column("ts"), a),
|
| 121 |
+
pc.less_equal(frames_tbl.column("ts"), b))))
|
| 122 |
+
if len(sel) < 2:
|
| 123 |
+
return None
|
| 124 |
+
pi = np.linspace(0, len(sel) - 1, min(NF, len(sel))).round().astype(int)
|
| 125 |
+
try:
|
| 126 |
+
dec = FrameSet(store, "frames", sel.take(pi)).decode(width=224)
|
| 127 |
+
except Exception:
|
| 128 |
+
return None
|
| 129 |
+
fr = [f for _, f in sorted(dec)]
|
| 130 |
+
if len(fr) < 2:
|
| 131 |
+
return None
|
| 132 |
+
fs = [cv2.resize(f, (224, 224)) for f in fr]
|
| 133 |
+
x = (np.stack(fs).astype(np.float32) / 255.0 - V_MEAN) / V_STD
|
| 134 |
+
px = torch.from_numpy(x).permute(0, 3, 1, 2)[None].to(dev, m.dtype)
|
| 135 |
+
with torch.no_grad():
|
| 136 |
+
vis, _ = m.encode_vision(px, test=True)
|
| 137 |
+
np.save(p, vis.cpu().numpy())
|
| 138 |
+
_evict(store)
|
| 139 |
+
return vis
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
def itm_scores(store, text, keys):
|
| 143 |
+
"""Logit margin P(match) - P(no match) for each candidate episode.
|
| 144 |
+
|
| 145 |
+
Returns np.array aligned with `keys`; NaN where frames are
|
| 146 |
+
undecodable. The vision pass dominates and is query-independent, so
|
| 147 |
+
a repeated query over the same candidates is nearly free."""
|
| 148 |
+
import torch
|
| 149 |
+
head, m, dev = _head()
|
| 150 |
+
frames_tbl = store.table("frames").scan()
|
| 151 |
+
tok = m.tokenizer(text, padding="max_length", truncation=True,
|
| 152 |
+
max_length=m._config.max_txt_l,
|
| 153 |
+
return_tensors="pt").to(dev)
|
| 154 |
+
out = np.full(len(keys), np.nan, np.float32)
|
| 155 |
+
for i, k in enumerate(keys):
|
| 156 |
+
vis = _vision_tokens(store, frames_tbl, k)
|
| 157 |
+
if vis is None:
|
| 158 |
+
continue
|
| 159 |
+
with torch.no_grad():
|
| 160 |
+
vam = torch.ones(vis.shape[:2], dtype=torch.long, device=dev)
|
| 161 |
+
o = m.get_text_encoder()(
|
| 162 |
+
tok.input_ids, attention_mask=tok.attention_mask,
|
| 163 |
+
encoder_hidden_states=vis, encoder_attention_mask=vam,
|
| 164 |
+
return_dict=True, mode="multi_modal")
|
| 165 |
+
lg = head(o.last_hidden_state[:, 0]).float()[0]
|
| 166 |
+
out[i] = float(lg[1] - lg[0])
|
| 167 |
+
return out
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def rerank_depth(k_max, n_total):
|
| 171 |
+
"""Candidates to rerank. Measured: N=150 saturates every query with
|
| 172 |
+
support under ~200; only the 247-support query needed 500. Scaling
|
| 173 |
+
with the ceiling rather than the corpus keeps the cost proportional
|
| 174 |
+
to what the caller actually asked for."""
|
| 175 |
+
return int(min(max(150, 2 * k_max), n_total))
|
python/elidedb/iv2.py
ADDED
|
@@ -0,0 +1,133 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""InternVideo2-Stage2 1B (Wang et al., ECCV 2024, arXiv 2403.15377)
|
| 2 |
+
— the only true VIDEO-native text channel: stage-2 trains video-text
|
| 3 |
+
contrastively WITH temporal modeling, which the frame-pooled image
|
| 4 |
+
channels (PE, SigLIP2, X-CLIP pooling) structurally lack. Distilled
|
| 5 |
+
small variants are reported much weaker on retrieval — 1B or nothing.
|
| 6 |
+
|
| 7 |
+
One 512-d aligned vector per episode (4-frame clip, the f4
|
| 8 |
+
checkpoint's native temporal extent), table iv2_vectors. Weights:
|
| 9 |
+
ziyjiang/InternVideo2-1B (fp16 re-serialization of the official
|
| 10 |
+
OpenGVLab stage2 checkpoint, towers only); modeling code vendored
|
| 11 |
+
from VLM2Vec into models/iv2_stage2_1b by scripts/get_iv2.py (three
|
| 12 |
+
MPS patches: flash-attn import guarded, LayerScale gamma naming so
|
| 13 |
+
the checkpoint's 80 layer scales actually load, bert config resolved
|
| 14 |
+
next to the file)."""
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
|
| 19 |
+
_S = {}
|
| 20 |
+
|
| 21 |
+
MDIR = "models/iv2_stage2_1b"
|
| 22 |
+
|
| 23 |
+
# InternVideo2's own demo preprocessing (frames2tensor): ImageNet
|
| 24 |
+
# stats, 224 square, [B,T,C,H,W]
|
| 25 |
+
V_MEAN = np.array([0.485, 0.456, 0.406], np.float32)
|
| 26 |
+
V_STD = np.array([0.229, 0.224, 0.225], np.float32)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def load_model():
|
| 30 |
+
if "model" not in _S:
|
| 31 |
+
from transformers import AutoModel
|
| 32 |
+
|
| 33 |
+
from .device import pick, strip_vision
|
| 34 |
+
dev, dtype = pick()
|
| 35 |
+
import transformers
|
| 36 |
+
kw = dict(trust_remote_code=True, torch_dtype=dtype)
|
| 37 |
+
if int(transformers.__version__.split(".")[0]) < 5:
|
| 38 |
+
# the 8GB-container path; transformers>=5 meta-device init
|
| 39 |
+
# breaks this custom port's from_pretrained, and the fp16
|
| 40 |
+
# towers load fine without it there
|
| 41 |
+
kw["low_cpu_mem_usage"] = True
|
| 42 |
+
m = AutoModel.from_pretrained(MDIR, **kw).to(dev).eval()
|
| 43 |
+
m = strip_vision(m, "vision_encoder")
|
| 44 |
+
m._config.device = dev # get_txt_feat routes tokens here
|
| 45 |
+
_S["model"], _S["dev"], _S["dtype"] = m, dev, dtype
|
| 46 |
+
return _S["model"], _S["dev"]
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def text_vec(text):
|
| 50 |
+
cache = _S.setdefault("tcache", {})
|
| 51 |
+
if text in cache:
|
| 52 |
+
return cache[text]
|
| 53 |
+
m, _ = load_model()
|
| 54 |
+
v = m.get_txt_feat(text).float().cpu().numpy().reshape(-1)
|
| 55 |
+
if len(cache) > 256:
|
| 56 |
+
cache.clear()
|
| 57 |
+
cache[text] = v
|
| 58 |
+
return v
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def set_num_frames(n):
|
| 62 |
+
"""Re-fit the vision encoder's temporal position embeddings to n
|
| 63 |
+
frames. OPT-IN and global to the loaded model.
|
| 64 |
+
|
| 65 |
+
The checkpoint ships a 4-frame clip: pos_embed is (1, 1 + 4*256, C),
|
| 66 |
+
and feeding 8 frames raises "size of tensor a (2049) must match
|
| 67 |
+
tensor b (1025)". That is a shape, not a capability - the patch
|
| 68 |
+
embedding is Conv3d with a kernel of (1,14,14), so tubelet size is
|
| 69 |
+
1 and T frames simply need T temporal positions. The repo's own
|
| 70 |
+
`interpolate_pos_embed(orig_t_size=4)` does exactly this, but only
|
| 71 |
+
while loading a checkpoint.
|
| 72 |
+
|
| 73 |
+
FOUR embeddings need it, not one: pos_embed AND clip_pos_embed both
|
| 74 |
+
carry the video-length grid (the image variants are separate and
|
| 75 |
+
untouched). Interpolating only the first fails deeper in the
|
| 76 |
+
forward, at the CLIP-alignment branch.
|
| 77 |
+
|
| 78 |
+
Sanity after interpolation: cos(4-frame, 8-frame) on the same span
|
| 79 |
+
is 0.999, i.e. the representation is preserved rather than rebuilt.
|
| 80 |
+
"""
|
| 81 |
+
import torch
|
| 82 |
+
m, _ = load_model()
|
| 83 |
+
ve = m.vision_encoder
|
| 84 |
+
old = int(ve.num_frames)
|
| 85 |
+
if old == n:
|
| 86 |
+
return
|
| 87 |
+
def _interp(p):
|
| 88 |
+
cls, rest = p[:, :1, :], p[:, 1:, :]
|
| 89 |
+
C = rest.shape[-1]
|
| 90 |
+
L = rest.shape[1] // old
|
| 91 |
+
r = rest.view(1, old, L, C).permute(0, 3, 2, 1).float()
|
| 92 |
+
r = torch.nn.functional.interpolate(
|
| 93 |
+
r, size=(L, n), mode="bilinear", align_corners=False)
|
| 94 |
+
r = r.permute(0, 3, 2, 1).reshape(1, n * L, C).to(p.dtype)
|
| 95 |
+
return torch.nn.Parameter(torch.cat([cls, r], 1),
|
| 96 |
+
requires_grad=False)
|
| 97 |
+
L = (ve.pos_embed.shape[1] - 1) // old
|
| 98 |
+
for name in ("pos_embed", "clip_pos_embed"):
|
| 99 |
+
if hasattr(ve, name):
|
| 100 |
+
setattr(ve, name, _interp(getattr(ve, name).data))
|
| 101 |
+
ve.num_frames = n
|
| 102 |
+
ve.patch_embed.num_patches = n * L
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def clip_vec(frames_hwc):
|
| 106 |
+
"""Aligned 512-d vector for a clip (HWC uint8 RGB).
|
| 107 |
+
|
| 108 |
+
Length must match the encoder's current num_frames - 4 by default,
|
| 109 |
+
or whatever set_num_frames() last fitted."""
|
| 110 |
+
import cv2
|
| 111 |
+
import torch
|
| 112 |
+
m, dev = load_model()
|
| 113 |
+
fs = [cv2.resize(f, (224, 224)) for f in frames_hwc]
|
| 114 |
+
x = (np.stack(fs).astype(np.float32) / 255.0 - V_MEAN) / V_STD
|
| 115 |
+
px = torch.from_numpy(x).permute(0, 3, 1, 2)[None].to(
|
| 116 |
+
dev, _S["dtype"])
|
| 117 |
+
return m.get_vid_feat(px).float().cpu().numpy().reshape(-1)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def iv2_lookup(store, text):
|
| 121 |
+
from .embeddings import _vec_table
|
| 122 |
+
tbl, vecs = _vec_table(store, "iv2_vectors")
|
| 123 |
+
key = {}
|
| 124 |
+
for r, (s, a) in enumerate(zip(
|
| 125 |
+
tbl.column("stream").to_pylist(),
|
| 126 |
+
(int(v) for v in tbl.column("ts").to_pylist()))):
|
| 127 |
+
key[(str(s), a)] = r
|
| 128 |
+
sc = np.asarray(vecs) @ text_vec(text)
|
| 129 |
+
|
| 130 |
+
def lookup(s, a, b):
|
| 131 |
+
r = key.get((str(s), a))
|
| 132 |
+
return float(sc[r]) if r is not None else float("nan")
|
| 133 |
+
return lookup, None
|
python/elidedb/lexicon.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Closed-class English lexicon shared by query routing and training.
|
| 2 |
+
|
| 3 |
+
Dependency-free ON PURPOSE: the query path imports this (via
|
| 4 |
+
rerank.directional_swap), and it previously reached the same tables
|
| 5 |
+
through fdnnv2, whose module level imports mlx. That worked on a Mac
|
| 6 |
+
and crashed every Linux deployment. Lexicon is language knowledge, not
|
| 7 |
+
model code; it lives where nothing heavier than a list comprehension
|
| 8 |
+
runs.
|
| 9 |
+
|
| 10 |
+
Verb/direction swaps: GENERIC english antonyms, no dataset nouns, per
|
| 11 |
+
the no-hardwire rule.
|
| 12 |
+
"""
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
# ---------------------------------------------------------------------
|
| 16 |
+
# DERIVED OPPOSITES. The 81 pairs below and the 16 `un-` bases were
|
| 17 |
+
# hand-authored, and fold/stack/cover are conspicuously this corpus's
|
| 18 |
+
# vocabulary - a task prior in a file that claims to hold generic
|
| 19 |
+
# English. The swap-contrast MECHANISM is sound and measured; the list
|
| 20 |
+
# was not derived.
|
| 21 |
+
#
|
| 22 |
+
# derived_swaps() replaces it: opposites are found by reflection in an
|
| 23 |
+
# embedding space over the vocabulary the CORPUS attested. On a driving
|
| 24 |
+
# log that finds accelerate/brake without anyone having thought of them;
|
| 25 |
+
# here it finds open/close. The hand list stays only as a cold-start
|
| 26 |
+
# fallback for an empty store, and is labelled as such rather than
|
| 27 |
+
# presented as knowledge.
|
| 28 |
+
# ---------------------------------------------------------------------
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def derived_swaps(store, vec=None, min_count=2):
|
| 32 |
+
"""Opposite pairs this corpus supports, or () if it cannot say."""
|
| 33 |
+
from .derive import attested, opposite_pairs
|
| 34 |
+
try:
|
| 35 |
+
if "labels" not in store.tables():
|
| 36 |
+
return ()
|
| 37 |
+
vals = store.table("labels").scan().column("value").to_pylist()
|
| 38 |
+
vocab = attested(vals, min_count=min_count)
|
| 39 |
+
if len(vocab) < 8:
|
| 40 |
+
return ()
|
| 41 |
+
if vec is None:
|
| 42 |
+
from .sig2 import _text_vec as vec
|
| 43 |
+
return tuple((a, b) for a, b, _ in opposite_pairs(vocab, vec))
|
| 44 |
+
except Exception:
|
| 45 |
+
return ()
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
# THERE IS NO VERB TABLE. VERB_SWAPS (17 hand-authored pairs) and
|
| 49 |
+
# _UN_BASES (16 hand-authored stems that generated ~64 more) lived here
|
| 50 |
+
# as a "cold-start fallback", labelled a violation by their own comment
|
| 51 |
+
# and kept anyway. A fallback that always fires is not a fallback: on
|
| 52 |
+
# every store measured, derived_swaps returned 0 pairs, so the hand list
|
| 53 |
+
# WAS the direction mechanism, and it was also feeding the student's
|
| 54 |
+
# training negatives.
|
| 55 |
+
#
|
| 56 |
+
# Removed. derived_swaps() above is the only source of oppositions: what
|
| 57 |
+
# THIS corpus attests, scored in its own embedding space. A corpus that
|
| 58 |
+
# cannot express an opposition now yields None, and callers must handle
|
| 59 |
+
# that honestly rather than borrow English from a list. A forklift
|
| 60 |
+
# corpus gets forklift oppositions or it gets nothing - which is the
|
| 61 |
+
# rule, and the point.
|
python/elidedb/log.py
ADDED
|
@@ -0,0 +1,293 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Transaction log: the Delta-Lake idea at StreetDex scale.
|
| 2 |
+
|
| 3 |
+
A table's state is not "the files in the directory" — it is the fold of an
|
| 4 |
+
append-only log of JSON commits. That one move buys, with plain files:
|
| 5 |
+
|
| 6 |
+
- **snapshot isolation / time travel**: version N is immutable forever; a
|
| 7 |
+
reader at version N never sees version N+1's files.
|
| 8 |
+
- **atomic multi-file commits**: a commit lands as one exclusively-created
|
| 9 |
+
log entry (O_EXCL is the lock — single writer, many readers).
|
| 10 |
+
- **file-level zone maps**: every added file records rows/bytes/min_ts/max_ts,
|
| 11 |
+
so a time-window query prunes whole files from the log alone, before any
|
| 12 |
+
Parquet footer is opened. (Row-group pruning inside surviving files is
|
| 13 |
+
Parquet's own statistics — two layers, same idea.)
|
| 14 |
+
- **schema-on-log**: the schema travels with the commit, so evolution is an
|
| 15 |
+
append, never a rewrite.
|
| 16 |
+
"""
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import json
|
| 20 |
+
import os
|
| 21 |
+
import time
|
| 22 |
+
from dataclasses import dataclass, field
|
| 23 |
+
from pathlib import Path
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
@dataclass
|
| 27 |
+
class FileEntry:
|
| 28 |
+
path: str # relative to the table dir
|
| 29 |
+
rows: int
|
| 30 |
+
bytes: int
|
| 31 |
+
min_ts: int
|
| 32 |
+
max_ts: int
|
| 33 |
+
# Zone map for columns BEYOND ts: {column: [min, max]}. min_ts/max_ts
|
| 34 |
+
# are the same idea hard-coded for the one column every table has;
|
| 35 |
+
# this generalises it to whatever a table is clustered on, so
|
| 36 |
+
# "object_id == 7" can drop a file from the LOG - kilobytes of JSON
|
| 37 |
+
# already in memory - instead of opening its Parquet footer. That is
|
| 38 |
+
# a layer below Parquet: layer 0 costs nothing per file, the footer
|
| 39 |
+
# costs a seek and a read per surviving file, and at a thousand
|
| 40 |
+
# files the difference is the query.
|
| 41 |
+
#
|
| 42 |
+
# Only meaningful for a column the file is CLUSTERED on. A min/max
|
| 43 |
+
# over an unsorted column spans nearly the whole domain and prunes
|
| 44 |
+
# nothing, so the writer records these only for its sort keys -
|
| 45 |
+
# a zone map that never prunes is pure metadata cost.
|
| 46 |
+
zone: dict = field(default_factory=dict)
|
| 47 |
+
# END of the latest interval in this file, i.e. max(t1) - NOT
|
| 48 |
+
# max(ts). For point-like rows the two are the same and this is
|
| 49 |
+
# redundant; for INTERVALS they are not, and the difference is a
|
| 50 |
+
# wrong answer rather than a slow one.
|
| 51 |
+
#
|
| 52 |
+
# An object present from t=100 to t=900 in a file whose starts span
|
| 53 |
+
# 100..200 has max_ts=200. A query for [400, 500] overlaps that
|
| 54 |
+
# object, but `max_ts < t0` prunes the whole file and returns
|
| 55 |
+
# nothing. Verified: 0 rows where 1 was correct. Every table until
|
| 56 |
+
# now was point-like or short-lived so it never fired; presence
|
| 57 |
+
# intervals - "this object was on the table for four hours" - are
|
| 58 |
+
# exactly the shape that breaks it.
|
| 59 |
+
max_end: int = 0
|
| 60 |
+
|
| 61 |
+
def to_json(self):
|
| 62 |
+
d = dict(self.__dict__)
|
| 63 |
+
if not d["zone"]:
|
| 64 |
+
d.pop("zone") # old readers, and old files, see no change
|
| 65 |
+
if not d["max_end"]:
|
| 66 |
+
d.pop("max_end")
|
| 67 |
+
return d
|
| 68 |
+
|
| 69 |
+
@staticmethod
|
| 70 |
+
def from_json(d):
|
| 71 |
+
# a file written before max_end existed falls back to max_ts,
|
| 72 |
+
# which is CONSERVATIVE only for point-like rows. It is recorded
|
| 73 |
+
# as such rather than silently trusted: see overlaps().
|
| 74 |
+
return FileEntry(d["path"], d["rows"], d["bytes"], d["min_ts"],
|
| 75 |
+
d["max_ts"], d.get("zone", {}),
|
| 76 |
+
int(d.get("max_end", 0)))
|
| 77 |
+
|
| 78 |
+
def overlaps(self, t0, t1) -> bool:
|
| 79 |
+
"""True unless this file provably holds no interval meeting
|
| 80 |
+
[t0, t1]. Compares against the interval END, so a long-lived row
|
| 81 |
+
that starts before the window is not pruned away."""
|
| 82 |
+
end = self.max_end or self.max_ts
|
| 83 |
+
if t0 is not None and end < t0:
|
| 84 |
+
return False
|
| 85 |
+
if t1 is not None and self.min_ts > t1:
|
| 86 |
+
return False
|
| 87 |
+
return True
|
| 88 |
+
|
| 89 |
+
def may_contain(self, column, lo, hi) -> bool:
|
| 90 |
+
"""False only when this file PROVABLY holds nothing in [lo, hi].
|
| 91 |
+
|
| 92 |
+
Absent statistics must answer True: a missing zone map means
|
| 93 |
+
unknown, never empty. Every file written before zone maps
|
| 94 |
+
existed takes that branch, so the optimisation degrades to the
|
| 95 |
+
old behaviour instead of silently losing rows.
|
| 96 |
+
"""
|
| 97 |
+
z = self.zone.get(column)
|
| 98 |
+
if not z:
|
| 99 |
+
return True
|
| 100 |
+
return not (hi < z[0] or lo > z[1])
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
@dataclass
|
| 104 |
+
class TableState:
|
| 105 |
+
version: int = 0
|
| 106 |
+
kind: str = "timeseries"
|
| 107 |
+
schema: str = ""
|
| 108 |
+
files: list[FileEntry] = field(default_factory=list)
|
| 109 |
+
meta: dict = field(default_factory=dict)
|
| 110 |
+
|
| 111 |
+
@property
|
| 112 |
+
def rows(self):
|
| 113 |
+
return sum(f.rows for f in self.files)
|
| 114 |
+
|
| 115 |
+
@property
|
| 116 |
+
def bytes(self):
|
| 117 |
+
return sum(f.bytes for f in self.files)
|
| 118 |
+
|
| 119 |
+
@property
|
| 120 |
+
def min_ts(self):
|
| 121 |
+
return min((f.min_ts for f in self.files), default=0)
|
| 122 |
+
|
| 123 |
+
@property
|
| 124 |
+
def max_ts(self):
|
| 125 |
+
return max((f.max_ts for f in self.files), default=0)
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
CHECKPOINT_EVERY = 10 # Delta checkpoints every 10th commit; same dial here
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def _fsync_dir(path: Path):
|
| 132 |
+
"""Durability of a file's *existence* requires fsyncing its directory —
|
| 133 |
+
the metadata write that makes the entry reachable after power loss."""
|
| 134 |
+
fd = os.open(path, os.O_RDONLY)
|
| 135 |
+
try:
|
| 136 |
+
os.fsync(fd)
|
| 137 |
+
finally:
|
| 138 |
+
os.close(fd)
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def fsync_file(path: Path):
|
| 142 |
+
fd = os.open(path, os.O_RDONLY)
|
| 143 |
+
try:
|
| 144 |
+
os.fsync(fd)
|
| 145 |
+
finally:
|
| 146 |
+
os.close(fd)
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
class CommitConflict(RuntimeError):
|
| 150 |
+
"""A concurrent commit removed files this transaction depended on."""
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
class TableLog:
|
| 154 |
+
def __init__(self, table_dir: Path):
|
| 155 |
+
self.dir = Path(table_dir)
|
| 156 |
+
self.log_dir = self.dir / "_log"
|
| 157 |
+
|
| 158 |
+
def versions(self) -> list[int]:
|
| 159 |
+
if not self.log_dir.is_dir():
|
| 160 |
+
return []
|
| 161 |
+
return sorted(int(p.stem) for p in self.log_dir.glob("*.json")
|
| 162 |
+
if p.stem.isdigit())
|
| 163 |
+
|
| 164 |
+
def _checkpoints(self) -> list[int]:
|
| 165 |
+
if not self.log_dir.is_dir():
|
| 166 |
+
return []
|
| 167 |
+
return sorted(int(p.name.split(".")[0])
|
| 168 |
+
for p in self.log_dir.glob("*.checkpoint.json"))
|
| 169 |
+
|
| 170 |
+
def read_state(self, version: int | None = None) -> TableState:
|
| 171 |
+
st = TableState()
|
| 172 |
+
start = 0
|
| 173 |
+
# Checkpoints make the fold O(commits since checkpoint) instead of
|
| 174 |
+
# O(all commits) — the log stays an audit trail without becoming a
|
| 175 |
+
# read cost. Same move as Delta's _last_checkpoint.
|
| 176 |
+
for v in reversed(self._checkpoints()):
|
| 177 |
+
if version is None or v <= version:
|
| 178 |
+
c = json.loads(
|
| 179 |
+
(self.log_dir / f"{v:020d}.checkpoint.json").read_text())
|
| 180 |
+
st.version = c["version"]
|
| 181 |
+
st.kind = c["kind"]
|
| 182 |
+
st.schema = c["schema"]
|
| 183 |
+
st.meta = dict(c["meta"])
|
| 184 |
+
st.files = [FileEntry.from_json(f) for f in c["files"]]
|
| 185 |
+
start = v
|
| 186 |
+
break
|
| 187 |
+
for v in self.versions():
|
| 188 |
+
if v <= start:
|
| 189 |
+
continue
|
| 190 |
+
if version is not None and v > version:
|
| 191 |
+
break
|
| 192 |
+
entry = json.loads((self.log_dir / f"{v:020d}.json").read_text())
|
| 193 |
+
st.version = v
|
| 194 |
+
st.kind = entry.get("table_kind", st.kind)
|
| 195 |
+
st.schema = entry.get("schema", st.schema)
|
| 196 |
+
st.meta.update(entry.get("meta", {}))
|
| 197 |
+
removed = set(entry.get("remove", []))
|
| 198 |
+
if removed:
|
| 199 |
+
st.files = [f for f in st.files if f.path not in removed]
|
| 200 |
+
st.files += [FileEntry.from_json(f) for f in entry.get("add", [])]
|
| 201 |
+
return st
|
| 202 |
+
|
| 203 |
+
def commit(self, *, op: str, kind: str, schema: str = "",
|
| 204 |
+
add: list[FileEntry] = (), remove: list[str] = (),
|
| 205 |
+
meta: dict | None = None, retries: int = 5) -> int:
|
| 206 |
+
"""Optimistic concurrency, Delta-style: O_EXCL on the next log entry
|
| 207 |
+
is the lock. Losing the race means retrying at the next version —
|
| 208 |
+
append-vs-append never truly conflicts. Commits that REMOVE files
|
| 209 |
+
revalidate against the fresh state first: if a concurrent writer
|
| 210 |
+
already removed one of ours, that is a real conflict and we fail
|
| 211 |
+
cleanly instead of double-applying."""
|
| 212 |
+
self.log_dir.mkdir(parents=True, exist_ok=True)
|
| 213 |
+
for _ in range(retries):
|
| 214 |
+
version = (self.versions() or [0])[-1] + 1
|
| 215 |
+
if remove:
|
| 216 |
+
active = {f.path for f in self.read_state().files}
|
| 217 |
+
missing = [r for r in remove if r not in active]
|
| 218 |
+
if missing:
|
| 219 |
+
raise CommitConflict(
|
| 220 |
+
"files no longer active (concurrent rewrite?): "
|
| 221 |
+
f"{missing[:3]}")
|
| 222 |
+
entry = {
|
| 223 |
+
"version": version,
|
| 224 |
+
"ts_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
| 225 |
+
"op": op,
|
| 226 |
+
"table_kind": kind,
|
| 227 |
+
"schema": schema,
|
| 228 |
+
"add": [f.to_json() for f in add],
|
| 229 |
+
"remove": list(remove),
|
| 230 |
+
"meta": meta or {},
|
| 231 |
+
}
|
| 232 |
+
path = self.log_dir / f"{version:020d}.json"
|
| 233 |
+
# Torn-write-proof commit: the entry is fully written and fsynced
|
| 234 |
+
# under a temp name, then hard-LINKED to its final name. link()
|
| 235 |
+
# is atomic AND exclusive (fails if the name exists), so a
|
| 236 |
+
# partially-written entry can never appear under a version name —
|
| 237 |
+
# a crash leaves only an ignorable *.tmp. Atomicity + durability
|
| 238 |
+
# in one primitive; fsync of the directory makes the rename
|
| 239 |
+
# itself survive power loss.
|
| 240 |
+
tmp = self.log_dir / f".{version:020d}.{os.getpid()}.tmp"
|
| 241 |
+
fd = os.open(tmp, os.O_CREAT | os.O_WRONLY | os.O_TRUNC)
|
| 242 |
+
with os.fdopen(fd, "w") as f:
|
| 243 |
+
f.write(json.dumps(entry, indent=1))
|
| 244 |
+
f.flush()
|
| 245 |
+
os.fsync(f.fileno())
|
| 246 |
+
try:
|
| 247 |
+
os.link(tmp, path)
|
| 248 |
+
except FileExistsError:
|
| 249 |
+
os.unlink(tmp)
|
| 250 |
+
continue # lost the race — re-read state, take the next slot
|
| 251 |
+
os.unlink(tmp)
|
| 252 |
+
_fsync_dir(self.log_dir)
|
| 253 |
+
if version % CHECKPOINT_EVERY == 0:
|
| 254 |
+
st = self.read_state(version)
|
| 255 |
+
cp_tmp = self.log_dir / f".cp{version}.{os.getpid()}.tmp"
|
| 256 |
+
cp_tmp.write_text(
|
| 257 |
+
json.dumps({"version": st.version, "kind": st.kind,
|
| 258 |
+
"schema": st.schema, "meta": st.meta,
|
| 259 |
+
"files": [f.to_json() for f in st.files]}))
|
| 260 |
+
os.replace(cp_tmp, # checkpoints are derived: replace is fine
|
| 261 |
+
self.log_dir / f"{version:020d}.checkpoint.json")
|
| 262 |
+
# `_meta.json` beside the table: the human/tool-readable summary
|
| 263 |
+
# of CURRENT state (schema, rows, bytes, ts range) so a store
|
| 264 |
+
# browser or a pilot's script answers "what is in this table?"
|
| 265 |
+
# with one file read, not a log replay. Derived — a failure
|
| 266 |
+
# here must never fail the commit.
|
| 267 |
+
try:
|
| 268 |
+
st = self.read_state(version)
|
| 269 |
+
meta_tmp = self.dir / f".meta.{os.getpid()}.tmp"
|
| 270 |
+
meta_tmp.write_text(json.dumps({
|
| 271 |
+
"version": st.version, "op": op, "kind": st.kind,
|
| 272 |
+
"rows": sum(f.rows for f in st.files),
|
| 273 |
+
"bytes": sum(f.bytes for f in st.files),
|
| 274 |
+
"files": len(st.files),
|
| 275 |
+
"min_ts": min((f.min_ts for f in st.files), default=None),
|
| 276 |
+
"max_ts": max((f.max_ts for f in st.files), default=None),
|
| 277 |
+
"schema": st.schema, "meta": st.meta,
|
| 278 |
+
"ts_utc": entry["ts_utc"]}, indent=1))
|
| 279 |
+
os.replace(meta_tmp, self.dir / "_meta.json")
|
| 280 |
+
except Exception:
|
| 281 |
+
pass
|
| 282 |
+
return version
|
| 283 |
+
raise CommitConflict(f"lost the commit race {retries} times")
|
| 284 |
+
|
| 285 |
+
def history(self) -> list[dict]:
|
| 286 |
+
out = []
|
| 287 |
+
for v in self.versions():
|
| 288 |
+
e = json.loads((self.log_dir / f"{v:020d}.json").read_text())
|
| 289 |
+
out.append({"version": v, "ts_utc": e["ts_utc"], "op": e["op"],
|
| 290 |
+
"added_files": len(e.get("add", [])),
|
| 291 |
+
"added_rows": sum(a["rows"] for a in e.get("add", [])),
|
| 292 |
+
"meta": e.get("meta", {})})
|
| 293 |
+
return out
|
python/elidedb/migrate.py
ADDED
|
@@ -0,0 +1,163 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""v1 (custom SDX/SFI binaries) → v2 (Parquet lake) migration.
|
| 2 |
+
|
| 3 |
+
Reads the v1 store's own files — SDX columnar chunks, SFI frame tables,
|
| 4 |
+
embedding runs — and rewrites each as Parquet tables under a transaction
|
| 5 |
+
log. Raw media files stay exactly where they are; only indexes and sensor
|
| 6 |
+
rows move. This is the "no custom datatypes" pivot: the v1 formats keep
|
| 7 |
+
working for the C++ engine, but the go-forward store is open-format.
|
| 8 |
+
"""
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import json
|
| 12 |
+
import struct
|
| 13 |
+
import sys
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
import pyarrow as pa
|
| 18 |
+
|
| 19 |
+
from .store import Store
|
| 20 |
+
|
| 21 |
+
_SDX_TYPES = {0: ("i8", 8), 1: ("f8", 8), 2: ("f4", 4), 3: ("i8", 8), 4: ("i2", 2)}
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def read_sdx(path) -> pa.Table:
|
| 25 |
+
"""Minimal SDX v1 reader (layout: FORMAT.md §1)."""
|
| 26 |
+
data = Path(path).read_bytes()
|
| 27 |
+
assert data[:4] == b"SDX1" and data[-4:] == b"SDX1", f"bad SDX: {path}"
|
| 28 |
+
(flen,) = struct.unpack_from("<I", data, len(data) - 8)
|
| 29 |
+
off = len(data) - 8 - flen
|
| 30 |
+
def u16():
|
| 31 |
+
nonlocal off; (v,) = struct.unpack_from("<H", data, off); off += 2; return v
|
| 32 |
+
def u32():
|
| 33 |
+
nonlocal off; (v,) = struct.unpack_from("<I", data, off); off += 4; return v
|
| 34 |
+
def i64():
|
| 35 |
+
nonlocal off; (v,) = struct.unpack_from("<q", data, off); off += 8; return v
|
| 36 |
+
def s():
|
| 37 |
+
nonlocal off
|
| 38 |
+
n = u16()
|
| 39 |
+
v = data[off:off + n].decode(); off += n; return v
|
| 40 |
+
_ver, _flags = u16(), u16()
|
| 41 |
+
_sid, _units = s(), s()
|
| 42 |
+
_clock_off = i64()
|
| 43 |
+
_chunk_target = u32()
|
| 44 |
+
ncol = u16()
|
| 45 |
+
cols = []
|
| 46 |
+
for _ in range(ncol):
|
| 47 |
+
name = s(); ty = data[off]; off += 2
|
| 48 |
+
cols.append((name, ty))
|
| 49 |
+
nchunk = u32()
|
| 50 |
+
spans = [[] for _ in cols]
|
| 51 |
+
for _ in range(nchunk):
|
| 52 |
+
_rows = u32(); off_pad = u32() # noqa: F841
|
| 53 |
+
for c in range(ncol):
|
| 54 |
+
(doff, dlen, _mn, _mx) = struct.unpack_from("<QQQQ", data, off)
|
| 55 |
+
off += 32
|
| 56 |
+
spans[c].append((doff, dlen))
|
| 57 |
+
arrays = {}
|
| 58 |
+
for c, (name, ty) in enumerate(cols):
|
| 59 |
+
np_ty, _w = _SDX_TYPES[ty]
|
| 60 |
+
parts = [np.frombuffer(data, dtype=np_ty,
|
| 61 |
+
count=dlen // int(np_ty[1]),
|
| 62 |
+
offset=doff) for (doff, dlen) in spans[c]]
|
| 63 |
+
arrays["ts" if c == 0 else name] = np.concatenate(parts) if parts else \
|
| 64 |
+
np.array([], np_ty)
|
| 65 |
+
return pa.table({k: pa.array(v) for k, v in arrays.items()})
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def _sfi_frame_table(sfi_path, source_path, stream, width, height, codec):
|
| 69 |
+
sys.path.insert(0, str(Path(__file__).parents[2] / "ml"))
|
| 70 |
+
from sfi_reader import read_sfi
|
| 71 |
+
sfi = read_sfi(str(sfi_path))
|
| 72 |
+
fr = sfi.frames
|
| 73 |
+
n = len(fr)
|
| 74 |
+
return pa.table({
|
| 75 |
+
"ts": pa.array(fr["pts_ns"].astype("int64")),
|
| 76 |
+
"byte_offset": pa.array(fr["byte_offset"].astype("int64")),
|
| 77 |
+
"packet_size": pa.array(fr["packet_size"].astype("int32")),
|
| 78 |
+
"keyframe": pa.array((fr["flags"] & 1).astype(bool)),
|
| 79 |
+
"width": pa.array(np.full(n, width, "int32")),
|
| 80 |
+
"height": pa.array(np.full(n, height, "int32")),
|
| 81 |
+
"codec": pa.array([codec] * n),
|
| 82 |
+
"source": pa.array([str(source_path)] * n),
|
| 83 |
+
"stream": pa.array([stream] * n),
|
| 84 |
+
})
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def migrate_v1(v1_store: str, out_path: str, name: str,
|
| 88 |
+
run_id: str | None = None, verbose=True) -> Store:
|
| 89 |
+
v1 = Path(v1_store)
|
| 90 |
+
cur = int((v1 / "CURRENT").read_text().strip())
|
| 91 |
+
man = json.loads((v1 / "manifests" / f"manifest-{cur}.json").read_text())
|
| 92 |
+
db = Store.create(out_path, name)
|
| 93 |
+
|
| 94 |
+
# video: every SFI → rows of ONE frame_index table ("frames")
|
| 95 |
+
for vs in man.get("video_streams", []):
|
| 96 |
+
for seg in vs["segments"]:
|
| 97 |
+
t = _sfi_frame_table(v1 / seg["sfi_path"], seg["source_path"],
|
| 98 |
+
vs["stream_id"], vs.get("width", 0),
|
| 99 |
+
vs.get("height", 0), "mjpeg")
|
| 100 |
+
db.table("frames").append(t, kind="frame_index",
|
| 101 |
+
meta={"migrated_from": seg["sfi_path"]})
|
| 102 |
+
if verbose:
|
| 103 |
+
print(f" frames << {vs['stream_id']} ({len(t)} rows)")
|
| 104 |
+
|
| 105 |
+
# sensors: every SDX → its own timeseries table
|
| 106 |
+
for ss in man.get("sensor_streams", []):
|
| 107 |
+
t = read_sdx(v1 / ss["sdx_path"])
|
| 108 |
+
tname = ss["stream_id"].replace("/", "_").replace(" ", "_").lower()
|
| 109 |
+
db.table(tname).append(t, meta={"migrated_from": ss["sdx_path"]})
|
| 110 |
+
if verbose:
|
| 111 |
+
print(f" {tname} << {len(t):,} rows")
|
| 112 |
+
|
| 113 |
+
# embeddings run → embeddings + centroids tables
|
| 114 |
+
if run_id is None and man.get("semantic_run"):
|
| 115 |
+
run_id = man["semantic_run"]["run_id"]
|
| 116 |
+
if run_id:
|
| 117 |
+
run = v1 / "ml" / run_id
|
| 118 |
+
win = json.loads((run / "windows.json").read_text())
|
| 119 |
+
dim = win["dim"]
|
| 120 |
+
vecs = np.fromfile(run / "embeddings.f32", dtype=np.float32) \
|
| 121 |
+
.reshape(-1, dim)
|
| 122 |
+
labels = None
|
| 123 |
+
if (run / "clusters.json").exists():
|
| 124 |
+
cj = json.loads((run / "clusters.json").read_text())
|
| 125 |
+
labels = np.array(cj["labels"], "int32")
|
| 126 |
+
cents = np.array(cj["centroids"], "float32")
|
| 127 |
+
t = pa.table({
|
| 128 |
+
"ts": pa.array([w["t0_ns"] for w in win["windows"]], pa.int64()),
|
| 129 |
+
"t1": pa.array([w["t1_ns"] for w in win["windows"]], pa.int64()),
|
| 130 |
+
"stream": pa.array([w["stream_id"] for w in win["windows"]]),
|
| 131 |
+
"vector": pa.array([v.tolist() for v in vecs],
|
| 132 |
+
pa.list_(pa.float32(), dim)),
|
| 133 |
+
**({"cluster": pa.array(labels)} if labels is not None else {}),
|
| 134 |
+
})
|
| 135 |
+
meta = json.loads((run / "meta.json").read_text())
|
| 136 |
+
db.table("embeddings").append(
|
| 137 |
+
t, kind="embeddings",
|
| 138 |
+
meta={"model": meta["model"], "dim": dim,
|
| 139 |
+
"window_s": meta.get("window_s"), "migrated_from": run_id})
|
| 140 |
+
if labels is not None and len(cents):
|
| 141 |
+
ct = pa.table({
|
| 142 |
+
"ts": pa.array([0] * len(cents), pa.int64()),
|
| 143 |
+
"cluster": pa.array(range(len(cents)), pa.int32()),
|
| 144 |
+
"vector": pa.array([c.tolist() for c in cents],
|
| 145 |
+
pa.list_(pa.float32(), dim)),
|
| 146 |
+
})
|
| 147 |
+
db.table("centroids").append(ct, kind="centroids")
|
| 148 |
+
if verbose:
|
| 149 |
+
print(f" embeddings << {len(t)} windows (dim {dim})")
|
| 150 |
+
return db
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
if __name__ == "__main__":
|
| 154 |
+
import argparse
|
| 155 |
+
ap = argparse.ArgumentParser()
|
| 156 |
+
ap.add_argument("v1_store")
|
| 157 |
+
ap.add_argument("out")
|
| 158 |
+
ap.add_argument("--name", required=True)
|
| 159 |
+
args = ap.parse_args()
|
| 160 |
+
db = migrate_v1(args.v1_store, args.out, args.name)
|
| 161 |
+
for d in db.describe():
|
| 162 |
+
print(f"{d['table']:14s} {d['kind']:12s} {d['rows']:>12,} rows "
|
| 163 |
+
f"{d['bytes'] / 1e6:9.1f} MB v{d['version']}")
|