From 5cbb21c88e9bef4d1f3ec095085cedb24981ff6f Mon Sep 17 00:00:00 2001 From: RDCLder Date: Fri, 6 Nov 2026 13:16:27 -0600 Subject: [PATCH 1/6] Add the folder index file contract jepa-scope.index is a safetensors file with a validated JSON header, like .latent. It holds every clip's window vectors in clip order, their start times, the clip each row belongs to and the window-weighted mean, plus a header with the representation block, the run settings, the root folder, skipped files and one entry per clip. Validation covers dtypes, finite values, ranges that tile the rows, strictly increasing times within a clip, portable relative paths, strict JSON and the stored mean. Saving goes through a unique temporary file beside the target. Signed-off-by: RDCLder --- ...6-weeks3-4-folder-search-trajectory-log.md | 21 + ...11-06-weeks3-4-folder-search-trajectory.md | 10 +- src/jepa_scope/index_file.py | 334 +++++++++ tests/test_index.py | 638 ++++++++++++++++++ 4 files changed, 1000 insertions(+), 3 deletions(-) create mode 100644 src/jepa_scope/index_file.py create mode 100644 tests/test_index.py diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md index 27a3323..81b495a 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md @@ -106,3 +106,24 @@ The plan is `2026-11-06-weeks3-4-folder-search-trajectory.md`. Entries are newes pass). On the real model the identity is unchanged (`550d387c4f4a1862`, hash `e8eb61d31255e782`), so nothing re-embeds. - **Permanent decision record needed:** No beyond the T10 entry. + +- **2026-11-06 — PR 1 merged (T0 and T1, ryangswe/jepa-scope#5); old cache deleted; T2 starts.** + - **Evidence:** the owner merged #5 (merge commit `ec26b1a`). 201 version-1 cache files (4.03 GB), + T0's set plus one earlier clip's, were deleted; the one version-2 file was kept. + - **Impact:** T2–T5 go on branch `week3-index` from `main`. + - **Resolution:** T2 delegated to an executor, as the plan routes it. + - **Permanent decision record needed:** No. + +- **2026-11-06 — T2 complete: the index file contract (`index_file.py`).** + - **Evidence:** an executor (Sonnet) wrote the module and its tests; a fresh reviewer (Opus) did + two rounds. Round 1 found eight gaps that would have been hard to tighten once users have + index files: a relative `root`; non-finite or out-of-range clip numbers and bare `NaN` in the + JSON (which would break `/api/meta` for a whole folder); negative clip ids; path holes and + duplicates; loosely shaped `skipped` and `sha256`; an unchecked `mean`; and a fixed temporary + name. Round 2 found that its own suggested ":" rule rejected legitimate timestamped names, plus + a fixed file mode and a temporary file placed beside a symlink's target. All are fixed; the + path rule is now public as `is_valid_clip_path`, and T4 skips failing files instead of + crashing. The fast suite passes (235). Every check was break-tested by removing it. + - **Impact:** T3 builds on `Index`, `clip_rows` and `clip_vector`. + - **Resolution:** committed on `week3-index`. + - **Permanent decision record needed:** Yes, at T10 (the index file as a versioned contract). diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md index 2da4383..3257865 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md @@ -6,7 +6,8 @@ adversarial review and one outside review (see "Plan review") - **Spec / source:** `docs/roadmap.md`, "Next" item 1 and the project section's MVP schedule; the 2026-11-05 process decision that milestones follow the roadmap's weekly order -- **Branch:** `week3-folder-search` (T0–T7), then `week4-trajectory` (T8–T10) +- **Branch:** one per pull request: `week3-folder-search` (T0–T1, merged as #5), `week3-index` + (T2–T5), then branches for T6–T7 and T8–T10 ## Goal @@ -565,7 +566,7 @@ flowchart LR ### T2 — Week 3. The index file contract -- **Status:** Pending +- **Status:** Complete - **Dependencies:** T1 - **Parallel group:** Sequential - **Complexity:** ROUTINE @@ -645,7 +646,10 @@ flowchart LR - **Expected files/workspaces:** `src/jepa_scope/indexing.py` (new), `src/jepa_scope/cache.py` (new), `src/jepa_scope/cli.py`, `tests/test_index.py`, `README.md` - **Objective:** `build_or_refresh(folder, encoder_factory, settings, out, recursive, patches, - centre, verify, log)` as in Approach. Move `cli._cache_path` and `_cache_is_valid` into a small + centre, verify, log)` as in Approach. Files whose relative path fails + `index_file.is_valid_clip_path` (a backslash, a NUL or a Windows drive prefix) are listed under + `skipped` with the reason, so one oddly named file never stops the index from being written + (added after T2's review). Move `cli._cache_path` and `_cache_is_valid` into a small `cache.py` so `indexing.py` does not import the CLI; the refresh recomputes the cache path from the current settings for both the with-patches and without-patches variants. The encoder is created lazily, only when a clip must be embedded, so a refresh with nothing new never loads the diff --git a/src/jepa_scope/index_file.py b/src/jepa_scope/index_file.py new file mode 100644 index 0000000..c69df8d --- /dev/null +++ b/src/jepa_scope/index_file.py @@ -0,0 +1,334 @@ +"""The index file: window vectors for a whole folder of videos in one safetensors file. + +An index file is a plain safetensors file, built the same way as a ``.latent`` file. Its metadata +holds one JSON string under the key ``jepa_scope`` that records how the vectors were made (the +``representation`` block copied from the clips), which clips they came from, and how similarity is +to be computed. The tensors hold every clip's window vectors, in clip order. Any safetensors reader +can open the tensors; this module adds validation and a small typed view, ``Index``. +""" + +from __future__ import annotations + +import json +import math +import os +import re +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +from safetensors import SafetensorError, safe_open +from safetensors.numpy import load_file, save_file + +from . import latent_file + +INDEX_FORMAT = "jepa-scope.index" +INDEX_FORMAT_VERSION = 1 +META_KEY = "jepa_scope" + +REQUIRED_HEADER_KEYS = ( + "format", + "format_version", + "representation", + "run", + "embedding", + "root", + "skipped", + "clips", +) +RUN_KEYS = ("max_seconds", "device", "dtype") +SIMILARITIES = ("cosine", "centred-cosine") +_DRIVE = re.compile(r"[A-Za-z]:(/|$)") # a Windows drive prefix, as in "C:/videos" +CLIP_KEYS = ("path", "sha256", "size", "mtime_ns", "duration", "width", "height", "windows") + +# The exact dtype each tensor must have. +TENSOR_DTYPES = { + "pooled": np.dtype(np.float32), + "window_start_time": np.dtype(np.float64), + "clip": np.dtype(np.int32), + "mean": np.dtype(np.float32), +} + + +class IndexFileError(ValueError): + """Raised when a file is not a valid index, or its tensors and header disagree.""" + + +@dataclass +class Index: + """A loaded index: the validated header and the four tensors. + + Row ``r`` of ``pooled`` is one window vector, ``window_start_time[r]`` is when that window + starts in its clip, and ``clip[r]`` is the position of its clip in ``header["clips"]``. + """ + + header: dict[str, Any] + pooled: np.ndarray + window_start_time: np.ndarray + clip: np.ndarray + mean: np.ndarray + + def clip_rows(self, i: int) -> slice: + """The rows of ``pooled`` that belong to clip ``i``, as a slice. + + Raises ``IndexError`` for a negative ``i`` or one past the last clip, instead of wrapping + around the way a Python list would. + """ + if not 0 <= i < len(self.header["clips"]): + raise IndexError(f"clip {i} is out of range for an index of {len(self.header['clips'])} clips") + start, stop = self.header["clips"][i]["windows"] + return slice(start, stop) + + def clip_vector(self, i: int) -> np.ndarray: + """The mean of clip ``i``'s raw window vectors, as float32 of shape ``[D]``. + + The rows are averaged exactly as stored. This function neither subtracts the index mean + nor normalises the result; centring and normalisation belong to search. + """ + return self.pooled[self.clip_rows(i)].mean(axis=0, dtype=np.float64).astype(np.float32) + + def resolve_root(self, index_path: str | Path, source_root: str | Path | None = None) -> Path: + """The folder that the clips' relative paths are resolved against. + + ``source_root``, when given, wins over whatever the index recorded; this is for a folder + that has moved. Otherwise a recorded root of ``"."`` means the directory that holds the + index file itself, so a folder with its colocated index can be moved or copied. Any other + recorded root is an absolute path and is used as it is. + """ + if source_root is not None: + return Path(source_root) + root = self.header["root"] + if root == ".": + return Path(index_path).resolve().parent + return Path(root) + + +def save(path: str | Path, tensors: dict[str, np.ndarray], header: dict[str, Any]) -> None: + """Write an index file in the current format, after checking it. + + The current ``format`` and ``format_version`` always win over the caller's header, so saving + never writes an old version number. The file is written under a temporary name and then + renamed, so an interrupted write never leaves a half file behind. + """ + header = {**header, "format": INDEX_FORMAT, "format_version": INDEX_FORMAT_VERSION} + check(tensors, header) + contiguous = {k: np.ascontiguousarray(v) for k, v in tensors.items()} + try: + text = json.dumps(header, sort_keys=True, allow_nan=False) + except ValueError as exc: + raise IndexFileError(f"the header cannot be written as JSON: {exc}") from exc + # Write to a uniquely named file in the target's directory and rename it, so an interrupted + # write never leaves a half file behind and two writers never share a temporary name. + # The temporary file sits beside the name being replaced (not beside a symlink's target), and + # gets the permissions an ordinary new file would (mkstemp creates it private). + fd, tmp = tempfile.mkstemp(dir=Path(path).absolute().parent, prefix=Path(path).name + ".", suffix=".tmp") + os.close(fd) + try: + save_file(contiguous, tmp, metadata={META_KEY: text}) + os.chmod(tmp, 0o666 & ~_umask()) + os.replace(tmp, path) + except BaseException: + Path(tmp).unlink(missing_ok=True) + raise + + +def read_header(path: str | Path) -> dict[str, Any]: + """Read and check only the header, without loading the tensors.""" + try: + with safe_open(str(path), framework="numpy") as f: + meta = f.metadata() or {} + except SafetensorError as exc: + raise IndexFileError(f"{path} is not a readable safetensors file: {exc}") from exc + if META_KEY not in meta: + raise IndexFileError(f"{path} is a safetensors file but has no '{META_KEY}' header") + try: + header = json.loads(meta[META_KEY], parse_constant=_refuse_constant) + except ValueError as exc: + raise IndexFileError(f"{path} has a header that is not valid JSON: {exc}") from exc + if not isinstance(header, dict): + raise IndexFileError(f"{path} has a header that is not a JSON object") + check_header(header) + return header + + +def _refuse_constant(name: str) -> None: + raise ValueError(f"the constant {name} is not allowed in the header") + + +def load(path: str | Path) -> Index: + """Load and fully check an index file.""" + header = read_header(path) + try: + tensors = load_file(str(path)) + except SafetensorError as exc: + raise IndexFileError(f"{path} is damaged: {exc}") from exc + check(tensors, header) + return Index( + header=header, + pooled=tensors["pooled"], + window_start_time=tensors["window_start_time"], + clip=tensors["clip"], + mean=tensors["mean"], + ) + + +def check(tensors: dict[str, np.ndarray], header: dict[str, Any]) -> None: + """Check a header and its tensors together, for callers that build both in memory. + + An index with no clips is valid: it has zero rows and a ``mean`` of zeros. + """ + check_header(header) + missing = [k for k in TENSOR_DTYPES if k not in tensors] + if missing: + raise IndexFileError(f"missing tensors {missing}") + for key, dtype in TENSOR_DTYPES.items(): + if tensors[key].dtype != dtype: + raise IndexFileError(f"{key} must be {dtype}, not {tensors[key].dtype}") + pooled = tensors["pooled"] + if pooled.ndim != 2: + raise IndexFileError(f"pooled must be 2-dimensional, not shape {tuple(pooled.shape)}") + rows, dim = pooled.shape + hidden = header["representation"].get("hidden_size") + if not _is_int(hidden) or dim != hidden: + raise IndexFileError(f"pooled has width {dim} but the representation's hidden_size is {hidden}") + expected = {"window_start_time": (rows,), "clip": (rows,), "mean": (dim,)} + for key, shape in expected.items(): + if tuple(tensors[key].shape) != shape: + raise IndexFileError(f"{key} has shape {tuple(tensors[key].shape)}, expected {shape}") + for key in TENSOR_DTYPES: + if not np.isfinite(tensors[key]).all(): + raise IndexFileError(f"{key} holds NaN or infinite values") + _check_ranges(tensors, header["clips"], rows) + _check_mean(tensors["mean"], pooled) + + +def _check_mean(mean: np.ndarray, pooled: np.ndarray) -> None: + """``mean`` must be the window-weighted mean of ``pooled`` (zeros when there are no rows).""" + if len(pooled) == 0: + if mean.any(): + raise IndexFileError("mean must be all zeros for an index with no windows") + return + expected = pooled.mean(axis=0, dtype=np.float64) + atol = 1e-5 * max(1.0, float(np.abs(pooled).max())) + if not np.allclose(mean.astype(np.float64), expected, rtol=0.0, atol=atol): + raise IndexFileError("mean is not the mean of pooled over all windows") + + +def check_header(header: dict[str, Any]) -> None: + """Check the header on its own: its keys, the representation hash and every clip entry.""" + if header.get("format") != INDEX_FORMAT: + raise IndexFileError(f"not a {INDEX_FORMAT} file (format={header.get('format')!r})") + version = header.get("format_version") + if not _is_int(version) or version != INDEX_FORMAT_VERSION: + raise IndexFileError( + f"format_version {header.get('format_version')} is not supported " + f"(expected {INDEX_FORMAT_VERSION}); rebuild the index" + ) + missing = [k for k in REQUIRED_HEADER_KEYS if k not in header] + if missing: + raise IndexFileError(f"header is missing {missing}") + block = header["representation"] + if not isinstance(block, dict) or block.get("hash") != latent_file.representation_hash(block): + raise IndexFileError("the representation block does not match its hash") + run = header["run"] + if not isinstance(run, dict) or any(k not in run for k in RUN_KEYS): + raise IndexFileError(f"run must be an object with {list(RUN_KEYS)}") + max_seconds = run["max_seconds"] + if max_seconds is not None and not (_is_number(max_seconds) and math.isfinite(max_seconds) and max_seconds > 0): + raise IndexFileError("run.max_seconds must be null or a finite number above zero") + embedding = header["embedding"] + if not isinstance(embedding, dict) or embedding.get("kind") != "window_vector": + raise IndexFileError("embedding.kind must be 'window_vector'") + if embedding.get("similarity") not in SIMILARITIES: + raise IndexFileError(f"embedding.similarity must be one of {list(SIMILARITIES)}") + root = header["root"] + if not isinstance(root, str) or not (root == "." or Path(root).is_absolute()): + raise IndexFileError("root must be '.' or an absolute path") + skipped = header["skipped"] + if not isinstance(skipped, list) or not all( + isinstance(e, dict) and isinstance(e.get("path"), str) and isinstance(e.get("reason"), str) for e in skipped + ): + raise IndexFileError("skipped must be a list of objects with a string path and a string reason") + if not isinstance(header["clips"], list): + raise IndexFileError("clips must be a list") + for i, entry in enumerate(header["clips"]): + _check_clip_entry(i, entry) + paths = [e["path"] for e in header["clips"]] + if len(set(paths)) != len(paths): + raise IndexFileError("two clips have the same path") + + +def _check_clip_entry(i: int, entry: Any) -> None: + if not isinstance(entry, dict): + raise IndexFileError(f"clips[{i}] must be an object") + missing = [k for k in CLIP_KEYS if k not in entry] + if missing: + raise IndexFileError(f"clips[{i}] is missing {missing}") + _check_path(i, entry["path"]) + if not isinstance(entry["sha256"], str) or not re.fullmatch(r"[0-9a-f]{64}", entry["sha256"]): + raise IndexFileError(f"clips[{i}].sha256 must be 64 lowercase hex characters") + for key, floor in (("size", 0), ("mtime_ns", 0), ("width", 1), ("height", 1)): + if not _is_int(entry[key]): + raise IndexFileError(f"clips[{i}].{key} must be an integer") + if entry[key] < floor: + raise IndexFileError(f"clips[{i}].{key} must be at least {floor}") + duration = entry["duration"] + if not _is_number(duration) or not math.isfinite(duration) or duration <= 0: + raise IndexFileError(f"clips[{i}].duration must be a finite number above zero") + windows = entry["windows"] + if not (isinstance(windows, list) and len(windows) == 2 and all(_is_int(v) for v in windows)): + raise IndexFileError(f"clips[{i}].windows must be [start, stop) with two integers") + + +def is_valid_clip_path(path: Any) -> bool: + """Whether ``path`` can be stored as a clip path: relative to the root, with forward slashes, + and no empty, ``.`` or ``..`` part. A backslash, a NUL or a Windows drive prefix (``C:/``) is + refused, because the path would mean something else on another system. A colon elsewhere, as + in a timestamped name, is fine. The indexer skips files whose paths fail this, rather than + failing the whole index.""" + if not isinstance(path, str) or not path or "\\" in path or "\x00" in path: + return False + if _DRIVE.match(path): + return False + return not any(part in ("", ".", "..") for part in path.split("/")) + + +def _check_path(i: int, path: Any) -> None: + if not is_valid_clip_path(path): + raise IndexFileError(f"clips[{i}].path must be a plain relative path with forward slashes") + + +def _umask() -> int: + mask = os.umask(0) + os.umask(mask) + return mask + + +def _is_number(value: Any) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _check_ranges(tensors: dict[str, np.ndarray], clips: list[dict[str, Any]], rows: int) -> None: + """The clips' window ranges must tile ``[0, rows)`` in order and agree with the tensors.""" + expected_start = 0 + times = tensors["window_start_time"] + owner = tensors["clip"] + for i, entry in enumerate(clips): + start, stop = entry["windows"] + if stop <= start: + raise IndexFileError(f"clips[{i}].windows [{start}, {stop}) is empty") + if start != expected_start: + raise IndexFileError(f"clips[{i}].windows starts at {start}, expected {expected_start}") + if (owner[start:stop] != i).any(): + raise IndexFileError(f"clip does not say {i} for every row of clips[{i}]") + if stop - start > 1 and (np.diff(times[start:stop]) <= 0).any(): + raise IndexFileError(f"window_start_time does not strictly increase within clips[{i}]") + expected_start = stop + if expected_start != rows: + raise IndexFileError(f"the clips cover {expected_start} rows but pooled has {rows}") diff --git a/tests/test_index.py b/tests/test_index.py new file mode 100644 index 0000000..f46c71a --- /dev/null +++ b/tests/test_index.py @@ -0,0 +1,638 @@ +"""The index file contract: round trips, rejection of malformed files, and root resolution. + +Everything here uses small synthetic arrays; no model is involved. +""" + +from __future__ import annotations + +import copy +import json +import os +from pathlib import Path + +import numpy as np +import pytest +from safetensors.numpy import save_file + +from jepa_scope import index_file, latent_file +from jepa_scope.index_file import IndexFileError + +DIM = 8 +WINDOWS_PER_CLIP = (2, 3, 1) + + +def make_representation(hidden_size: int = DIM) -> dict: + return latent_file.representation( + model_id="test/model", + weights="sha256:abc", + hidden_size=hidden_size, + preprocess_hash="pre123", + window=16, + sample_fps=8.0, + output_version=2, + ) + + +def make_index(counts=WINDOWS_PER_CLIP, dim: int = DIM, similarity: str = "cosine"): + """A valid synthetic index with one clip per entry of ``counts``, as (tensors, header).""" + rng = np.random.default_rng(0) + rows = sum(counts) + clip_ids = np.repeat(np.arange(len(counts), dtype=np.int32), counts) + times = np.concatenate([np.arange(c, dtype=np.float64) * 1.0 for c in counts]) if counts else np.zeros(0) + pooled = rng.normal(size=(rows, dim)).astype(np.float32) + mean = pooled.mean(axis=0) if rows else np.zeros(dim, dtype=np.float32) + tensors = { + "pooled": pooled, + "window_start_time": times.astype(np.float64), + "clip": clip_ids, + "mean": mean.astype(np.float32), + } + clips = [] + start = 0 + for i, count in enumerate(counts): + clips.append( + { + "path": f"sub/clip{i}.mp4", + "sha256": f"{i:064x}", + "size": 1000 + i, + "mtime_ns": 1_700_000_000_000_000_000 + i, + "duration": 4.0 + i, + "width": 320, + "height": 240, + "windows": [start, start + count], + } + ) + start += count + header = { + "format": index_file.INDEX_FORMAT, + "format_version": index_file.INDEX_FORMAT_VERSION, + "representation": make_representation(dim), + "run": {"max_seconds": 30.0, "device": "cpu", "dtype": "float32"}, + "embedding": {"kind": "window_vector", "similarity": similarity}, + "root": ".", + "skipped": [{"path": "bad.mp4", "reason": "could not decode"}], + "clips": clips, + } + return tensors, header + + +def test_round_trip(tmp_path: Path): + tensors, header = make_index() + path = tmp_path / "jepa-scope.index" + index_file.save(path, tensors, header) + assert not (tmp_path / "jepa-scope.index.tmp").exists() + + index = index_file.load(path) + assert index.header == header + for key, value in tensors.items(): + assert getattr(index, key).dtype == value.dtype + np.testing.assert_array_equal(getattr(index, key), value) + assert index_file.read_header(path) == header + + +def test_empty_index_round_trips(tmp_path: Path): + tensors, header = make_index(counts=()) + assert tensors["pooled"].shape == (0, DIM) + assert not tensors["mean"].any() + path = tmp_path / "empty.index" + index_file.save(path, tensors, header) + + index = index_file.load(path) + assert index.pooled.shape == (0, DIM) + assert index.window_start_time.shape == (0,) + assert index.clip.shape == (0,) + assert index.mean.shape == (DIM,) + assert index.header["clips"] == [] + + +def test_save_stamps_current_format_and_version(tmp_path: Path): + tensors, header = make_index() + header["format_version"] = 0 + header["format"] = "something-else" + path = tmp_path / "x.index" + index_file.save(path, tensors, header) + loaded = index_file.read_header(path) + assert loaded["format"] == index_file.INDEX_FORMAT + assert loaded["format_version"] == index_file.INDEX_FORMAT_VERSION + + +def test_clip_rows_and_clip_vector(): + tensors, header = make_index() + index = index_file.Index(header=header, **tensors) + assert index.clip_rows(0) == slice(0, 2) + assert index.clip_rows(1) == slice(2, 5) + assert index.clip_rows(2) == slice(5, 6) + expected = tensors["pooled"][2:5].mean(axis=0) + got = index.clip_vector(1) + assert got.dtype == np.float32 + assert got.shape == (DIM,) + np.testing.assert_allclose(got, expected, rtol=1e-6) + # A one-window clip's vector is that row, unchanged and not normalised. + np.testing.assert_array_equal(index.clip_vector(2), tensors["pooled"][5]) + + +def test_resolve_root(tmp_path: Path): + tensors, header = make_index() + index = index_file.Index(header=header, **tensors) + index_path = tmp_path / "folder" / "jepa-scope.index" + (tmp_path / "folder").mkdir() + + # "." means the directory of the index file itself. + assert index.resolve_root(index_path) == (tmp_path / "folder").resolve() + + # An absolute recorded root is used as it is, wherever the index file is. + absolute = copy.deepcopy(header) + absolute["root"] = "/data/videos" + index_abs = index_file.Index(header=absolute, **tensors) + assert index_abs.resolve_root(index_path) == Path("/data/videos") + + # source_root overrides either kind of root. + override = tmp_path / "moved" + assert index.resolve_root(index_path, source_root=override) == override + assert index_abs.resolve_root(index_path, source_root=override) == override + + +def with_ranges(ranges): + tensors, header = make_index() + for entry, window in zip(header["clips"], ranges, strict=True): + entry["windows"] = list(window) + return tensors, header + + +@pytest.mark.parametrize( + ("ranges", "message"), + [ + pytest.param([(0, 2), (1, 5), (5, 6)], r"clips\[1\].windows starts at 1, expected 2", id="overlap"), + pytest.param([(0, 2), (3, 5), (5, 6)], r"clips\[1\].windows starts at 3, expected 2", id="gap"), + pytest.param([(0, 2), (2, 2), (2, 6)], r"clips\[1\].windows \[2, 2\) is empty", id="empty-range"), + pytest.param([(0, 2), (2, 5), (5, 5)], r"clips\[2\].windows \[5, 5\) is empty", id="empty-last-range"), + pytest.param([(1, 2), (2, 5), (5, 6)], r"clips\[0\].windows starts at 1, expected 0", id="not-from-zero"), + pytest.param([(0, 2), (2, 5), (5, 7)], r"cover 7 rows but pooled has 6", id="runs-past-the-rows"), + pytest.param([(0, 2), (2, 4), (4, 5)], r"clip does not say 2 for every row of clips\[2\]", id="wrong-owner"), + ], +) +def test_rejects_bad_window_ranges(ranges, message): + tensors, header = with_ranges(ranges) + with pytest.raises(IndexFileError, match=message): + index_file.check(tensors, header) + + +def test_rejects_ranges_that_stop_before_the_last_row(): + tensors, header = make_index() + del header["clips"][2] + tensors["clip"][5] = 1 # keep every remaining row owned by a listed clip + with pytest.raises(IndexFileError, match="cover 5 rows but pooled has 6"): + index_file.check(tensors, header) + + +def test_rejects_an_empty_clip_between_clips(): + # The rows and the clip tensor are consistent; only the empty range in the header is wrong. + tensors, header = make_index(counts=(2, 4)) + tensors["clip"][2:] = 2 + empty = copy.deepcopy(header["clips"][1]) + empty["path"] = "sub/empty.mp4" + empty["windows"] = [2, 2] + header["clips"].insert(1, empty) + header["clips"][2]["windows"] = [2, 6] + with pytest.raises(IndexFileError, match="is empty"): + index_file.check(tensors, header) + + +def test_rejects_clip_disagreeing_with_ranges(): + tensors, header = make_index() + tensors["clip"][1] = 2 + with pytest.raises(IndexFileError, match="clip does not say"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("bad_times", [[0.0, 1.0, 1.0], [0.0, 2.0, 1.0]], ids=["equal", "decreasing"]) +def test_rejects_non_increasing_times_within_a_clip(bad_times): + tensors, header = make_index() + tensors["window_start_time"][2:5] = bad_times + with pytest.raises(IndexFileError, match="strictly increase"): + index_file.check(tensors, header) + + +def test_times_may_restart_between_clips(): + tensors, header = make_index() + # The default fixture already restarts at 0 for every clip; this states the intent. + assert tensors["window_start_time"][2] == 0.0 + index_file.check(tensors, header) + + +def test_rejects_wrong_hidden_size(): + tensors, header = make_index() + header["representation"] = make_representation(hidden_size=DIM + 1) + with pytest.raises(IndexFileError, match="hidden_size"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize( + ("key", "bad"), + [ + ("pooled", np.float64), + ("pooled", np.float16), + ("window_start_time", np.float32), + ("clip", np.int64), + ("clip", np.uint8), + ("mean", np.float64), + ], +) +def test_rejects_wrong_dtype(key, bad): + tensors, header = make_index() + tensors[key] = tensors[key].astype(bad) + with pytest.raises(IndexFileError, match="must be"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("key", ["pooled", "window_start_time", "mean"]) +@pytest.mark.parametrize("bad", [np.nan, np.inf, -np.inf], ids=["nan", "inf", "-inf"]) +def test_rejects_non_finite_values(key, bad): + tensors, header = make_index() + tensors[key][0 if key != "pooled" else (0, 0)] = bad + with pytest.raises(IndexFileError, match="NaN or infinite"): + index_file.check(tensors, header) + + +def test_rejects_missing_tensor(): + tensors, header = make_index() + del tensors["mean"] + with pytest.raises(IndexFileError, match="missing tensors"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("key", ["window_start_time", "clip", "mean"]) +def test_rejects_wrong_tensor_shape(key): + tensors, header = make_index() + tensors[key] = tensors[key][:-1] + with pytest.raises(IndexFileError, match="shape"): + index_file.check(tensors, header) + + +def test_rejects_non_matrix_pooled(): + tensors, header = make_index() + tensors["pooled"] = tensors["pooled"].reshape(-1) + with pytest.raises(IndexFileError, match="2-dimensional"): + index_file.check(tensors, header) + + +def test_rejects_a_plain_safetensors_file(tmp_path: Path): + path = tmp_path / "plain.safetensors" + save_file({"x": np.zeros(3, dtype=np.float32)}, str(path)) + with pytest.raises(IndexFileError, match="no 'jepa_scope' header"): + index_file.load(path) + with pytest.raises(IndexFileError, match="no 'jepa_scope' header"): + index_file.read_header(path) + + +def test_rejects_a_latent_file_header(tmp_path: Path): + tensors, header = make_index() + header_with_latent_format = {**header, "format": latent_file.FORMAT} + # save() stamps the right format, so write the wrong one directly. + path = tmp_path / "wrong.index" + save_file(tensors, str(path), metadata={"jepa_scope": json.dumps(header_with_latent_format)}) + with pytest.raises(IndexFileError, match="not a jepa-scope.index file"): + index_file.load(path) + + +def test_rejects_other_format_version(tmp_path: Path): + tensors, header = make_index() + header["format_version"] = index_file.INDEX_FORMAT_VERSION + 1 + with pytest.raises(IndexFileError, match="format_version"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("key", index_file.REQUIRED_HEADER_KEYS) +def test_rejects_missing_header_key(key): + tensors, header = make_index() + del header[key] + with pytest.raises(IndexFileError): + index_file.check(tensors, header) + + +def test_rejects_representation_with_wrong_hash(): + tensors, header = make_index() + header["representation"]["pooling"] = "max" + with pytest.raises(IndexFileError, match="representation block"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("similarity", ["dot", "euclidean", None]) +def test_rejects_unknown_similarity(similarity): + tensors, header = make_index() + header["embedding"]["similarity"] = similarity + with pytest.raises(IndexFileError, match="similarity"): + index_file.check(tensors, header) + + +def test_accepts_centred_cosine(): + tensors, header = make_index(similarity="centred-cosine") + index_file.check(tensors, header) + + +def test_rejects_unknown_embedding_kind(): + tensors, header = make_index() + header["embedding"]["kind"] = "tubelet_vector" + with pytest.raises(IndexFileError, match="embedding.kind"): + index_file.check(tensors, header) + + +def test_rejects_run_without_settings(): + tensors, header = make_index() + del header["run"]["dtype"] + with pytest.raises(IndexFileError, match="run must be"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("missing", index_file.CLIP_KEYS) +def test_rejects_clip_entry_missing_a_field(missing): + tensors, header = make_index() + del header["clips"][0][missing] + with pytest.raises(IndexFileError, match="missing"): + index_file.check(tensors, header) + + +@pytest.mark.parametrize("path", ["/abs/clip.mp4", "a\\b.mp4", "../up.mp4", "a/../b.mp4", ""]) +def test_rejects_unportable_clip_paths(path): + tensors, header = make_index() + header["clips"][0]["path"] = path + with pytest.raises(IndexFileError, match="path"): + index_file.check(tensors, header) + + +def test_rejects_non_integer_clip_size(): + tensors, header = make_index() + header["clips"][0]["size"] = 1.5 + with pytest.raises(IndexFileError, match="size"): + index_file.check(tensors, header) + + +def test_save_refuses_an_invalid_index(tmp_path: Path): + tensors, header = make_index() + tensors["pooled"][0, 0] = np.nan + path = tmp_path / "bad.index" + with pytest.raises(IndexFileError): + index_file.save(path, tensors, header) + assert not path.exists() + assert not (tmp_path / "bad.index.tmp").exists() + + +def test_rejects_a_mean_that_is_not_the_mean_of_pooled(): + tensors, header = make_index() + tensors["mean"] = tensors["mean"] + 100 + with pytest.raises(IndexFileError, match="mean is not the mean"): + index_file.check(tensors, header) + + +def test_rejects_a_nonzero_mean_for_an_empty_index(): + tensors, header = make_index(counts=()) + tensors["mean"] = np.ones(DIM, dtype=np.float32) + with pytest.raises(IndexFileError, match="all zeros"): + index_file.check(tensors, header) + + +def test_mean_tolerance_scales_with_the_data(): + tensors, header = make_index() + tensors["pooled"] = tensors["pooled"] * np.float32(1e4) + tensors["mean"] = tensors["pooled"].mean(axis=0, dtype=np.float64).astype(np.float32) + index_file.check(tensors, header) + + +@pytest.mark.parametrize("i", [-1, 3, 100]) +def test_clip_rows_and_vector_reject_out_of_range_clips(i): + tensors, header = make_index() + index = index_file.Index(header=header, **tensors) + with pytest.raises(IndexError): + index.clip_rows(i) + with pytest.raises(IndexError): + index.clip_vector(i) + + +def set_clip(key, value): + def mutate(header): + header["clips"][0][key] = value + + return mutate + + +def set_top(key, value): + def mutate(header): + header[key] = value + + return mutate + + +def set_run(key, value): + def mutate(header): + header["run"][key] = value + + return mutate + + +def duplicate_paths(header): + header["clips"][1]["path"] = header["clips"][0]["path"] + + +BAD_HEADER_FIELDS = [ + pytest.param(set_top("root", "videos"), "root must be", id="root-relative"), + pytest.param(set_top("root", "~/videos"), "root must be", id="root-tilde"), + pytest.param(set_top("root", ""), "root must be", id="root-empty"), + pytest.param(set_top("root", 5), "root must be", id="root-not-a-string"), + pytest.param(set_top("format_version", True), "format_version", id="format-version-bool"), + pytest.param(set_top("skipped", "x"), "skipped must be", id="skipped-not-a-list"), + pytest.param(set_top("skipped", {}), "skipped must be", id="skipped-empty-object"), + pytest.param(set_top("skipped", [1, "x"]), "skipped must be", id="skipped-not-objects"), + pytest.param(set_top("skipped", [{"path": "a.mp4"}]), "skipped must be", id="skipped-no-reason"), + pytest.param(set_top("skipped", [{"path": 3, "reason": "r"}]), "skipped must be", id="skipped-path-type"), + pytest.param(set_top("clips", {}), "clips must be a list", id="clips-not-a-list"), + pytest.param(set_top("clips", [5]), r"clips\[0\] must be an object", id="clip-not-an-object"), + pytest.param(set_top("run", []), "run must be", id="run-not-an-object"), + pytest.param(set_run("max_seconds", float("nan")), "max_seconds", id="max-seconds-nan"), + pytest.param(set_run("max_seconds", float("inf")), "max_seconds", id="max-seconds-inf"), + pytest.param(set_run("max_seconds", 0), "max_seconds", id="max-seconds-zero"), + pytest.param(set_run("max_seconds", "30"), "max_seconds", id="max-seconds-string"), + pytest.param(set_run("max_seconds", True), "max_seconds", id="max-seconds-bool"), + pytest.param(set_clip("sha256", 5), "sha256", id="sha256-not-a-string"), + pytest.param(set_clip("sha256", "zz"), "sha256", id="sha256-short"), + pytest.param(set_clip("sha256", "A" * 64), "sha256", id="sha256-uppercase"), + pytest.param(set_clip("size", -1), "size", id="size-negative"), + pytest.param(set_clip("size", True), "size", id="size-bool"), + pytest.param(set_clip("mtime_ns", 1.5), "mtime_ns", id="mtime-float"), + pytest.param(set_clip("mtime_ns", -1), "mtime_ns", id="mtime-negative"), + pytest.param(set_clip("width", 0), "width", id="width-zero"), + pytest.param(set_clip("width", 3.0), "width", id="width-float"), + pytest.param(set_clip("height", -5), "height", id="height-negative"), + pytest.param(set_clip("duration", float("nan")), "duration", id="duration-nan"), + pytest.param(set_clip("duration", float("inf")), "duration", id="duration-inf"), + pytest.param(set_clip("duration", -1.0), "duration", id="duration-negative"), + pytest.param(set_clip("duration", 0), "duration", id="duration-zero"), + pytest.param(set_clip("duration", "4"), "duration", id="duration-string"), + pytest.param(set_clip("duration", True), "duration", id="duration-bool"), + pytest.param(set_clip("windows", [0]), "windows must be", id="windows-one-value"), + pytest.param(set_clip("windows", [0, 2.0]), "windows must be", id="windows-float"), + pytest.param(set_clip("windows", [True, 2]), "windows must be", id="windows-bool"), + pytest.param(set_clip("windows", "0:2"), "windows must be", id="windows-string"), + pytest.param(set_clip("path", "./a.mp4"), "path", id="path-dot-segment"), + pytest.param(set_clip("path", "a//b.mp4"), "path", id="path-empty-segment"), + pytest.param(set_clip("path", "sub/"), "path", id="path-trailing-slash"), + pytest.param(set_clip("path", "."), "path", id="path-dot"), + pytest.param(set_clip("path", "C:/x.mp4"), "path", id="path-drive-letter"), + pytest.param(set_clip("path", "a\x00b.mp4"), "path", id="path-nul"), + pytest.param(set_clip("path", 7), "path", id="path-not-a-string"), + pytest.param(duplicate_paths, "same path", id="paths-duplicate"), +] + + +@pytest.mark.parametrize(("mutate", "message"), BAD_HEADER_FIELDS) +def test_rejects_malformed_header_field(mutate, message): + tensors, header = make_index() + mutate(header) + with pytest.raises(IndexFileError, match=message): + index_file.check(tensors, header) + + +def test_rejects_non_integer_hidden_size(): + tensors, header = make_index() + block = {**header["representation"], "hidden_size": float(DIM)} + header["representation"] = {**block, "hash": latent_file.representation_hash(block)} + with pytest.raises(IndexFileError, match="hidden_size"): + index_file.check(tensors, header) + + +def test_paths_with_spaces_and_unicode_are_accepted(): + tensors, header = make_index() + header["clips"][0]["path"] = "Ünï dir/clip 🎬.mp4" + index_file.check(tensors, header) + + +def test_save_refuses_a_nan_in_the_header(tmp_path: Path): + tensors, header = make_index() + header["clips"][0]["duration"] = float("nan") + with pytest.raises(IndexFileError, match="duration"): + index_file.save(tmp_path / "x.index", tensors, header) + + +@pytest.mark.parametrize("constant", ["NaN", "Infinity", "-Infinity"]) +def test_read_header_refuses_json_constants(tmp_path: Path, constant): + tensors, header = make_index() + text = json.dumps(header, sort_keys=True).replace('"duration": 4.0', f'"duration": {constant}', 1) + assert constant in text + path = tmp_path / "nan.index" + save_file(tensors, str(path), metadata={"jepa_scope": text}) + with pytest.raises(IndexFileError, match="not allowed|not valid JSON"): + index_file.read_header(path) + + +@pytest.mark.parametrize("text", ["[1, 2]", "{broken"], ids=["not-an-object", "not-json"]) +def test_read_header_refuses_a_non_object_or_broken_header(tmp_path: Path, text): + tensors, _ = make_index() + path = tmp_path / "odd.index" + save_file(tensors, str(path), metadata={"jepa_scope": text}) + with pytest.raises(IndexFileError, match="JSON"): + index_file.read_header(path) + + +def test_read_header_of_a_missing_file_raises_like_latent_file(tmp_path: Path): + # latent_file.read_header lets the missing-file error through; the index does the same. + with pytest.raises(FileNotFoundError): + latent_file.read_header(tmp_path / "missing.latent") + with pytest.raises(FileNotFoundError): + index_file.read_header(tmp_path / "missing.index") + + +def test_failed_save_leaves_no_temporary_file(tmp_path: Path): + tensors, header = make_index() + target = tmp_path / "target.index" + target.mkdir() # os.replace of a file onto a directory fails + with pytest.raises(OSError): + index_file.save(target, tensors, header) + assert sorted(p.name for p in tmp_path.iterdir()) == ["target.index"] + + +def test_saves_use_distinct_temporary_names(tmp_path: Path, monkeypatch): + tensors, header = make_index() + seen = [] + real_replace = os.replace + + def spy(src, dst): + seen.append(Path(src).name) + real_replace(src, dst) + + monkeypatch.setattr(os, "replace", spy) + index_file.save(tmp_path / "a.index", tensors, header) + index_file.save(tmp_path / "a.index", tensors, header) + assert len(set(seen)) == 2 + assert sorted(p.name for p in tmp_path.iterdir()) == ["a.index"] + + +def test_resolve_root_makes_the_colocated_path_absolute(tmp_path: Path, monkeypatch): + tensors, header = make_index() + index = index_file.Index(header=header, **tensors) + (tmp_path / "folder").mkdir() + monkeypatch.chdir(tmp_path) + root = index.resolve_root(Path("folder") / "jepa-scope.index") + assert root.is_absolute() + assert root == (tmp_path / "folder").resolve() + + +def test_save_writes_strict_json_even_if_the_check_is_bypassed(tmp_path: Path, monkeypatch): + # The check already rejects a NaN; this pins the second line of defence in the writer itself. + monkeypatch.setattr(index_file, "check", lambda tensors, header: None) + tensors, header = make_index() + header["clips"][0]["duration"] = float("nan") + with pytest.raises(IndexFileError, match="cannot be written as JSON"): + index_file.save(tmp_path / "x.index", tensors, header) + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.parametrize( + ("path", "ok"), + [ + ("cam 2026-11-06 10:15:00.mp4", True), # a colon in a top-level timestamped name + ("day/10:15:00.mp4", True), + ("a:b.mp4", True), + ("C:/x.mp4", False), + ("C:", False), + ("sub\\x.mp4", False), + ], +) +def test_is_valid_clip_path(path, ok): + assert index_file.is_valid_clip_path(path) is ok + + +def test_saved_index_follows_the_umask(tmp_path): + import os + + tensors, header = make_index() + old = os.umask(0o077) + try: + index_file.save(tmp_path / "private.index", tensors, header) + finally: + os.umask(old) + index_file.save(tmp_path / "shared.index", tensors, header) + assert (tmp_path / "private.index").stat().st_mode & 0o777 == 0o600 + assert (tmp_path / "shared.index").stat().st_mode & 0o777 == 0o666 & ~old + + +def test_saving_through_a_symlink_writes_beside_the_link(tmp_path, monkeypatch): + import tempfile + + made_in = [] + real_mkstemp = tempfile.mkstemp + + def recording_mkstemp(*args, **kwargs): + made_in.append(Path(kwargs["dir"])) + return real_mkstemp(*args, **kwargs) + + monkeypatch.setattr(index_file.tempfile, "mkstemp", recording_mkstemp) + tensors, header = make_index() + target_dir = tmp_path / "elsewhere" + target_dir.mkdir() + (target_dir / "real.index").write_bytes(b"old") + link = tmp_path / "jepa-scope.index" + link.symlink_to(target_dir / "real.index") + index_file.save(link, tensors, header) + assert not link.is_symlink() and index_file.load(link).pooled.shape == tensors["pooled"].shape + assert not list(tmp_path.glob("*.tmp")) and not list(target_dir.glob("*.tmp")) + # Beside the link, so a link to another filesystem (an external drive) still renames in place. + assert made_in == [tmp_path] From 70ffac9015aea3f97cbf7349d4f7c0ee0a859e5c Mon Sep 17 00:00:00 2001 From: RDCLder Date: Fri, 6 Nov 2026 13:40:16 -0600 Subject: [PATCH 2/6] Add brute-force search over the folder index prepare turns an index into normalised window vectors, centred on the folder mean when the index asks for it; vectors that are zero or not finite after that are unavailable rather than ranked. nearest_moments finds windows like a given window, in other clips by default, keeping only non-overlapping hits per clip; nearest_clips compares clips by the mean of their window vectors. Candidates found with a float32 product are rescored in float64 so equal scores order by clip and window on any maths library, and overlap is decided by window number, not timestamps. Signed-off-by: RDCLder --- ...6-weeks3-4-folder-search-trajectory-log.md | 21 + ...11-06-weeks3-4-folder-search-trajectory.md | 2 +- src/jepa_scope/search.py | 425 ++++++++++++++ tests/test_search.py | 529 ++++++++++++++++++ 4 files changed, 976 insertions(+), 1 deletion(-) create mode 100644 src/jepa_scope/search.py create mode 100644 tests/test_search.py diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md index 81b495a..4c61078 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md @@ -127,3 +127,24 @@ The plan is `2026-11-06-weeks3-4-folder-search-trajectory.md`. Entries are newes - **Impact:** T3 builds on `Index`, `clip_rows` and `clip_vector`. - **Resolution:** committed on `week3-index`. - **Permanent decision record needed:** Yes, at T10 (the index file as a versioned contract). + +- **2026-11-06 — T3 complete: brute-force search (`search.py`).** + - **Evidence:** an executor (Sonnet) wrote the module; a fresh reviewer (Opus) did two rounds. + Round 1 found two real defects. First, ties were not ordered by `(clip, window)` at the + model's width on this Mac: Accelerate's float32 product gives exact duplicates scores 1 ulp + apart depending on row position (16 of 100 seeds out of order); candidates are now rescored + in float64 and ranked on that. Second, overlap was decided by start-time difference, which + real timestamp jitter breaks (about half of non-overlapping pairs looked overlapping); it is + now decided by window number (`|w1 - w2| * stride < frames`). It also had the centring mean + recomputed in float64 so identical clips are reliably "unavailable", fixed the reason + wording, and added `DEFAULT_SIMILARITY` and `nearest_window` for T4. Round 2 passed; its + three small follow-ups (a suppression test, a clipping test with teeth, and `MAX_K = 1000` + so one request cannot allocate hundreds of megabytes) were added by the lead. + - **Measured (estimates for the plan, this laptop, 30,000 × 1024):** `prepare` about 0.1 s; + a moment query about 1.5 ms; a clip query about 1 ms. The plan's 30 ms estimate was ten times + too pessimistic. + - **Deviation:** the search tests are in `tests/test_search.py`, not `tests/test_index.py`. + - **Impact:** T4's `search` command uses `query_vector`, `search_moments`, `search_clips` and + `nearest_window`; T6 maps `k > MAX_K` to a 400. T4 should repeat the window-weighted-mean + check against the indexer's output. + - **Permanent decision record needed:** Yes, at T10 (search as plain cosine over window means). diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md index 3257865..121e361 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md @@ -594,7 +594,7 @@ flowchart LR ### T3 — Week 3. Search on hand-built vectors -- **Status:** Pending +- **Status:** Complete - **Dependencies:** T2 - **Parallel group:** Sequential - **Complexity:** ROUTINE diff --git a/src/jepa_scope/search.py b/src/jepa_scope/search.py new file mode 100644 index 0000000..83f91ec --- /dev/null +++ b/src/jepa_scope/search.py @@ -0,0 +1,425 @@ +"""Brute-force search over an index: moments like a moment, and clips like a clip. + +Every window vector in an index is L2-normalised once, in ``prepare``, and a query is one +matrix-vector product. When the index says ``centred-cosine`` (or the caller asks for it), the +index mean is subtracted from every vector first, so a direction shared by the whole folder does not +squeeze all the scores into a narrow band. Centring is an explicit, recorded choice: the plain +cosine over the stored window vectors is the default. + +Two kinds of query share that primitive. + +* Moments: the query is one window, and the hits are other windows. Within one clip, a hit that + overlaps a higher-ranked hit is dropped, so a list of ten hits is ten distinct, non-overlapping + moments. +* Clips: the query is a clip's vector, the mean of its raw (or centred) window vectors normalised + after averaging, and the hits are other clips. + +A vector whose norm is below ``MIN_NORM`` (after centring when centring is on), or that is not +finite, is "unavailable". It is never a candidate, and a query made from it returns an empty +result that says why. Equal scores are ordered by clip and then by window, so duplicates rank the +same way every time. + +Two windows of one clip overlap when their window numbers are closer than the window length allows: +``abs(w1 - w2) * stride < frames``, with ``stride`` and ``frames`` from the index header. This is +decided by window number, never by start times, which are real frame timestamps and jitter. + +Scores are computed in float32 to pick candidates and then recomputed in float64 for the +candidates, so rows that are exact duplicates get exactly equal scores and tie in +(clip, window) order whatever their position in the matrix. + +Window numbers in this module are clip-local, counted from 0 within a clip. Every moment hit also +carries ``row``, the window's row in the index's tensors. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Generic, TypeVar + +import numpy as np + +from .index_file import Index + +# The similarity a new index is built with unless the caller asks for centring. It is the choice +# recorded in the execution log of the folder-search plan (T0): plain cosine over the stored window +# means. "centred-cosine" is the other accepted value of the header's ``embedding.similarity``. +DEFAULT_SIMILARITY = "cosine" + +MIN_NORM = 1e-6 +# The most hits one search may ask for. Rescoring candidates in float64 copies them, so an +# unbounded k (one API request) could allocate hundreds of megabytes. +MAX_K = 1000 +_SCORE_SLACK = 1e-6 # float32 products of duplicate rows can differ by an ulp or two; keep both +_CHUNK = 4096 # rows normalised at a time, so float64 working copies stay small + +T = TypeVar("T") + + +@dataclass(frozen=True) +class Hit: + """One moment: window ``window`` of clip ``clip``, at index-wide row ``row``. + + ``start_time`` is when the window starts in its clip, in seconds; ``score`` is the cosine + similarity to the query (of centred vectors when the search is centred). + """ + + clip: int + window: int + row: int + start_time: float + score: float + + +@dataclass(frozen=True) +class ClipHit: + """One clip in a "clips like this clip" result, with its cosine similarity to the query.""" + + clip: int + score: float + + +@dataclass(frozen=True) +class Result(Generic[T]): # noqa: UP046 (kept as a plain Generic so the module reads simply) + """What a search returns: ``hits`` best first, and ``reason``. + + ``reason`` is None unless the result is empty, in which case it says why in a sentence (the + query is unavailable, the index has no windows, or no candidate is left after exclusion). A + non-empty result with fewer than k hits means the index ran out of candidates. + """ + + hits: list[T] + reason: str | None = None + + +@dataclass(frozen=True) +class Query: + """A normalised query vector, or None with the reason it is unavailable.""" + + vector: np.ndarray | None + reason: str | None = None + + +@dataclass +class SearchSpace: + """An index prepared for search. Build it with ``prepare``; the fields are read-only in use. + + ``matrix`` holds the normalised window vectors (float32, ``[N, D]``; unavailable rows are + zero). ``available`` marks the rows that may be candidates. ``clip``, ``window``, + ``start_time`` describe each row (``window`` is clip-local). ``clip_start`` and ``clip_stop`` + give each clip's row range. ``clip_vectors`` and ``clip_available`` are the per-clip vectors. + ``window_seconds`` is the length of one window (used to find a window at a time). ``frames`` + and ``stride`` are the window length and step in frames; two windows of a clip overlap when + ``abs(w1 - w2) * stride < frames``. ``mean`` is the vector subtracted before normalising, or None + for plain cosine. + """ + + matrix: np.ndarray + available: np.ndarray + clip: np.ndarray + window: np.ndarray + start_time: np.ndarray + clip_start: np.ndarray + clip_stop: np.ndarray + clip_vectors: np.ndarray + clip_available: np.ndarray + window_seconds: float + frames: int + stride: int + mean: np.ndarray | None + + @property + def centred(self) -> bool: + """Whether the vectors were centred on the index mean before normalising.""" + return self.mean is not None + + @property + def n_clips(self) -> int: + return len(self.clip_start) + + def overlap(self, w1: int, w2: int) -> bool: + """Whether clip-local windows ``w1`` and ``w2`` of one clip share frames.""" + return abs(int(w1) - int(w2)) * self.stride < self.frames + + +def _normalise(x: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Normalise rows of float64 ``x``; return (float32 rows, available flags). + + A row that is not finite, or whose norm is below ``MIN_NORM``, becomes a zero row and is not + available. + """ + finite = np.isfinite(x).all(axis=1) + safe = np.where(finite[:, None], x, 0.0) + norms = np.sqrt((safe * safe).sum(axis=1)) + ok = finite & (norms >= MIN_NORM) + out = np.zeros(x.shape, dtype=np.float32) + out[ok] = (safe[ok] / norms[ok, None]).astype(np.float32) + return out, ok + + +def prepare(index: Index, centre: bool | None = None) -> SearchSpace: + """Prepare ``index`` for search. + + Vectors are centred (the index ``mean`` is subtracted, in float64) when ``centre`` is True, or + when ``centre`` is None and the header's ``embedding.similarity`` is ``"centred-cosine"``; + otherwise they are used as stored. The centring mean is the window-weighted mean of the rows, + recomputed in float64. Rows are then L2-normalised. Clip vectors are the mean of + the clip's raw or centred rows, normalised after averaging (never an average of normalised + rows). + """ + if centre is None: + centre = index.header["embedding"]["similarity"] == "centred-cosine" + pooled = index.pooled + n, dim = pooled.shape + # The centring mean is recomputed in float64 from the rows, not read from the file: the stored + # mean may have been accumulated in float32, and ten identical rows must centre to exactly zero. + mean = pooled.mean(axis=0, dtype=np.float64) if centre and n else None + if centre and mean is None: + mean = np.zeros(dim, dtype=np.float64) + matrix = np.zeros((n, dim), dtype=np.float32) + available = np.zeros(n, dtype=bool) + for lo in range(0, n, _CHUNK): + block = pooled[lo : lo + _CHUNK].astype(np.float64) + if mean is not None: + block -= mean + matrix[lo : lo + _CHUNK], available[lo : lo + _CHUNK] = _normalise(block) + + clips = index.header["clips"] + clip_start = np.array([c["windows"][0] for c in clips], dtype=np.int64) + clip_stop = np.array([c["windows"][1] for c in clips], dtype=np.int64) + raw_means = np.zeros((len(clips), dim), dtype=np.float64) + for i in range(len(clips)): + raw_means[i] = pooled[clip_start[i] : clip_stop[i]].mean(axis=0, dtype=np.float64) + if mean is not None: + raw_means -= mean # the mean of centred rows is the centred mean of the rows + clip_vectors, clip_available = _normalise(raw_means) + + window_info = index.header["representation"]["window"] + window_seconds = float(window_info["frames"]) / float(window_info["sample_fps"]) + local = np.arange(n, dtype=np.int64) - clip_start[index.clip] if n else np.zeros(0, dtype=np.int64) + return SearchSpace( + matrix=matrix, + available=available, + clip=index.clip.astype(np.int64), + window=local, + start_time=index.window_start_time, + clip_start=clip_start, + clip_stop=clip_stop, + clip_vectors=clip_vectors, + clip_available=clip_available, + window_seconds=window_seconds, + frames=int(window_info["frames"]), + stride=int(window_info["stride"]), + mean=mean, + ) + + +def _check_k(k: int) -> None: + if k < 1: + raise ValueError(f"k must be at least 1, not {k}") + if k > MAX_K: + raise ValueError(f"k must be at most {MAX_K}, not {k}") + + +def _check_clip(space: SearchSpace, clip: int) -> None: + if not 0 <= clip < space.n_clips: + raise IndexError(f"clip {clip} is out of range for an index of {space.n_clips} clips") + + +def _row_of(space: SearchSpace, clip: int, window: int) -> int: + _check_clip(space, clip) + count = int(space.clip_stop[clip] - space.clip_start[clip]) + if not 0 <= window < count: + raise IndexError(f"window {window} is out of range for clip {clip}, which has {count} windows") + return int(space.clip_start[clip]) + window + + +def nearest_window(start_times: np.ndarray, window_seconds: float, t: float) -> int: + """The position in ``start_times`` of the window whose centre is nearest to time ``t``. + + A window's centre is its start time plus half of ``window_seconds``. When ``t`` is exactly + halfway between two centres the earlier window wins, and a time before the first centre or after + the last one gives the first or last window. ``start_times`` must be non-empty and increasing. + Raises ``ValueError`` when ``t`` is not finite or there are no windows. + """ + starts = np.asarray(start_times, dtype=np.float64) + if len(starts) == 0: + raise ValueError("there are no windows to choose from") + if not np.isfinite(t): + raise ValueError(f"t must be finite, not {t}") + centres = starts + window_seconds / 2 + return int(np.argmin(np.abs(centres - t))) # argmin returns the first minimum: the earlier window + + +def window_at(space: SearchSpace, clip: int, t: float) -> int: + """The clip-local window of ``clip`` whose centre is nearest to time ``t`` (seconds). + + See ``nearest_window`` for the tie and clamping rules. Raises ``IndexError`` for an unknown + clip and ``ValueError`` when ``t`` is not finite. + """ + _check_clip(space, clip) + return nearest_window(space.start_time[space.clip_start[clip] : space.clip_stop[clip]], space.window_seconds, t) + + +def _rank(space: SearchSpace, scores: np.ndarray, vector: np.ndarray, eligible: np.ndarray, k: int): + """The best ``k`` non-overlapping moments among the ``eligible`` rows, as (row, score) pairs. + + ``scores`` (float32) only chooses candidates; the candidates are then rescored in float64 and + cast to float32, so duplicate rows tie exactly. Rows are taken best first, ties by + (clip, window), and a row is dropped when it overlaps a kept row of the same clip. The candidate + set starts at 2k rows and doubles until k rows are kept or every eligible row has been seen, so + the result equals a full sort. + """ + n = len(eligible) + m = 2 * k + v64 = vector.astype(np.float64) + while True: + if m >= n: + cand = eligible + else: + top = scores[eligible] + threshold = np.partition(top, n - m)[n - m] + # Everything at or above the cutoff (less a little slack for rounding), so a tie never + # depends on partition order or on a row's position in the matrix. + cand = eligible[top >= threshold - _SCORE_SLACK] + exact = np.clip((space.matrix[cand].astype(np.float64) @ v64).astype(np.float32), -1.0, 1.0) + order = np.lexsort((space.window[cand], space.clip[cand], -exact)) + kept: list[tuple[int, float]] = [] + kept_windows: dict[int, list[int]] = {} + for pos in order: + row = int(cand[pos]) + c = int(space.clip[row]) + w = int(space.window[row]) + if any(space.overlap(w, other) for other in kept_windows.get(c, ())): + continue + kept.append((row, float(exact[pos]))) + kept_windows.setdefault(c, []).append(w) + if len(kept) == k: + return kept + if m >= n: + return kept + m *= 2 + + +def _moments( + space: SearchSpace, vector: np.ndarray, k: int, query_clip: int | None, query_row: int | None, same_clip: bool +) -> Result[Hit]: + if len(space.matrix) == 0: + return Result([], "the index has no windows") + scores = space.matrix @ vector + mask = space.available.copy() + if query_clip is not None: + in_clip = space.clip == query_clip + if not same_clip: + mask &= ~in_clip + else: + assert query_row is not None + near = np.abs(space.window - space.window[query_row]) * space.stride < space.frames + mask &= ~(in_clip & near) # the query window itself is in this set too + eligible = np.flatnonzero(mask) + if len(eligible) == 0: + return Result([], "no other window is available to compare with") + hits = [ + Hit(int(space.clip[r]), int(space.window[r]), r, float(space.start_time[r]), score) + for r, score in _rank(space, scores, vector, eligible, k) + ] + return Result(hits) + + +def _unavailable(space: SearchSpace, what: str) -> str: + """The reason given for an unavailable query, worded for plain or centred search.""" + if space.centred: + return f"{what} is unavailable: its vector is zero (or not finite) after subtracting the index mean" + return f"{what} is unavailable: its vector is zero or not finite" + + +def nearest_moments(space: SearchSpace, clip: int, window: int, k: int, same_clip: bool = False) -> Result[Hit]: + """The ``k`` windows most like window ``window`` (clip-local) of clip ``clip``. + + The query's own clip is left out entirely unless ``same_clip`` is True; with ``same_clip``, + the query window and every window of its clip that overlaps it are left out. Within one clip a + hit that overlaps a higher-ranked hit is dropped. Fewer than ``k`` hits come back only when + the index has no more candidates. Equal scores are ordered by (clip, window). + + An unavailable query window gives an empty result with a reason. Raises ``ValueError`` for + ``k < 1`` and ``IndexError`` for an unknown clip or window. + """ + _check_k(k) + row = _row_of(space, clip, window) + if not space.available[row]: + return Result([], _unavailable(space, "the query window")) + return _moments(space, space.matrix[row], k, clip, row, same_clip) + + +def nearest_clips(space: SearchSpace, clip: int, k: int) -> Result[ClipHit]: + """The ``k`` clips most like clip ``clip``, other clips only, by their clip vectors. + + Clips whose vector is unavailable are never hits; an unavailable query clip gives an empty + result with a reason. Equal scores are ordered by clip. Raises ``ValueError`` for ``k < 1`` and + ``IndexError`` for an unknown clip. + """ + _check_k(k) + _check_clip(space, clip) + if not space.clip_available[clip]: + return Result([], _unavailable(space, "the query clip")) + return _clips(space, space.clip_vectors[clip], k, clip) + + +def _clips(space: SearchSpace, vector: np.ndarray, k: int, query_clip: int | None) -> Result[ClipHit]: + mask = space.clip_available.copy() + if query_clip is not None: + mask[query_clip] = False + eligible = np.flatnonzero(mask) + if len(eligible) == 0: + return Result([], "no other clip is available to compare with") + # Float64 products, cast to float32, so identical clip vectors tie exactly in any position. + scores = np.clip( + (space.clip_vectors[eligible].astype(np.float64) @ vector.astype(np.float64)).astype(np.float32), -1.0, 1.0 + ) + order = np.lexsort((eligible, -scores))[:k] + return Result([ClipHit(int(eligible[i]), float(scores[i])) for i in order]) + + +def query_vector(space: SearchSpace, pooled_rows: np.ndarray) -> Query: + """A query vector for a clip that is not in the index, from its raw window vectors. + + ``pooled_rows`` is ``[W, D]`` (the clip's window vectors as ``.latent`` stores them). They are + averaged, centred when the space is, and normalised, the same way ``prepare`` makes a clip + vector. If the result is unavailable (no rows, not finite, or a zero vector) the returned + ``Query`` has ``vector`` None and a reason. Raises ``ValueError`` when the width is not the + index's. + """ + rows = np.asarray(pooled_rows) + dim = space.matrix.shape[1] + if rows.ndim != 2 or rows.shape[1] != dim: + raise ValueError(f"pooled_rows must have shape [W, {dim}], not {tuple(rows.shape)}") + if len(rows) == 0: + return Query(None, "the query has no windows") + avg = rows.mean(axis=0, dtype=np.float64) + if space.mean is not None: + avg = avg - space.mean + vec, ok = _normalise(avg[None, :]) + if not ok[0]: + return Query(None, _unavailable(space, "the query vector")) + return Query(vec[0]) + + +def search_moments(space: SearchSpace, query: Query, k: int) -> Result[Hit]: + """The ``k`` windows of the index most like an external ``query`` (from ``query_vector``). + + No clip is excluded, since the query is not in the index; overlapping hits within one clip are + still suppressed. An unavailable query gives an empty result with its reason. + """ + _check_k(k) + if query.vector is None: + return Result([], query.reason) + return _moments(space, query.vector, k, None, None, True) + + +def search_clips(space: SearchSpace, query: Query, k: int) -> Result[ClipHit]: + """The ``k`` clips of the index most like an external ``query`` (from ``query_vector``).""" + _check_k(k) + if query.vector is None: + return Result([], query.reason) + if space.n_clips == 0: + return Result([], "the index has no windows") + return _clips(space, query.vector, k, None) diff --git a/tests/test_search.py b/tests/test_search.py new file mode 100644 index 0000000..9602bd2 --- /dev/null +++ b/tests/test_search.py @@ -0,0 +1,529 @@ +"""Search over hand-built vectors: ranking rules, exclusion, centring and error handling. + +Every index here is built from vectors written in the test; no encoder is involved. A window is +2 seconds long (16 frames at 8 fps) and a window starts every second (stride 8), so only +neighbouring windows of a clip overlap. Each test says in its name or docstring whether it runs +centred or plain. +""" + +from __future__ import annotations + +import dataclasses + +import numpy as np +import pytest + +from jepa_scope import index_file, search +from jepa_scope.index_file import Index +from tests.test_index import make_representation + +DIM = 8 + + +def build(clips, similarity: str = "cosine", dim: int | None = None) -> Index: + """An index from ``clips``, a list of ``(vectors, start_times)``; vectors are ``[W, D]``.""" + dim = dim or (len(clips[0][0][0]) if clips else DIM) + if clips: + pooled = np.concatenate([np.asarray(v, dtype=np.float32).reshape(-1, dim) for v, _ in clips]) + times = np.concatenate([np.asarray(t, dtype=np.float64) for _, t in clips]) + else: + pooled, times = np.zeros((0, dim), np.float32), np.zeros(0) + counts = [len(t) for _, t in clips] + owner = np.repeat(np.arange(len(clips), dtype=np.int32), counts) + mean = pooled.mean(axis=0, dtype=np.float64).astype(np.float32) if len(pooled) else np.zeros(dim, np.float32) + tensors = {"pooled": pooled, "window_start_time": times, "clip": owner, "mean": mean} + entries, start = [], 0 + for i, c in enumerate(counts): + entries.append( + { + "path": f"clip{i}.mp4", + "sha256": f"{i:064x}", + "size": 1, + "mtime_ns": 1, + "duration": 10.0, + "width": 8, + "height": 8, + "windows": [start, start + c], + } + ) + start += c + header = { + "format": index_file.INDEX_FORMAT, + "format_version": index_file.INDEX_FORMAT_VERSION, + "representation": make_representation(dim), + "run": {"max_seconds": None, "device": "cpu", "dtype": "float32"}, + "embedding": {"kind": "window_vector", "similarity": similarity}, + "root": ".", + "skipped": [], + "clips": entries, + } + index_file.check(tensors, header) + return Index(header=header, pooled=pooled, window_start_time=times, clip=owner, mean=mean) + + +def unit(i: int, dim: int = DIM) -> np.ndarray: + v = np.zeros(dim) + v[i] = 1.0 + return v + + +def random_clips(n_clips: int, windows: int, seed: int = 0, spacing: float = 1.0): + rng = np.random.default_rng(seed) + return [(rng.normal(size=(windows, DIM)), np.arange(windows) * spacing) for _ in range(n_clips)] + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_duplicate_scores_one_and_ranks_first(centre): + clips = random_clips(4, 3, seed=1) + clips[2] = (np.concatenate([clips[2][0][:2], clips[0][0][1:2]]), clips[2][1]) # clip 2 window 2 = clip 0 window 1 + space = search.prepare(build(clips), centre=centre) + res = search.nearest_moments(space, 0, 1, k=5) + assert (res.hits[0].clip, res.hits[0].window) == (2, 2) + assert res.hits[0].score == pytest.approx(1.0, abs=1e-6) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_two_duplicates_tie_in_clip_window_order(centre): + clips = random_clips(5, 2, seed=2) + q = clips[0][0][0] + for c in (4, 1, 3): # plant the duplicate in clips 1, 3 and 4, window 1 + clips[c][0][1] = q + space = search.prepare(build(clips), centre=centre) + res = search.nearest_moments(space, 0, 0, k=3) + assert [(h.clip, h.window) for h in res.hits] == [(1, 1), (3, 1), (4, 1)] + assert len({h.score for h in res.hits}) == 1 + + +def test_scaled_row_scores_one_plain(): + clips = random_clips(3, 2, seed=3) + clips[1][0][0] = 2.0 * clips[0][0][0] + space = search.prepare(build(clips), centre=False) + hit = search.nearest_moments(space, 0, 0, k=1).hits[0] + assert (hit.clip, hit.window) == (1, 0) + assert hit.score == pytest.approx(1.0, abs=1e-6) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_same_clip_excludes_query_and_overlaps(centre): + # Clip 0 has windows every second, so each window overlaps its neighbours (2 s long). + clips = [(np.random.default_rng(4).normal(size=(8, DIM)), np.arange(8) * 1.0)] + random_clips(2, 2, seed=5) + space = search.prepare(build(clips), centre=centre) + res = search.nearest_moments(space, 0, 3, k=20, same_clip=True) + own = sorted(h.window for h in res.hits if h.clip == 0) + assert all(abs(w - 3) >= 2 for w in own) # starts differ by at least the 2 s window length + assert 3 not in own and 2 not in own and 4 not in own + assert any(h.clip != 0 for h in res.hits) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_not_same_clip_excludes_query_clip_entirely(centre): + space = search.prepare(build(random_clips(4, 3, seed=6)), centre=centre) + res = search.nearest_moments(space, 1, 2, k=50) + assert res.hits and all(h.clip != 1 for h in res.hits) + assert len(res.hits) >= 3 # at least one non-overlapping window from each of the other clips + single = search.prepare(build(random_clips(4, 1, seed=6)), centre=centre) + assert sorted(h.clip for h in search.nearest_moments(single, 1, 0, k=50).hits) == [0, 2, 3] + + +def test_overlapping_hits_in_one_clip_collapse_to_the_better_plain(): + q = unit(0) + better, worse = unit(0) + 0.1 * unit(1), unit(0) + 0.5 * unit(1) + other = unit(0) + 0.9 * unit(1) + clips = [([q], [0.0]), ([worse, better], [0.0, 1.0]), ([other], [0.0])] # clip 1 windows overlap + space = search.prepare(build(clips), centre=False) + res = search.nearest_moments(space, 0, 0, k=5) + assert [(h.clip, h.window) for h in res.hits] == [(1, 1), (2, 0)] + + +def test_widening_still_returns_k_when_suppression_removes_most_plain(): + q = unit(0) + # Clip 1 has twelve windows that all outrank the other clips, in an order (1, 0, 2, 4, 3, 5, ...) + # where the best 2k=10 candidates keep only four non-overlapping hits, so k=5 must widen. + order = [1, 0, 2, 4, 3, 5, 7, 6, 8, 10, 9, 11] + heavy = np.zeros((12, DIM)) + for rank, w in enumerate(order): + heavy[w] = unit(0) + (0.05 + 0.01 * rank) * unit(1) + clips = [([q], [0.0]), (heavy, np.arange(12) * 1.0)] + for i in range(4): # four single-window clips ranked lower + clips.append(([unit(0) + (1.0 + i) * unit(2)], [0.0])) + space = search.prepare(build(clips), centre=False) + res = search.nearest_moments(space, 0, 0, k=5) + assert len(res.hits) == 5 + assert [h.window for h in res.hits if h.clip == 1] == [1, 4, 7, 10] + assert res.hits[4].clip == 2 + full = search.nearest_moments(space, 0, 0, k=100) + assert [h.row for h in res.hits] == [h.row for h in full.hits[:5]] + + +def test_k_larger_than_candidates_returns_all_that_exist(): + space = search.prepare(build(random_clips(3, 3, seed=8)), centre=False) + res = search.nearest_moments(space, 0, 0, k=100) + assert len(res.hits) == 4 and res.reason is None # windows 0 and 2 of each of the two other clips + + +def test_common_offset_centred_recovers_structure_and_plain_does_not(): + """A large offset shared by every vector hides the planted structure from plain cosine.""" + rng = np.random.default_rng(0) + offset = 10.0 * unit(7) + clips = [] + for _ in range(5): # clips 1 to 4 carry the query clip's structure, with noise + clips.append((offset + unit(0) + rng.normal(scale=0.4, size=(1, DIM)), [0.0])) + for _ in range(4): # clips 5 to 8 are almost pure offset + clips.append((offset + 0.3 * unit(1) + rng.normal(scale=0.05, size=(1, DIM)), [0.0])) + index = build(clips) + centred = search.nearest_moments(search.prepare(index, centre=True), 0, 0, k=4) + plain = search.nearest_moments(search.prepare(index, centre=False), 0, 0, k=4) + assert {h.clip for h in centred.hits} == {1, 2, 3, 4} + assert any(h.clip >= 5 for h in plain.hits) + + +def test_header_decides_centring_when_not_given(): + index = build(random_clips(3, 2, seed=10), similarity="centred-cosine") + assert search.prepare(index).centred + assert not search.prepare(index, centre=False).centred + assert not search.prepare(build(random_clips(3, 2, seed=10))).centred + + +def test_ten_identical_clips_centred_are_unavailable_plain_works(): + v = np.random.default_rng(11).normal(size=(1, DIM)) + index = build([(v, [0.0]) for _ in range(10)]) + centred = search.prepare(index, centre=True) + assert not centred.available.any() and not centred.clip_available.any() + res = search.nearest_moments(centred, 0, 0, k=3) + assert res.hits == [] and res.reason + clips = search.nearest_clips(centred, 0, k=3) + assert clips.hits == [] and clips.reason + plain = search.prepare(index, centre=False) + res = search.nearest_moments(plain, 0, 0, k=3) + assert len(res.hits) == 3 and res.hits[0].score == pytest.approx(1.0, abs=1e-6) + + +def test_unavailable_row_is_never_a_candidate_centred(): + # Rows x, -x and 0 average to zero, so the zero row is unavailable after centring. + x = unit(0) + clips = [([x], [0.0]), ([-x], [0.0]), ([0 * x], [0.0])] + index = build(clips) + space = search.prepare(index, centre=True) + zero_row = 2 + assert not space.available[zero_row] + res = search.nearest_moments(space, 0, 0, k=10) + assert zero_row not in [h.row for h in res.hits] + assert search.nearest_moments(space, 2, 0, k=3).reason + assert [h.row for h in res.hits] == [1] + + +def test_near_identical_clips_centred_give_finite_deterministic_scores(): + rng = np.random.default_rng(12) + v = rng.normal(size=(3, DIM)) + noisy = v + 1e-4 * rng.normal(size=v.shape) + index = build([(v, np.arange(3) * 2.0), (noisy, np.arange(3) * 2.0)]) + space = search.prepare(index, centre=True) + a = search.nearest_moments(space, 0, 0, k=3) + b = search.nearest_moments(space, 0, 0, k=3) + assert a == b + assert all(np.isfinite(h.score) and -1.0 - 1e-6 <= h.score <= 1.0 + 1e-6 for h in a.hits) + c = search.nearest_clips(space, 0, k=1) + assert np.isfinite(c.hits[0].score) + + +def test_index_mean_is_window_weighted_and_long_clip_dominates(): + """Documents the weighting using this file's ``build`` helper, not product code. + + The real indexer's mean is checked against the same property in the indexer's own tests. + """ + rng = np.random.default_rng(13) + long_clip = unit(0) * 5 + rng.normal(scale=0.1, size=(10, DIM)) + shorts = [(unit(i) * 5 + rng.normal(scale=0.1, size=(1, DIM)), [0.0]) for i in (1, 2, 3)] + index = build([(long_clip, np.arange(10) * 2.0)] + shorts) + weighted = sum(len(index.pooled[index.clip_rows(i)]) * index.clip_vector(i) for i in range(4)) / 13 + np.testing.assert_allclose(index.mean, weighted, atol=1e-5) + unweighted = np.mean([index.clip_vector(i) for i in range(4)], axis=0) + assert not np.allclose(index.mean, unweighted, atol=1e-2) + to_long = np.linalg.norm(index.mean - index.clip_vector(0)) + assert all(to_long < np.linalg.norm(index.mean - index.clip_vector(i)) for i in (1, 2, 3)) + + +def test_clip_vector_averages_raw_rows_then_normalises(): + rows = np.array([10 * unit(0), unit(1)]) + clips = [(rows, [0.0, 2.0]), ([unit(2) + 0.0 * unit(3)], [0.0]), ([unit(1) + unit(2)], [0.0])] + index = build(clips) + plain = search.prepare(index, centre=False) + raw_mean = rows.mean(axis=0) + np.testing.assert_allclose(plain.clip_vectors[0], raw_mean / np.linalg.norm(raw_mean), atol=1e-6) + wrong = (plain.matrix[:2]).mean(axis=0) # the average of normalised rows would be [.5, .5, ...] + assert not np.allclose(plain.clip_vectors[0], wrong / np.linalg.norm(wrong), atol=1e-3) + centred = search.prepare(index, centre=True) + cm = raw_mean - index.mean.astype(np.float64) + np.testing.assert_allclose(centred.clip_vectors[0], cm / np.linalg.norm(cm), atol=1e-5) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_rows_are_unit_length(centre): + space = search.prepare(build(random_clips(3, 4, seed=14)), centre=centre) + norms = np.linalg.norm(space.matrix, axis=1) + np.testing.assert_allclose(norms, 1.0, atol=1e-6) + np.testing.assert_allclose(np.linalg.norm(space.clip_vectors, axis=1), 1.0, atol=1e-6) + + +def test_window_at_nearest_centre_tie_and_clamp(): + # Windows start at 0, 1, 2, 3 and last 2 s, so their centres are 1, 2, 3, 4. + clips = [(np.eye(DIM)[:4], [0.0, 1.0, 2.0, 3.0])] + space = search.prepare(build(clips), centre=False) + assert search.window_at(space, 0, 1.4) == 0 + assert search.window_at(space, 0, 1.5) == 0 # exact midpoint: the earlier window + assert search.window_at(space, 0, 1.6) == 1 # second half of the overlap: the later window + assert search.window_at(space, 0, 99.0) == 3 + assert search.window_at(space, 0, -5.0) == 0 + with pytest.raises(ValueError): + search.window_at(space, 0, float("nan")) + with pytest.raises(IndexError): + search.window_at(space, 5, 1.0) + + +def test_nearest_window_rule_directly(): + starts = np.array([0.0, 1.0, 2.0, 3.0]) + assert search.nearest_window(starts, 2.0, 1.5) == 0 # exact midpoint of centres 1 and 2: earlier + assert search.nearest_window(starts, 2.0, 1.5000001) == 1 + assert search.nearest_window(starts, 2.0, 100.0) == 3 + assert search.nearest_window(starts, 2.0, -100.0) == 0 + assert search.nearest_window(np.array([5.0]), 2.0, 0.0) == 0 + with pytest.raises(ValueError): + search.nearest_window(starts, 2.0, float("inf")) + with pytest.raises(ValueError): + search.nearest_window(np.array([]), 2.0, 1.0) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_nearest_clips_excludes_query_clip_and_orders_ties(centre): + clips = random_clips(5, 2, seed=15) + clips[3] = (clips[1][0], clips[1][1]) # clips 1 and 3 are identical + clips[4] = (clips[1][0], clips[1][1]) + space = search.prepare(build(clips), centre=centre) + res = search.nearest_clips(space, 0, k=10) + assert [h.clip for h in res.hits if h.clip != 2] == sorted(h.clip for h in res.hits if h.clip != 2) + assert 0 not in [h.clip for h in res.hits] and len(res.hits) == 4 + from_one = search.nearest_clips(space, 1, k=2) + assert [h.clip for h in from_one.hits] == [3, 4] + assert from_one.hits[0].score == pytest.approx(1.0, abs=1e-6) + + +def test_errors_for_k_and_unknown_clip_or_window(): + space = search.prepare(build(random_clips(2, 2, seed=16)), centre=False) + for bad in (0, -1): + with pytest.raises(ValueError): + search.nearest_moments(space, 0, 0, k=bad) + with pytest.raises(ValueError): + search.nearest_clips(space, 0, k=bad) + for clip, window in ((0, 2), (0, -1), (2, 0), (-1, 0)): + with pytest.raises(IndexError): + search.nearest_moments(space, clip, window, k=1) + for clip in (2, -1): + with pytest.raises(IndexError): + search.nearest_clips(space, clip, k=1) + + +def test_empty_index_returns_empty_results_with_reason(): + space = search.prepare(build([], dim=DIM), centre=False) + assert space.n_clips == 0 and space.matrix.shape == (0, DIM) + query = search.query_vector(space, np.ones((1, DIM))) + moments = search.search_moments(space, query, k=3) + clips = search.search_clips(space, query, k=3) + assert moments.hits == [] and moments.reason + assert clips.hits == [] and clips.reason + with pytest.raises(IndexError): + search.nearest_moments(space, 0, 0, k=1) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_hits_carry_clip_local_window_and_index_wide_row(centre): + index = build(random_clips(4, 3, seed=17, spacing=3.0)) + space = search.prepare(index, centre=centre) + res = search.nearest_moments(space, 2, 1, k=8) + assert res.hits + for h in res.hits: + assert h.row == index.clip_rows(h.clip).start + h.window + assert index.clip[h.row] == h.clip + assert h.start_time == index.window_start_time[h.row] + assert h.score == pytest.approx( + float(space.matrix[h.row] @ space.matrix[index.clip_rows(2).start + 1]), abs=1e-6 + ) + scores = [h.score for h in res.hits] + assert scores == sorted(scores, reverse=True) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_query_vector_matches_in_index_clip_vector(centre): + index = build(random_clips(4, 3, seed=18)) + space = search.prepare(index, centre=centre) + for i in range(4): + q = search.query_vector(space, index.pooled[index.clip_rows(i)]) + np.testing.assert_allclose(q.vector, space.clip_vectors[i], atol=1e-6) + hits = search.search_clips(space, search.query_vector(space, index.pooled[index.clip_rows(1)]), k=4) + assert hits.hits[0].clip == 1 and hits.hits[0].score == pytest.approx(1.0, abs=1e-6) + moments = search.search_moments(space, search.query_vector(space, index.pooled[index.clip_rows(1)]), k=3) + assert len(moments.hits) == 3 + + +def test_query_vector_unavailable_and_shape(): + space = search.prepare(build(random_clips(2, 2, seed=19)), centre=False) + zero = search.query_vector(space, np.zeros((2, DIM))) + assert zero.vector is None and zero.reason + assert search.search_moments(space, zero, k=2) == search.Result([], zero.reason) + assert search.query_vector(space, np.zeros((0, DIM))).vector is None + with pytest.raises(ValueError): + search.query_vector(space, np.ones((2, DIM + 1))) + with pytest.raises(ValueError): + search.search_clips(space, zero, k=0) + + +JITTER = np.array([0.0, 0.003, 0.0, -0.003, 0.0, 0.0]) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_same_clip_exclusion_ignores_start_time_jitter(centre): + """Windows 1 and 3 share no frames, but jitter puts their start times 1.994 s apart.""" + rng = np.random.default_rng(20) + vecs = rng.normal(size=(6, DIM)) * 0.3 + vecs[3] = unit(0) + vecs[1] = unit(0) + 0.1 * unit(1) # the best match for window 3 inside its own clip + clips = [(vecs, np.arange(6) * 1.0 + JITTER)] + random_clips(1, 2, seed=21) + index = build(clips) + assert index.window_start_time[3] - index.window_start_time[1] < 2.0 + space = search.prepare(index, centre=centre) + own = [h.window for h in search.nearest_moments(space, 0, 3, k=20, same_clip=True).hits if h.clip == 0] + assert 1 in own and 2 not in own and 3 not in own and 4 not in own + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +def test_suppression_ignores_start_time_jitter(centre): + """Windows 0 and 2 of a clip do not overlap even though jitter puts them 1.994 s apart.""" + q = unit(0) + near = unit(0) + 0.05 * unit(1) + far = unit(2) + clip1 = np.array([near, far, near + 0.01 * unit(3)]) + clips = [([q], [0.0]), (clip1, np.array([0.003, 1.0, 1.997])), ([unit(1) * -1.0], [0.0])] + index = build(clips) + assert index.window_start_time[2] - index.window_start_time[0] < 2.0 + space = search.prepare(index, centre=centre) + res = search.nearest_moments(space, 0, 0, k=2) + assert sorted((h.clip, h.window) for h in res.hits) == [(1, 0), (1, 2)] + + +def _tie_index(seed: int, dim: int = 1024, n: int = 37): + """37 single-window clips of width 1024 with exact copies of clip 0 in clips 7, 8 and 36.""" + rng = np.random.default_rng(seed) + vecs = rng.normal(size=(n, dim)).astype(np.float32) + for c in (7, 8, n - 1): + vecs[c] = vecs[0] + return build([(vecs[i : i + 1], [0.0]) for i in range(n)], dim=dim) + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +@pytest.mark.parametrize("seed", range(40)) +def test_exact_duplicates_tie_in_order_at_full_width(seed, centre): + """Duplicate rows at matrix positions 7, 8 and the last must score equal and rank by clip.""" + space = search.prepare(_tie_index(seed), centre=centre) + res = search.nearest_moments(space, 0, 0, k=3) + assert [(h.clip, h.window) for h in res.hits] == [(7, 0), (8, 0), (36, 0)] + assert len({h.score for h in res.hits}) == 1 + + +@pytest.mark.parametrize("centre", [False, True], ids=["plain", "centred"]) +@pytest.mark.parametrize("seed", range(40)) +def test_exact_duplicate_clips_tie_in_order_at_full_width(seed, centre): + space = search.prepare(_tie_index(seed), centre=centre) + res = search.nearest_clips(space, 0, k=3) + assert [h.clip for h in res.hits] == [7, 8, 36] + assert len({h.score for h in res.hits}) == 1 and res.hits[0].score <= 1.0 + + +def _many_duplicates(seed: int, copies: int, dim: int = 1024, n: int = 37): + rng = np.random.default_rng(seed) + vecs = rng.normal(size=(n, dim)).astype(np.float32) + vecs[1 : copies + 1] = vecs[0] + return build([(vecs[i : i + 1], [0.0]) for i in range(n)], dim=dim) + + +@pytest.mark.parametrize("seed", range(40)) +def test_cutoff_keeps_one_ulp_twins_of_the_best_score_plain(seed): + """With twenty exact copies of the query and k=1, float32 scores differ by an ulp by position.""" + space = search.prepare(_many_duplicates(seed, 20), centre=False) + assert [h.clip for h in search.nearest_moments(space, 0, 0, k=1).hits] == [1] + + +@pytest.mark.parametrize("seed", range(40)) +def test_many_duplicate_clips_tie_in_clip_order_plain(seed): + space = search.prepare(_many_duplicates(seed, 12), centre=False) + assert [h.clip for h in search.nearest_clips(space, 0, k=3).hits] == [1, 2, 3] + + +def test_scores_are_clipped_to_unit_range(): + space = search.prepare(build(random_clips(3, 1, seed=22)), centre=False) + for h in search.nearest_moments(space, 0, 0, k=3).hits: + assert -1.0 <= h.score <= 1.0 + + +def test_identical_rows_with_float32_accumulated_mean_are_unavailable_centred(): + """A stored mean off by a few float32 ulps must not make identical rows look distinct.""" + v = np.random.default_rng(23).normal(size=(1, DIM)).astype(np.float32) * 40 + index = build([(v, [0.0]) for _ in range(10)]) + sloppy = (v[0] + 3e-5).astype(np.float32) # what a float32 accumulation might leave behind + index = dataclasses.replace(index, mean=sloppy) + index_file.check( + {"pooled": index.pooled, "window_start_time": index.window_start_time, "clip": index.clip, "mean": index.mean}, + index.header, + ) + space = search.prepare(index, centre=True) + assert not space.available.any() and not space.clip_available.any() + assert search.nearest_moments(space, 0, 0, k=2).reason + + +def test_min_norm_threshold_separates_tiny_residues(): + """Nine identical rows and one that differs by 4e-6: after centring the nine have norm 4e-7.""" + base = np.random.default_rng(24).normal(size=DIM).astype(np.float32) + odd = base.copy() + odd[0] += np.float32(4e-6) + vecs = np.array([base] * 9 + [odd]) + space = search.prepare(build([(vecs[i : i + 1], [0.0]) for i in range(10)]), centre=True) + assert space.available.tolist() == [False] * 9 + [True] + + +def test_unavailable_reason_names_centring_only_when_centred(): + x = unit(0) + index = build([([x], [0.0]), ([-x], [0.0]), ([0 * x], [0.0])]) + centred = search.nearest_moments(search.prepare(index, centre=True), 2, 0, k=1).reason + plain = search.nearest_moments(search.prepare(index, centre=False), 2, 0, k=1).reason + assert "index mean" in centred and "index mean" not in plain + + +def test_default_similarity_is_plain_cosine(): + assert search.DEFAULT_SIMILARITY == "cosine" + + +def test_a_suppressed_hit_never_suppresses_another_plain(): + """Windows 0, 1 and 2 of one clip rank in that order; 1 overlaps both others and is dropped, + but 2 does not overlap 0, so it stays.""" + q = unit(0) + w0, w1, w2 = (unit(0) + d * unit(1) for d in (0.1, 0.2, 0.3)) + space = search.prepare(build([([q], [0.0]), ([w0, w1, w2], [0.0, 1.0, 2.0])]), centre=False) + assert [(h.clip, h.window) for h in search.nearest_moments(space, 0, 0, k=5).hits] == [(1, 0), (1, 2)] + + +def test_scores_above_one_are_clipped_plain(): + """A stored row a little longer than unit length would score above 1 without the clip.""" + space = search.prepare(build([([unit(0)], [0.0]), ([unit(0)], [0.0])]), centre=False) + matrix = space.matrix.copy() + matrix[1] *= np.float32(1 + 1e-5) + long_row = dataclasses.replace(space, matrix=matrix) + assert float(matrix[1].astype(np.float64) @ matrix[0].astype(np.float64)) > 1.0 + (hit,) = search.nearest_moments(long_row, 0, 0, k=1).hits + assert hit.score == 1.0 + + +def test_k_is_capped(): + space = search.prepare(build(random_clips(3, 2, seed=3)), centre=False) + assert len(search.nearest_moments(space, 0, 0, k=search.MAX_K).hits) == 2 # one per other clip + with pytest.raises(ValueError, match="at most"): + search.nearest_moments(space, 0, 0, k=search.MAX_K + 1) + with pytest.raises(ValueError, match="at most"): + search.nearest_clips(space, 0, k=search.MAX_K + 1) From 35bd26171880248b54bfb2231a1761bbdc085357 Mon Sep 17 00:00:00 2001 From: RDCLder Date: Sun, 8 Nov 2026 16:41:31 -0600 Subject: [PATCH 3/6] Add the folder indexer and the index and search commands jepa-scope index embeds every video in a folder through the result cache and writes jepa-scope.index beside the videos. A rerun keeps files whose size and modification time are unchanged without reading them, hashes the rest, and embeds only what the cache does not hold; byte-identical files share one cache file. --verify hashes everything. An index made with other settings, from another folder, or in an unreadable format is rebuilt with the reason printed. A bad file is skipped and recorded with its reason instead of stopping the run, and is not retried until it changes. The model loads only when a clip has to be embedded. jepa-scope search lists the most similar moments or clips. A video that is not in the index is embedded with the index's own model and settings and refused when this machine would produce a different representation. --source-root locates a folder that moved away from an index written elsewhere with --out. The cache rule moves from the CLI into cache.py with the same file names, so existing cache files keep hitting; a test pins the names. Signed-off-by: RDCLder --- README.md | 50 +++ src/jepa_scope/cache.py | 171 ++++++++++ src/jepa_scope/cli.py | 249 +++++++++++---- src/jepa_scope/indexing.py | 454 ++++++++++++++++++++++++++ tests/retrieval_eval.py | 13 +- tests/test_indexing.py | 639 +++++++++++++++++++++++++++++++++++++ tests/test_pipeline.py | 63 ++-- 7 files changed, 1550 insertions(+), 89 deletions(-) create mode 100644 src/jepa_scope/cache.py create mode 100644 src/jepa_scope/indexing.py create mode 100644 tests/test_indexing.py diff --git a/README.md b/README.md index 4fb836e..af08793 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,56 @@ proportionally more memory and disk. recorded when it was made. A video with fewer frames per second than `--fps` is refused too, because each window would silently cover more time than the header says. +## Indexing and searching a folder + +```bash +jepa-scope index ~/videos # embed every video (cached) and write ~/videos/jepa-scope.index +jepa-scope search ~/videos beach.mp4 --at 12 +jepa-scope search ~/videos beach.mp4 --by clip +jepa-scope search ~/videos ~/elsewhere/query.mp4 +``` + +`jepa-scope index` embeds every video in a folder through the same cache as the other commands and +writes one index file, `jepa-scope.index`, beside the videos. It looks for `.mp4`, `.mov`, `.m4v`, +`.mkv`, `.webm` and `.avi` files in any letter case, ignores names that start with a dot, and only +enters subfolders with `--recursive`, which never follows a symlinked folder. It takes the same +embedding options as `jepa-scope embed`. Unlike `embed`, it leaves the large `patches` tensor out of +the cache files it writes, because search does not use it, and `--patches` keeps it. + +Running it again refreshes the index. A video whose size and modification time match the index, and +whose cache file is still valid, is kept without being read. Any other video is hashed again and +embedded only if the cache holds nothing for its contents, so renaming or touching a file costs one +hash, and a file that was deleted is dropped from the index. A rerun with nothing new to embed does +not load the model. `--verify` hashes every video whatever its size and time say, which catches a +file replaced by another of the same size and time. If the model, the embedding settings, the +device or the dtype change, the index is rebuilt from scratch, and the program says why. + +One bad file does not stop a run. A video that is too short for one window, has fewer frames per +second than `--fps`, has no video stream, cannot be decoded, or makes the model produce invalid +values is reported, skipped and listed in the index with the reason. It is not tried again until it +changes or `--verify` is given. + +`jepa-scope search` takes the folder (or the index file) and a query. A query that is one of the +indexed clips, given by its path inside the folder or by a path to the file, is searched with the +vectors already in the index. Any other video is embedded through the cache with the model and +settings the index was built with, and it is refused if this machine's model files would produce a +different representation. By default the results are the moments in other clips most like the +window at `--at` seconds, with overlapping windows of one clip collapsed to the best one; +`--same-clip` also lists moments from the query's own clip, and `--by clip` compares whole clips. Each +window is two seconds long at the defaults and a new one starts every second. Similarity is plain +cosine between the mean tokens of two windows, and `jepa-scope index --centre` subtracts the +folder's mean vector first. + +Only an index written beside the videos moves with the folder: it records the folder as "the +directory I am in", so the folder can be copied or moved and searched again as is. An index written +elsewhere with `--out` records the folder's absolute path. If that folder moves, pass its new +location to `jepa-scope search` with `--source-root`. + +Each clip is embedded in one piece, so its decoded frames and the tokens of all its windows are in +memory at once. By our estimates, clips of up to about 10 minutes each fit on a laptop with 32 GB +of memory, and up to about 5 minutes with 16 GB. Use `--max-seconds` to index only the start of +longer clips. + ## Development ```bash diff --git a/src/jepa_scope/cache.py b/src/jepa_scope/cache.py new file mode 100644 index 0000000..1178414 --- /dev/null +++ b/src/jepa_scope/cache.py @@ -0,0 +1,171 @@ +"""The result cache: where a video's ``.latent`` file lives for a set of embedding settings, and +whether a file found there can be reused. + +Every command that embeds goes through this rule, so a clip embedded once by ``jepa-scope clip.mp4`` +is reused by ``jepa-scope index`` and ``jepa-scope search`` with the same settings, and the other +way round. The cache file name is a function of the video's SHA-256 and the settings that change +the output; it does not depend on where the video is or what it is called. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from . import OUTPUT_VERSION, latent_file + +if TYPE_CHECKING: + import torch + + from .model import Encoder + +CACHE_DIR = Path.home() / ".cache" / "jepa-scope" + + +def _print(message: str) -> None: + print(message, flush=True) + + +@dataclass(frozen=True) +class EmbedSettings: + """The embedding settings a command was given. + + ``device`` and ``dtype`` are kept as given (None means "the best available"); ``resolved`` + says what they resolve to on this machine. ``batch`` changes speed and memory, not the output, + so it is not part of the cache key. + """ + + model: str + fps: float + window: int + batch: int = 2 + max_seconds: float | None = None + device: str | None = None + dtype: str | None = None + + def __post_init__(self) -> None: + # The cache key writes these with json.dumps, where 8 and 8.0 differ. The CLI always passes + # floats; normalising keeps settings built elsewhere (from an index header) on the same key. + object.__setattr__(self, "fps", float(self.fps)) + object.__setattr__(self, "window", int(self.window)) + if self.max_seconds is not None: + object.__setattr__(self, "max_seconds", float(self.max_seconds)) + + @classmethod + def from_args(cls, args: Any) -> EmbedSettings: + """Settings from parsed command line arguments (``--model``, ``--fps`` and so on).""" + return cls( + model=args.model, + fps=args.fps, + window=args.window, + batch=args.batch, + max_seconds=args.max_seconds, + device=args.device, + dtype=args.dtype, + ) + + def _resolve(self) -> tuple[torch.device, torch.dtype]: + from .model import pick_device, pick_dtype + + device = pick_device(self.device) + return device, pick_dtype(device, self.dtype) + + def run(self) -> dict[str, Any]: + """The run settings an index records (``index_file.RUN_KEYS``), with the device type and + dtype resolved, so "best available" and naming the same device give the same block.""" + device, dtype = self._resolve() + return { + "max_seconds": self.max_seconds, + "device": device.type, + "dtype": str(dtype).removeprefix("torch."), + } + + +def cache_path(video_sha256: str, settings: EmbedSettings, patches: bool, cache_dir: Path | None = None) -> Path: + """Cache file for this video and these settings, with or without the ``patches`` tensor. + + The key covers everything that changes the output: the output version (bumped whenever the + energy, preprocessing or windowing changes), the model, sampling, and the device and dtype + actually used, so a CPU float32 run and an MPS float16 run never share a file. ``cache_dir`` + defaults to ``CACHE_DIR``, read when the function is called. + """ + device, dtype = settings._resolve() + key = { + "output_version": OUTPUT_VERSION, + "model": settings.model, + "fps": settings.fps, + "window": settings.window, + "max_seconds": settings.max_seconds, + "device": device.type, + "dtype": str(dtype), + "patches": patches, + } + digest = hashlib.sha256(json.dumps(key, sort_keys=True).encode()).hexdigest()[:12] + return (CACHE_DIR if cache_dir is None else cache_dir) / f"{video_sha256[:16]}-{digest}.latent" + + +class PretrainedModel: + """Access to a pre-trained encoder that loads the weights only when they are needed. + + ``representation()`` describes what the settings would produce from the model's config, + processor and file names, without loading weights, unless the weight files are not in the + local snapshot yet; then it loads the model to find out. ``encoder()`` loads the model on + first use and keeps it. Both remember their result. + """ + + def __init__(self, settings: EmbedSettings, log: Callable[[str], None] = _print): + self.settings = settings + self.log = log + self._encoder: Encoder | None = None + self._block: dict[str, Any] | None = None + + def representation(self) -> dict[str, Any]: + if self._block is None: + from .model import UNKNOWN_WEIGHTS, Encoder + + s = self.settings + block = Encoder.describe(s.model, None, s.window, s.fps) + if block["weights"] == UNKNOWN_WEIGHTS: + block = self.encoder().representation(s.window, s.fps) + self._block = block + return self._block + + def encoder(self) -> Encoder: + if self._encoder is None: + from .model import Encoder + + s = self.settings + self.log(f"loading {s.model} ...") + self._encoder = Encoder.from_pretrained(s.model, device=s.device, dtype=s.dtype) + return self._encoder + + +def expected_representation(settings: EmbedSettings) -> str: + """Hash of the representation these settings would produce, read without loading weights when + the weight files are already in the local snapshot, and by loading the model otherwise.""" + return PretrainedModel(settings, log=lambda _: None).representation()["hash"] + + +def cache_is_valid(path: Path, expected_representation: str, log: Callable[[str], None] = _print) -> bool: + """Whether a cache file can be reused: readable in the current format, made by this output + version, and holding the representation these settings would produce (same weights, + preprocessing, pooling and windowing). Otherwise it is re-embedded over the same path, and + ``log`` is told why (nothing is logged for a file that does not exist).""" + if not path.exists(): + return False + try: + header = latent_file.read_header(path) + except (latent_file.LatentFileError, OSError, ValueError) as exc: + log(f"re-embedding: the cache file {path.name} cannot be used ({exc})") + return False + if header["run"].get("output_version") != OUTPUT_VERSION: + log(f"re-embedding: the cache file {path.name} is from another output version") + return False + if header["representation"]["hash"] != expected_representation: + log(f"re-embedding: the cache file {path.name} holds a different representation") + return False + return True diff --git a/src/jepa_scope/cli.py b/src/jepa_scope/cli.py index 37c3ec7..541107f 100644 --- a/src/jepa_scope/cli.py +++ b/src/jepa_scope/cli.py @@ -3,25 +3,26 @@ jepa-scope clip.mp4 embed (cached) and open the viewer jepa-scope embed clip.mp4 -o x write a .latent file jepa-scope view x.latent open the viewer for an existing .latent file +jepa-scope index folder embed every video in a folder (cached) and write its index +jepa-scope search folder clip list the moments or clips in the index most like a clip """ from __future__ import annotations import argparse -import hashlib -import json +import math import sys import threading import webbrowser from pathlib import Path -from . import OUTPUT_VERSION, __version__, latent_file +from . import __version__, cache, latent_file +from .cache import EmbedSettings -SUBCOMMANDS = {"embed", "view"} -CACHE_DIR = Path.home() / ".cache" / "jepa-scope" +SUBCOMMANDS = {"embed", "view", "index", "search"} -def _add_embed_args(p: argparse.ArgumentParser) -> None: +def _add_embed_args(p: argparse.ArgumentParser, patches_flag: str = "--no-patches") -> None: from .model import DEFAULT_MODEL p.add_argument("--model", default=DEFAULT_MODEL, help="Hugging Face model id (default: %(default)s)") @@ -31,11 +32,18 @@ def _add_embed_args(p: argparse.ArgumentParser) -> None: p.add_argument("--max-seconds", type=float, default=None, help="only embed the first N seconds") p.add_argument("--device", default=None, help="cuda, mps or cpu (default: best available)") p.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default=None) - p.add_argument( - "--no-patches", - action="store_true", - help="skip the full patch features (about 4 MB per second of video); the viewer does not need them", - ) + if patches_flag == "--no-patches": + p.add_argument( + "--no-patches", + action="store_true", + help="skip the full patch features (about 4 MB per second of video); the viewer does not need them", + ) + else: + p.add_argument( + "--patches", + action="store_true", + help="keep the full patch features of newly embedded clips (about 4 MB per second of video)", + ) def _add_view_args(p: argparse.ArgumentParser) -> None: @@ -80,60 +88,139 @@ def _serve(latent: Path, video: Path | None, port: int, open_browser: bool, vide uvicorn.run(app, host="127.0.0.1", port=port, log_level="warning") -def _cache_path(video_sha256: str, args) -> Path: - """Cache file for this video and these settings. +def _model_provider(settings: EmbedSettings) -> cache.PretrainedModel: + """How the index and search commands reach the model (replaced in tests).""" + return cache.PretrainedModel(settings) - The key covers everything that changes the output: the output version (bumped whenever the - energy, preprocessing or windowing changes), the model, sampling, and the device and dtype actually - used, so a CPU float32 run and an MPS float16 run never share a file. - """ - from .model import pick_device, pick_dtype - - device = pick_device(args.device) - settings = { - "output_version": OUTPUT_VERSION, - "model": args.model, - "fps": args.fps, - "window": args.window, - "max_seconds": args.max_seconds, - "device": device.type, - "dtype": str(pick_dtype(device, args.dtype)), - "patches": not args.no_patches, - } - digest = hashlib.sha256(json.dumps(settings, sort_keys=True).encode()).hexdigest()[:12] - return CACHE_DIR / f"{video_sha256[:16]}-{digest}.latent" - - -def _expected_representation(args) -> str: - """Hash of the representation these settings would produce, read without loading weights when - the weight files are already in the local snapshot, and by loading the model otherwise.""" - from .model import UNKNOWN_WEIGHTS, Encoder - - block = Encoder.describe(args.model, None, args.window, args.fps) - if block["weights"] == UNKNOWN_WEIGHTS: - encoder = Encoder.from_pretrained(args.model, device=args.device, dtype=args.dtype) - block = encoder.representation(args.window, args.fps) - return block["hash"] - - -def _cache_is_valid(path: Path, expected_representation: str) -> bool: - """Whether a cache file can be reused: readable in the current format, made by this output - version, and holding the representation these settings would produce (same weights, - preprocessing, pooling and windowing). Otherwise it is re-embedded over the same path.""" - if not path.exists(): - return False + +def _index(args) -> None: + from .indexing import build_or_refresh + + if not args.folder.is_dir(): + raise SystemExit(f"error: {args.folder} is not a folder") + settings = EmbedSettings.from_args(args) + build_or_refresh( + args.folder, + _model_provider(settings), + settings, + out=args.out, + recursive=args.recursive, + patches=args.patches, + centre=args.centre, + verify=args.verify, + log=lambda m: print(m, flush=True), + ) + + +def _search(args) -> int: + from . import index_file, search + from .indexing import INDEX_NAME, ClipSkipped, load_query + + if args.k < 1 or args.k > search.MAX_K: + raise SystemExit(f"error: -k must be between 1 and {search.MAX_K}") + if args.by == "clip" and (args.at is not None or args.same_clip): + raise SystemExit("error: --at and --same-clip apply to --by moment only") + if args.at is not None and not math.isfinite(args.at): + raise SystemExit("error: --at must be a finite number of seconds") + index_path = args.index / INDEX_NAME if args.index.is_dir() else args.index + if not index_path.exists(): + hint = f"; run `jepa-scope index {args.index}` first" if args.index.is_dir() else "" + raise SystemExit(f"error: {index_path} does not exist{hint}") try: - header = latent_file.read_header(path) - except (latent_file.LatentFileError, OSError, ValueError) as exc: - print(f"re-embedding: the cache file {path.name} cannot be used ({exc})", flush=True) - return False - if header["run"].get("output_version") != OUTPUT_VERSION: - print(f"re-embedding: the cache file {path.name} is from another output version", flush=True) - return False - if header["representation"]["hash"] != expected_representation: - print(f"re-embedding: the cache file {path.name} holds a different representation", flush=True) - return False - return True + index = index_file.load(index_path) + except (index_file.IndexFileError, OSError) as exc: + raise SystemExit(f"error: {exc}") from None + root = index.resolve_root(index_path, args.source_root) + if not root.is_dir(): + if args.source_root is not None: + raise SystemExit(f"error: --source-root {root} is not a folder") + raise SystemExit( + f"error: the folder this index was built from, {root}, is not there; " + "if it moved, pass its new location with --source-root DIR" + ) + clips = index.header["clips"] + if not clips: + print("nothing was indexed, so there is nothing to search") + return 0 + space = search.prepare(index) + names = [c["path"] for c in clips] + at = 0.0 if args.at is None else args.at + + clip = _indexed_clip(args.query, names, root) + if clip is not None: + if args.by == "clip": + print(f"clips like {names[clip]}:") + result = search.nearest_clips(space, clip, args.k) + else: + window = search.window_at(space, clip, at) + start = float(space.start_time[space.clip_start[clip] + window]) + print(f"moments like {names[clip]} at {at:g}s (window {window}, from {start:.2f}s):") + result = search.nearest_moments(space, clip, window, args.k, same_clip=args.same_clip) + else: + if not args.query.is_file(): + raise SystemExit(f"error: {args.query} is neither a clip in the index nor a video file") + if args.same_clip: + raise SystemExit("error: --same-clip applies to a clip in the index only") + rep = index.header["representation"] + run = index.header["run"] + settings = EmbedSettings( + model=rep["model_id"], + fps=rep["window"]["sample_fps"], + window=rep["window"]["frames"], + max_seconds=run["max_seconds"], + device=run["device"], + dtype=run["dtype"], + ) + provider = _model_provider(settings) + got = provider.representation()["hash"] + if got != rep["hash"]: + raise SystemExit( + f"error: {args.query} cannot be embedded the way this index was: the model here gives " + f"representation {got}, the index holds {rep['hash']} (other weights or preprocessing); " + "rebuild the index with `jepa-scope index` or search with the model it was built with" + ) + try: + query = load_query(args.query, provider, settings, log=lambda m: print(m, flush=True)) + except ClipSkipped as exc: + raise SystemExit(f"error: {args.query} cannot be embedded: {exc}") from None + if args.by == "clip": + print(f"clips like {args.query}:") + result = search.search_clips(space, search.query_vector(space, query.pooled), args.k) + else: + window = search.nearest_window(query.window_start_time, space.window_seconds, at) + start = float(query.window_start_time[window]) + print(f"moments like {args.query} at {at:g}s (window {window}, from {start:.2f}s):") + vector = search.query_vector(space, query.pooled[window : window + 1]) + result = search.search_moments(space, vector, args.k) + + if not result.hits: + print(f"no results: {result.reason}") + return 0 + print(f"{'rank':>4} {'score':>7} {'time':>8} clip") + for rank, hit in enumerate(result.hits, 1): + time = f"{hit.start_time:.2f}s" if isinstance(hit, search.Hit) else "-" + print(f"{rank:>4} {hit.score:>7.4f} {time:>8} {names[hit.clip]}") + return 0 + + +def _indexed_clip(query: Path, names: list[str], root: Path) -> int | None: + """The position of ``query`` among the index's clips, or None when it is not one of them. + + ``query`` matches a clip when it is written exactly as the clip's path relative to the root + (``sub/clip.mp4``), or when it is a path to an existing file that lies under the root and whose + path relative to the root is a clip's. + """ + text = query.as_posix() + if text in names: + return names.index(text) + if query.is_file(): + try: + rel = query.resolve().relative_to(root.resolve()).as_posix() + except ValueError: + return None + if rel in names: + return names.index(rel) + return None def main(argv: list[str] | None = None) -> int: @@ -149,9 +236,10 @@ def main(argv: list[str] | None = None) -> int: if not args.video.exists(): p.error(f"{args.video} does not exist") video_sha256 = latent_file.file_sha256(args.video) - out = _cache_path(video_sha256, args) + settings = EmbedSettings.from_args(args) + out = cache.cache_path(video_sha256, settings, patches=not args.no_patches) # The representation is described (a few seconds of imports) only when a file exists to check. - if args.refresh or not out.exists() or not _cache_is_valid(out, _expected_representation(args)): + if args.refresh or not out.exists() or not cache.cache_is_valid(out, cache.expected_representation(settings)): _embed(args.video, out, args) else: print(f"using cached {out}", flush=True) @@ -174,6 +262,33 @@ def main(argv: list[str] | None = None) -> int: pv.add_argument("--video", type=Path, default=None, help="source video, if it moved since embedding") _add_view_args(pv) + pi = sub.add_parser( + "index", + help="embed every video in a folder (cached) and write the folder's index", + description="Embed every video in a folder through the cache and write jepa-scope.index beside " + "them. A rerun only embeds what is new or changed.", + ) + pi.add_argument("folder", type=Path) + pi.add_argument("--out", type=Path, default=None, help="index path (default: /jepa-scope.index)") + pi.add_argument("--recursive", action="store_true", help="include subfolders (symlinked folders are never entered)") + pi.add_argument("--centre", action="store_true", help="search with cosine similarity centred on the folder mean") + pi.add_argument("--verify", action="store_true", help="hash every file, even when its size and time are unchanged") + _add_embed_args(pi, patches_flag="--patches") + + ps = sub.add_parser( + "search", + help="list the moments or clips in an index most like a clip", + description="Search a folder's index with one of its clips, or with any video, which is then " + "embedded the way the index was.", + ) + ps.add_argument("index", type=Path, help="an indexed folder, or an index file") + ps.add_argument("query", type=Path, help="a clip in the index (its path relative to the folder) or a video file") + ps.add_argument("--at", type=float, default=None, help="time in the query clip, in seconds (default: 0)") + ps.add_argument("-k", type=int, default=10, help="number of results (default: %(default)s)") + ps.add_argument("--by", choices=["moment", "clip"], default="moment", help="what to compare (default: %(default)s)") + ps.add_argument("--same-clip", action="store_true", help="also list moments from the query's own clip") + ps.add_argument("--source-root", type=Path, default=None, help="where the indexed folder is now, if it moved") + args = p.parse_args(argv) if args.cmd == "embed": if not args.video.exists(): @@ -181,6 +296,10 @@ def main(argv: list[str] | None = None) -> int: _embed(args.video, args.out or Path(args.video.stem + ".latent"), args) elif args.cmd == "view": _serve(args.latent, args.video, args.port, not args.no_browser) + elif args.cmd == "index": + _index(args) + elif args.cmd == "search": + return _search(args) return 0 diff --git a/src/jepa_scope/indexing.py b/src/jepa_scope/indexing.py new file mode 100644 index 0000000..fd47d2a --- /dev/null +++ b/src/jepa_scope/indexing.py @@ -0,0 +1,454 @@ +"""Build or refresh the index of a folder of videos. + +``build_or_refresh`` walks a folder for videos, finds or makes each one's ``.latent`` file through +the result cache (``cache.py``), and writes one index file (``index_file.py``) holding every clip's +window vectors. + +A refresh is incremental. A clip whose size and modification time match its entry in the existing +index, and whose cache file is still there and valid, is kept without reading the video again. +Any other clip is hashed, and it is embedded only when no valid cache file exists for that hash. +Byte-identical files share one cache file and are embedded at most once per run. An existing index +made with other settings, from another folder, or in a format this version cannot read is rebuilt +from scratch, which re-embeds only what the cache does not already hold. + +One bad file never stops the run: it is reported and listed under ``skipped`` in the index with +the reason. The model is reached through a provider object (``ModelProvider``) so that a refresh +with nothing new to embed never loads the weights. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any, Protocol + +import av +import numpy as np +from safetensors import SafetensorError, safe_open + +from . import index_file, latent_file +from .cache import EmbedSettings, cache_is_valid, cache_path +from .embed import embed_video +from .search import DEFAULT_SIMILARITY + +if TYPE_CHECKING: + from .model import Encoder + +INDEX_NAME = "jepa-scope.index" +VIDEO_EXTENSIONS = (".mp4", ".mov", ".m4v", ".mkv", ".webm", ".avi") # compared case-insensitively + +# What reading or embedding one video can raise when the file itself is the problem. PyAV raises +# av.FFmpegError subclasses for files it cannot open or decode (InvalidDataError is also a +# ValueError, FileNotFoundError also an OSError, but others are neither). decode raises ValueError +# for no video stream, no frames, or a source below the sampling rate; embed_video raises ValueError +# for a clip shorter than one window and RuntimeError for NaN or inf model output (and +# Encoder.preprocess raises RuntimeError when preprocessing gives an unexpected shape); OSError +# covers a file that cannot be read. +CLIP_ERRORS = (av.FFmpegError, ValueError, RuntimeError, OSError) + + +class ModelProvider(Protocol): + """How the indexer reaches the model. + + ``representation()`` is the representation block the settings produce, ideally found without + loading weights; ``encoder()`` returns the loaded encoder. Both should remember their result. + ``cache.PretrainedModel`` is the implementation the command line uses. + """ + + def representation(self) -> dict[str, Any]: ... + + def encoder(self) -> Encoder: ... + + +@dataclass(frozen=True) +class ClipVectors: + """What the index needs from one clip's ``.latent`` file (never its ``patches``).""" + + cache_file: Path + pooled: np.ndarray # float32 [W, D] + window_start_time: np.ndarray # float64 [W] + duration: float + width: int + height: int + + +class ClipSkipped(Exception): + """A clip cannot be indexed; the message is the reason recorded under ``skipped``. + + ``sha256`` is set when the file could be hashed, so that the skip can be remembered. + """ + + sha256: str | None = None + + +@dataclass +class Report: + """What one ``build_or_refresh`` run did.""" + + path: Path # the index file written + clips: int = 0 + windows: int = 0 + skipped: list[dict[str, str]] = field(default_factory=list) + embedded: int = 0 # clips embedded in this run + hashed: int = 0 # files whose SHA-256 was computed in this run + rebuilt: str | None = None # why an existing index was rebuilt from scratch, if it was + + +def find_videos(folder: Path, recursive: bool = False) -> list[tuple[str, Path]]: + """Videos under ``folder`` as ``(relative POSIX path, path)``, sorted by relative path. + + A video is a name ending in one of ``VIDEO_EXTENSIONS`` (any case) that is not a directory. + Any name starting with ``.``, file or directory, is skipped. Subfolders are entered only when + ``recursive`` is True, and symlinked directories never are, so a link cycle cannot occur. + Symlinked files are included. + """ + found: list[tuple[str, Path]] = [] + for dirpath, dirnames, filenames in os.walk(folder, followlinks=False): + dirnames[:] = [d for d in dirnames if not d.startswith(".")] if recursive else [] + for name in filenames: + if name.startswith(".") or not name.lower().endswith(VIDEO_EXTENSIONS): + continue + path = Path(dirpath) / name + found.append((path.relative_to(folder).as_posix(), path)) + # Emptying dirnames stops os.walk from descending. It lists a symlink to a directory under + # dirnames, never filenames, so it is never mistaken for a video, and with followlinks=False + # it is not entered either. + return sorted(found, key=lambda item: item[0]) + + +def read_clip_vectors(path: Path, video_sha256: str, hidden_size: int) -> ClipVectors: + """Read the window vectors and the clip facts from a ``.latent`` file, without its patches. + + Raises ``ClipSkipped`` when the file belongs to another video (its ``source.sha256`` is not + ``video_sha256``; the cache name only carries a 16-character prefix) or its vectors would not + make a valid index row set. + """ + try: + header = latent_file.read_header(path) + with safe_open(str(path), framework="numpy") as f: + pooled = f.get_tensor("pooled") + starts = f.get_tensor("window_start_time") + except (latent_file.LatentFileError, SafetensorError, OSError, ValueError, KeyError) as exc: + raise ClipSkipped(f"the cache file {path.name} cannot be read ({exc})") from exc + source = header["source"] + if source.get("sha256") != video_sha256: + raise ClipSkipped(f"the cache file {path.name} belongs to another video") + if pooled.dtype != np.float32 or pooled.ndim != 2 or pooled.shape[1] != hidden_size or len(pooled) == 0: + raise ClipSkipped(f"the cache file {path.name} has pooled of shape {pooled.shape} ({pooled.dtype})") + if starts.dtype != np.float64 or starts.shape != (len(pooled),): + raise ClipSkipped(f"the cache file {path.name} has window_start_time that does not match pooled") + if not (np.isfinite(pooled).all() and np.isfinite(starts).all()): + raise ClipSkipped(f"the cache file {path.name} holds NaN or infinite values") + if len(starts) > 1 and (np.diff(starts) <= 0).any(): + raise ClipSkipped(f"the cache file {path.name} has window start times that do not increase") + return ClipVectors( + cache_file=path, + pooled=pooled, + window_start_time=starts, + duration=float(source["duration"]), + width=int(source["width"]), + height=int(source["height"]), + ) + + +class ClipCache: + """Finds or makes the ``.latent`` file for a video hash, at most once per hash per run.""" + + def __init__( + self, + model: ModelProvider, + settings: EmbedSettings, + patches: bool, + cache_dir: Path | None, + log: Callable[[str], None], + ): + self.model = model + self.settings = settings + self.patches = patches + self.cache_dir = cache_dir + self.log = log + self.embedded = 0 + self._done: dict[str, ClipVectors | ClipSkipped] = {} + self._missing: set[str] = set() # hashes already looked up with no valid cache file + + def _expected(self) -> dict[str, Any]: + return self.model.representation() + + def find(self, video_sha256: str) -> ClipVectors | None: + """A valid cache file for this hash, with patches first, then without; None if neither.""" + done = self._done.get(video_sha256) + if isinstance(done, ClipVectors): + return done + if done is not None or video_sha256 in self._missing: + return None + expected = self._expected() + for patches in (True, False): + path = cache_path(video_sha256, self.settings, patches, self.cache_dir) + if not cache_is_valid(path, expected["hash"], self.log): + continue + try: + vectors = read_clip_vectors(path, video_sha256, expected["hidden_size"]) + except ClipSkipped as exc: + self.log(f"re-embedding: {exc}") + continue + self._done[video_sha256] = vectors + return vectors + self._missing.add(video_sha256) + return None + + def get(self, video: Path, video_sha256: str) -> ClipVectors: + """The vectors for ``video`` (whose hash is ``video_sha256``), embedding it when no valid + cache file exists. Raises ``ClipSkipped`` when the video cannot be embedded.""" + done = self._done.get(video_sha256) + if isinstance(done, ClipSkipped): + raise done + found = self.find(video_sha256) + if found is not None: + return found + try: + vectors = self._embed(video, video_sha256) + except ClipSkipped as exc: + self._done[video_sha256] = exc + raise + self._done[video_sha256] = vectors + return vectors + + def _embed(self, video: Path, video_sha256: str) -> ClipVectors: + s = self.settings + expected = self._expected() + encoder = self.model.encoder() # outside the try: a model that cannot load stops the run + try: + tensors, header = embed_video( + video, + encoder, + sample_fps=s.fps, + window=s.window, + batch_size=s.batch, + max_seconds=s.max_seconds, + keep_patches=self.patches, + log=lambda _: None, + ) + except av.FFmpegError as exc: + # PyAV's message repeats the absolute path; the index entry already names the file. + raise ClipSkipped(f"the file cannot be decoded ({exc.strerror or type(exc).__name__})") from exc + except CLIP_ERRORS as exc: + raise ClipSkipped(str(exc) or type(exc).__name__) from exc + if header["representation"]["hash"] != expected["hash"]: + # Every clip would disagree in the same way, so this is not one bad file. + raise RuntimeError( + "the loaded encoder's representation does not match the one its description gave " + f"({header['representation']['hash']} against {expected['hash']})" + ) + if header["source"]["sha256"] != video_sha256: + raise ClipSkipped("the file changed while it was being indexed") + out = cache_path(video_sha256, s, self.patches, self.cache_dir) + out.parent.mkdir(parents=True, exist_ok=True) + latent_file.save(out, tensors, header) # a cache that cannot be written stops the run + self.embedded += 1 + return read_clip_vectors(out, video_sha256, expected["hidden_size"]) + + +def _root_for(folder: Path, out: Path) -> str: + """``"."`` when the index sits directly in the folder, otherwise the folder's absolute path.""" + folder = folder.resolve() + return "." if out.resolve().parent == folder else str(folder) + + +def _previous( + out: Path, folder: Path, representation: dict[str, Any], run: dict[str, Any], log: Callable[[str], None] +) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]], str | None]: + """The clip entries and the skipped entries of the existing index at ``out``, each by path, + and the reason the index is rebuilt from scratch instead (then both are empty; the reason is + None when there was no index).""" + if not out.exists(): + return {}, {}, None + try: + header = index_file.load(out).header + except (index_file.IndexFileError, OSError) as exc: + reason = f"the existing index cannot be read ({exc})" + else: + old_root = Path(header["root"]) if header["root"] != "." else out.resolve().parent + old_run = {k: header["run"][k] for k in index_file.RUN_KEYS} + if header["representation"]["hash"] != representation["hash"]: + reason = ( + f"the representation changed ({header['representation']['hash']} in the index, " + f"{representation['hash']} now)" + ) + elif old_run != run: + changed = ", ".join(f"{k} {old_run[k]} -> {run[k]}" for k in index_file.RUN_KEYS if old_run[k] != run[k]) + reason = f"the run settings changed ({changed})" + elif old_root.resolve() != folder.resolve(): + reason = f"the existing index was built from another folder ({old_root})" + else: + return {e["path"]: e for e in header["clips"]}, {e["path"]: e for e in header["skipped"]}, None + log(f"rebuilding the index from scratch: {reason}") + return {}, {}, reason + + +def build_or_refresh( + folder: str | Path, + model: ModelProvider, + settings: EmbedSettings, + out: str | Path | None = None, + recursive: bool = False, + patches: bool = False, + centre: bool = False, + verify: bool = False, + log: Callable[[str], None] = print, + cache_dir: Path | None = None, +) -> Report: + """Index the videos in ``folder`` and write the index to ``out`` (default + ``/jepa-scope.index``); return what was done. + + ``patches`` keeps the ``patches`` tensor in newly embedded cache files. ``centre`` records + ``centred-cosine`` as the index's similarity (otherwise the default, plain cosine); changing + only that rewrites the index without embedding. ``verify`` hashes every file whatever its size + and modification time say. ``cache_dir`` overrides the result cache directory. + """ + folder = Path(folder) + if not folder.is_dir(): + raise NotADirectoryError(f"{folder} is not a folder") + out = Path(out) if out is not None else folder / INDEX_NAME + t0 = time.perf_counter() + representation = model.representation() + run = settings.run() + previous, previous_skipped, rebuilt = _previous(out, folder, representation, run, log) + clips = ClipCache(model, settings, patches, cache_dir, log) + report = Report(path=out, rebuilt=rebuilt) + + videos = find_videos(folder, recursive) + entries: list[dict[str, Any]] = [] + rows: list[ClipVectors] = [] + n_rows = 0 + for i, (rel, path) in enumerate(videos, 1): + prefix = f"[{i}/{len(videos)}] {rel}" + if not index_file.is_valid_clip_path(rel): + reason = "the path cannot be stored portably (it holds a backslash, a NUL or a drive prefix)" + report.skipped.append({"path": rel, "reason": reason}) + log(f"{prefix}: skipped, {reason}") + continue + started = time.perf_counter() + embedded_before = clips.embedded + try: + st = _stat(path) + was = previous_skipped.get(rel) + if not verify and was is not None and _unchanged(was, st) and "sha256" in was: + # A file that failed before and has not changed since fails again; trying would + # load the model for nothing. --verify tries every file again. + report.skipped.append(was) + log(f"{prefix}: skipped as before (unchanged), {was['reason']}") + continue + vectors, sha, how = _clip(path, st, previous.get(rel), clips, verify, report) + except ClipSkipped as exc: + entry = {"path": rel, "reason": str(exc)} + if exc.sha256 is not None: + entry.update(sha256=exc.sha256, size=st.st_size, mtime_ns=st.st_mtime_ns) + report.skipped.append(entry) + log(f"{prefix}: skipped, {exc}") + continue + if clips.embedded > embedded_before: + how = f"embedded in {time.perf_counter() - started:.1f}s" + start, n_rows = n_rows, n_rows + len(vectors.pooled) + rows.append(vectors) + entries.append( + { + "path": rel, + "sha256": sha, + "size": st.st_size, + "mtime_ns": st.st_mtime_ns, + "duration": vectors.duration, + "width": vectors.width, + "height": vectors.height, + "windows": [start, n_rows], + } + ) + log(f"{prefix}: {how}, {len(vectors.pooled)} windows") + report.embedded = clips.embedded + + dim = int(representation["hidden_size"]) + if rows: + pooled = np.concatenate([r.pooled for r in rows]).astype(np.float32) + starts = np.concatenate([r.window_start_time for r in rows]).astype(np.float64) + mean = pooled.mean(axis=0, dtype=np.float64).astype(np.float32) + else: + pooled, starts, mean = np.zeros((0, dim), np.float32), np.zeros(0, np.float64), np.zeros(dim, np.float32) + owner = np.repeat(np.arange(len(rows), dtype=np.int32), [len(r.pooled) for r in rows]) + header = { + "representation": representation, + "run": run, + "embedding": {"kind": "window_vector", "similarity": "centred-cosine" if centre else DEFAULT_SIMILARITY}, + "root": _root_for(folder, out), + "skipped": report.skipped, + "clips": entries, + } + out.parent.mkdir(parents=True, exist_ok=True) + index_file.save(out, {"pooled": pooled, "window_start_time": starts, "clip": owner, "mean": mean}, header) + report.clips, report.windows = len(entries), len(pooled) + elapsed = time.perf_counter() - t0 + if entries: + log( + f"indexed {report.clips} clips ({report.windows} windows) into {out} in {elapsed:.1f}s; " + f"embedded {report.embedded}, skipped {len(report.skipped)}" + ) + else: + what = f"all {len(report.skipped)} videos were skipped" if report.skipped else "no videos were found" + log(f"nothing was indexed ({what}); wrote an empty index to {out}") + return report + + +def _clip( + path: Path, + st: os.stat_result, + prev: dict[str, Any] | None, + clips: ClipCache, + verify: bool, + report: Report, +) -> tuple[ClipVectors, str, str]: + """The vectors, hash and a word on how they were found, for one video.""" + if not verify and prev is not None and _unchanged(prev, st): + found = clips.find(prev["sha256"]) + if found is not None: + return found, prev["sha256"], "unchanged" + sha = _sha256(path) + report.hashed += 1 + try: + vectors = clips.get(path, sha) + except ClipSkipped as exc: + exc.sha256 = sha + raise + return vectors, sha, "reused the cache" + + +def _unchanged(entry: dict[str, Any], st: os.stat_result) -> bool: + """Whether a file's size and modification time are the ones ``entry`` recorded.""" + return entry.get("size") == st.st_size and entry.get("mtime_ns") == st.st_mtime_ns + + +def _stat(path: Path) -> os.stat_result: + try: + return path.stat() + except OSError as exc: + raise ClipSkipped(f"the file cannot be read ({exc})") from exc + + +def _sha256(path: Path) -> str: + try: + return latent_file.file_sha256(path) + except OSError as exc: + raise ClipSkipped(f"the file cannot be read ({exc})") from exc + + +def load_query( + video: Path, + model: ModelProvider, + settings: EmbedSettings, + log: Callable[[str], None] = print, + cache_dir: Path | None = None, +) -> ClipVectors: + """The window vectors of a video that is not in an index, through the cache (embedding it, + without patches, only when no valid cache file exists). Raises ``ClipSkipped`` with the reason + when the video cannot be read or embedded.""" + return ClipCache(model, settings, False, cache_dir, log).get(video, _sha256(video)) diff --git a/tests/retrieval_eval.py b/tests/retrieval_eval.py index cdde3dc..b36f842 100644 --- a/tests/retrieval_eval.py +++ b/tests/retrieval_eval.py @@ -119,8 +119,8 @@ def even(frame: np.ndarray) -> np.ndarray: class CachedEmbedder: - """Embeds through ``cli._cache_path`` and ``cli._cache_is_valid`` at the CLI defaults, so a rerun - embeds nothing. The model loads only when a file is missing from the cache.""" + """Embeds through ``cache.cache_path`` and ``cache.cache_is_valid`` at the CLI defaults, so a + rerun embeds nothing. The model loads only when a file is missing from the cache.""" def __init__(self, log=print): from jepa_scope.model import DEFAULT_MODEL @@ -155,13 +155,14 @@ def shortest_edge(self) -> int: return int(processor.size["shortest_edge"]) def load(self, path: Path) -> tuple[dict[str, np.ndarray], dict]: - from jepa_scope import cli, latent_file + from jepa_scope import cache, latent_file from jepa_scope.embed import embed_video - out = cli._cache_path(latent_file.file_sha256(path), self.args) + settings = cache.EmbedSettings.from_args(self.args) + out = cache.cache_path(latent_file.file_sha256(path), settings, patches=not self.args.no_patches) if self._expected is None: - self._expected = cli._expected_representation(self.args) - if not cli._cache_is_valid(out, self._expected): + self._expected = cache.expected_representation(settings) + if not cache.cache_is_valid(out, self._expected): self.log(f"embedding {path.name}") tensors, header = embed_video( path, diff --git a/tests/test_indexing.py b/tests/test_indexing.py new file mode 100644 index 0000000..bbb0237 --- /dev/null +++ b/tests/test_indexing.py @@ -0,0 +1,639 @@ +"""The folder indexer and the ``index`` and ``search`` commands, on the tiny random encoder. + +Every test indexes a fresh copy of one small folder: a moving square, an exact copy of it under +another name, a constant grey clip, and a half-second clip that is shorter than one window. They +are embedded with ``window=8`` at 8 fps, so a three-second clip gives 5 windows. The cache lives in +the test's own temporary folder. The random encoder supports no ranking claim beyond an exact copy +scoring 1.0; these tests check plumbing and the refresh rules. +""" + +from __future__ import annotations + +import dataclasses +import os +import shutil +import struct +from pathlib import Path + +import numpy as np +import pytest + +from jepa_scope import cache, cli, index_file, indexing, latent_file, search +from jepa_scope.cache import EmbedSettings + +from .conftest import write_video + +SETTINGS = EmbedSettings(model="test/tiny-vjepa2", fps=8.0, window=8, device="cpu", dtype="float32") +SQUARE, COPY, GREY, SHORT = "square.mp4", "square-copy.mp4", "grey.mp4", "short.mp4" +EMBED_ARGS = ["--model", "test/tiny-vjepa2", "--fps", "8", "--window", "8", "--device", "cpu", "--dtype", "float32"] + + +class TinyModel: + """A model provider around the tiny encoder. ``Encoder.describe`` cannot resolve its fake model + id, so its representation comes from the encoder itself. ``loads`` counts ``encoder()`` calls.""" + + def __init__(self, encoder, settings: EmbedSettings = SETTINGS): + self.enc = encoder + self.settings = settings + self.loads = 0 + + def representation(self): + return self.enc.representation(self.settings.window, self.settings.fps) + + def encoder(self): + self.loads += 1 + return self.enc + + +def _free_box(payload: bytes) -> bytes: + """An MP4 ``free`` box: padding any MP4 reader skips.""" + return struct.pack(">I", 8 + len(payload)) + b"free" + payload + + +def _square_frames() -> np.ndarray: + t, h, w = 72, 64, 96 + frames = np.zeros((t, h, w, 3), np.uint8) + for i in range(t): + x = int(i / t * (w - 16)) + frames[i, 24:40, x : x + 16] = 255 + return frames + + +@pytest.fixture(scope="session") +def template(tmp_path_factory) -> Path: + """The four videos, written once per session; tests copy them.""" + folder = tmp_path_factory.mktemp("template") + write_video(folder / SQUARE, _square_frames()) + shutil.copyfile(folder / SQUARE, folder / COPY) + write_video(folder / GREY, np.full((72, 64, 96, 3), 128, np.uint8)) + # A trailing free box, so a test can change the file's bytes without changing its size or + # what it decodes to. + with open(folder / GREY, "ab") as f: + f.write(_free_box(b"\0" * 8)) + write_video(folder / SHORT, _square_frames()[:12]) + return folder + + +@pytest.fixture +def folder(tmp_path, template) -> Path: + out = tmp_path / "videos" + shutil.copytree(template, out) + return out + + +@pytest.fixture +def cache_dir(tmp_path) -> Path: + return tmp_path / "cache" + + +@pytest.fixture +def embeds(monkeypatch) -> list[str]: + """The names of the videos ``embed_video`` is called on, in order.""" + calls: list[str] = [] + real = indexing.embed_video + + def counting(video, *args, **kwargs): + calls.append(Path(video).name) + return real(video, *args, **kwargs) + + monkeypatch.setattr(indexing, "embed_video", counting) + return calls + + +@pytest.fixture +def run(tiny_encoder, cache_dir): + """``run(folder, **kwargs)`` builds or refreshes an index, logging into ``run.lines``.""" + lines: list[str] = [] + + def go(folder, model=None, settings=SETTINGS, **kwargs): + return indexing.build_or_refresh( + folder, + model or TinyModel(tiny_encoder, settings), + settings, + log=lines.append, + cache_dir=cache_dir, + **kwargs, + ) + + go.lines = lines + return go + + +def _entry(index: index_file.Index, path: str) -> dict: + return next(c for c in index.header["clips"] if c["path"] == path) + + +def _clip_id(index: index_file.Index, path: str) -> int: + return [c["path"] for c in index.header["clips"]].index(path) + + +# ---------------------------------------------------------------------------------------------- +# Building + + +def test_index_holds_three_clips_and_skips_the_short_one(folder, run, embeds): + report = run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + assert [c["path"] for c in index.header["clips"]] == [GREY, COPY, SQUARE] # sorted by path + assert all(c["windows"][1] - c["windows"][0] == 5 for c in index.header["clips"]) + assert index.pooled.shape == (15, 96) + [skipped] = index.header["skipped"] + assert skipped["path"] == SHORT and "one window needs 8" in skipped["reason"] + assert index.header["root"] == "." + assert index.header["run"] == {"max_seconds": None, "device": "cpu", "dtype": "float32"} + assert index.header["embedding"] == {"kind": "window_vector", "similarity": "cosine"} + assert sorted(embeds) == sorted([GREY, SQUARE, SHORT]) or sorted(embeds) == sorted([GREY, COPY, SHORT]) + assert report.embedded == 2 and report.clips == 3 and report.windows == 15 + # The copy and the square share one cache file and its vectors. + assert _entry(index, COPY)["sha256"] == _entry(index, SQUARE)["sha256"] + assert np.array_equal(index.pooled[index.clip_rows(1)], index.pooled[index.clip_rows(2)]) + # The stored mean is the window-weighted mean of the rows the indexer wrote. + assert np.allclose(index.mean, index.pooled.mean(axis=0, dtype=np.float64), atol=1e-6) + # One progress line per clip. + assert sum(line.startswith("[") for line in run.lines) == 4 + + +def test_new_clips_are_embedded_without_patches_unless_asked(folder, run, cache_dir): + run(folder) + sha = _entry(index_file.load(folder / indexing.INDEX_NAME), SQUARE)["sha256"] + plain = cache.cache_path(sha, SETTINGS, False, cache_dir) + assert plain.exists() and "patches" not in latent_file.load(plain)[0] + assert not cache.cache_path(sha, SETTINGS, True, cache_dir).exists() + + +def test_patches_keeps_them(folder, run, cache_dir): + run(folder, patches=True) + sha = _entry(index_file.load(folder / indexing.INDEX_NAME), SQUARE)["sha256"] + assert "patches" in latent_file.load(cache.cache_path(sha, SETTINGS, True, cache_dir))[0] + + +def test_a_with_patches_cache_file_is_preferred(folder, run, cache_dir, embeds): + """With both cache files present, the with-patches one is used.""" + run(folder) + sha = _entry(index_file.load(folder / indexing.INDEX_NAME), SQUARE)["sha256"] + tensors, header = latent_file.load(cache.cache_path(sha, SETTINGS, False, cache_dir)) + marked = {**tensors, "pooled": tensors["pooled"] * 2} + latent_file.save(cache.cache_path(sha, SETTINGS, True, cache_dir), marked, header) + (folder / indexing.INDEX_NAME).unlink() + embeds.clear() + run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + assert embeds == [SHORT] # the index is gone, so only the clip with no cache file is tried + assert np.array_equal(index.pooled[index.clip_rows(_clip_id(index, SQUARE))], marked["pooled"]) + + +def test_a_without_patches_cache_file_is_reused_by_a_patches_run(folder, run, embeds): + run(folder) + embeds.clear() + run(folder, patches=True, verify=True) + assert embeds == [SHORT] # only the clip that has no cache file at all is tried again + + +def test_the_indexer_never_loads_whole_latent_files(folder, run, monkeypatch): + """It reads only the window vectors; loading every tensor would read the patches too.""" + run(folder, patches=True) + + def refuse(*args, **kwargs): + raise AssertionError("latent_file.load was called") + + monkeypatch.setattr(latent_file, "load", refuse) + run(folder, verify=True) + + +def test_a_rerun_embeds_nothing_and_never_loads_the_model(folder, run, embeds, tiny_encoder): + run(folder) + embeds.clear() + model = TinyModel(tiny_encoder) + report = run(folder, model=model) + assert embeds == [] and model.loads == 0 and report.hashed == 0 and report.embedded == 0 + index = index_file.load(folder / indexing.INDEX_NAME) + assert len(index.header["clips"]) == 3 + # The short clip is skipped again from its record, without trying (and loading the model). + [skipped] = index.header["skipped"] + assert skipped["path"] == SHORT and "one window needs 8" in skipped["reason"] + assert any("skipped as before (unchanged)" in line for line in run.lines) + + +def test_a_skipped_file_is_tried_again_when_it_changes_or_with_verify(folder, run, embeds): + run(folder) + embeds.clear() + run(folder, verify=True) + assert embeds == [SHORT] + embeds.clear() + write_video(folder / SHORT, _square_frames()[:36]) # now 1.5 s: 12 frames at 8 fps, 2 windows + run(folder) + assert embeds == [SHORT] + index = index_file.load(folder / indexing.INDEX_NAME) + assert index.header["skipped"] == [] and len(index.header["clips"]) == 4 + + +def test_touching_a_file_rehashes_but_does_not_re_embed(folder, run, embeds): + (folder / SHORT).unlink() + run(folder) + embeds.clear() + st = (folder / GREY).stat() + os.utime(folder / GREY, ns=(st.st_atime_ns, st.st_mtime_ns + 5_000_000_000)) + report = run(folder) + assert embeds == [] and report.hashed == 1 + assert _entry(index_file.load(folder / indexing.INDEX_NAME), GREY)["mtime_ns"] == st.st_mtime_ns + 5_000_000_000 + + +def test_rewriting_a_file_re_embeds_only_it_and_deleting_one_drops_it(folder, run, embeds): + (folder / SHORT).unlink() + run(folder) + before = index_file.load(folder / indexing.INDEX_NAME) + embeds.clear() + write_video(folder / GREY, np.full((72, 64, 96, 3), 40, np.uint8)) + run(folder) + assert embeds == [GREY] + after = index_file.load(folder / indexing.INDEX_NAME) + assert _entry(after, GREY)["sha256"] != _entry(before, GREY)["sha256"] + + embeds.clear() + (folder / COPY).unlink() + run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + assert embeds == [] and [c["path"] for c in index.header["clips"]] == [GREY, SQUARE] + assert index.pooled.shape == (10, 96) + + +def test_deleting_a_shared_cache_file_embeds_once_and_restores_both(folder, run, embeds, cache_dir): + (folder / SHORT).unlink() + run(folder) + sha = _entry(index_file.load(folder / indexing.INDEX_NAME), SQUARE)["sha256"] + cache.cache_path(sha, SETTINGS, False, cache_dir).unlink() + embeds.clear() + run(folder) + assert len(embeds) == 1 and embeds[0] in (SQUARE, COPY) + index = index_file.load(folder / indexing.INDEX_NAME) + assert {c["path"] for c in index.header["clips"]} == {GREY, SQUARE, COPY} + + +def test_a_cache_file_of_another_video_is_not_reused(folder, run, embeds, cache_dir): + """The cache name holds only 16 characters of the hash; the header's full hash must match.""" + (folder / SHORT).unlink() + run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + square = cache.cache_path(_entry(index, SQUARE)["sha256"], SETTINGS, False, cache_dir) + grey = cache.cache_path(_entry(index, GREY)["sha256"], SETTINGS, False, cache_dir) + shutil.copyfile(square, grey) + embeds.clear() + run(folder) + assert embeds == [GREY] + assert any("belongs to another video" in line for line in run.lines) + index = index_file.load(folder / indexing.INDEX_NAME) + assert not np.array_equal(index.pooled[index.clip_rows(0)], index.pooled[index.clip_rows(2)]) + + +def test_an_invalid_cache_file_is_re_embedded(folder, run, embeds, cache_dir): + (folder / SHORT).unlink() + run(folder) + grey = cache.cache_path( + _entry(index_file.load(folder / indexing.INDEX_NAME), GREY)["sha256"], SETTINGS, False, cache_dir + ) + grey.write_bytes(grey.read_bytes()[:-100]) + embeds.clear() + run(folder) + assert embeds == [GREY] + assert sum("re-embedding" in line for line in run.lines) == 1 + + +def _other_weights(tiny_encoder): + from jepa_scope.model import Encoder, ModelSource + + source = ModelSource("test/tiny-vjepa2", weights="random-seed-0-retrained", revision=None) + return Encoder( + tiny_encoder.model, + tiny_encoder.processor, + "test/tiny-vjepa2", + tiny_encoder.device, + tiny_encoder.dtype, + source=source, + ) + + +@pytest.mark.parametrize( + "change, reason", + [ + ("representation", "the representation changed"), + ("max_seconds", "max_seconds None -> 2.5"), + ("device", "device cpu -> meta"), + ("dtype", "dtype float32 -> float16"), + ], +) +def test_a_settings_change_rebuilds_everything(folder, run, embeds, tiny_encoder, change, reason): + (folder / SHORT).unlink() + run(folder) + embeds.clear() + settings, model = SETTINGS, None + if change == "representation": + model = TinyModel(_other_weights(tiny_encoder)) + else: + value = {"max_seconds": 2.5, "device": "meta", "dtype": "float16"}[change] + settings = dataclasses.replace(SETTINGS, **{change: value}) + report = run(folder, model=model, settings=settings) + assert report.rebuilt is not None and reason in report.rebuilt + assert any(line.startswith("rebuilding the index from scratch") and reason in line for line in run.lines) + assert sorted(embeds) in (sorted([GREY, SQUARE]), sorted([GREY, COPY])) + index = index_file.load(folder / indexing.INDEX_NAME) + assert len(index.header["clips"]) == 3 + assert index.header["run"] == settings.run() + if change == "max_seconds": + assert index.pooled.shape[0] == 3 * 4 # 2.5 s at 8 fps keeps 21 frames: 4 windows + + +def test_an_unreadable_or_older_index_is_rebuilt(folder, run, embeds): + (folder / SHORT).unlink() + out = folder / indexing.INDEX_NAME + out.write_bytes(b"not an index") + report = run(folder) + assert report.rebuilt is not None and "cannot be read" in report.rebuilt + assert report.hashed == 3 + + import json + + from safetensors.numpy import save_file + + index = index_file.load(out) + tensors = { + "pooled": index.pooled, + "window_start_time": index.window_start_time, + "clip": index.clip, + "mean": index.mean, + } + save_file(tensors, str(out), metadata={"jepa_scope": json.dumps({**index.header, "format_version": 99})}) + embeds.clear() + report = run(folder) + assert report.rebuilt is not None and "format_version 99" in report.rebuilt + assert embeds == [] and report.hashed == 3 # rebuilt from scratch: hashed again, cache reused + + +def test_an_index_of_another_folder_is_rebuilt(folder, run, tmp_path): + (folder / SHORT).unlink() + out = tmp_path / "shared.index" + run(folder, out=out) + other = tmp_path / "other" + shutil.copytree(folder, other) + report = run(other, out=out) + assert report.rebuilt is not None and "another folder" in report.rebuilt + + +def test_changing_only_centre_rewrites_the_similarity(folder, run, embeds): + (folder / SHORT).unlink() + run(folder) + embeds.clear() + report = run(folder, centre=True) + index = index_file.load(folder / indexing.INDEX_NAME) + assert index.header["embedding"]["similarity"] == "centred-cosine" + assert embeds == [] and report.rebuilt is None and report.hashed == 0 + run(folder) + assert index_file.load(folder / indexing.INDEX_NAME).header["embedding"]["similarity"] == "cosine" + + +def test_verify_finds_a_replacement_that_kept_size_and_mtime(folder, run, embeds): + (folder / SHORT).unlink() + run(folder) + before = _entry(index_file.load(folder / indexing.INDEX_NAME), GREY) + path = folder / GREY + st = path.stat() + data = bytearray(path.read_bytes()) + data[-1] = 1 # inside the trailing free box: same size, same frames, other bytes + path.write_bytes(bytes(data)) + os.utime(path, ns=(st.st_atime_ns, st.st_mtime_ns)) + assert path.stat().st_size == st.st_size + + embeds.clear() + report = run(folder) + assert embeds == [] and report.hashed == 0 + assert _entry(index_file.load(folder / indexing.INDEX_NAME), GREY)["sha256"] == before["sha256"] + + report = run(folder, verify=True) + assert embeds == [GREY] and report.hashed == 3 + assert _entry(index_file.load(folder / indexing.INDEX_NAME), GREY)["sha256"] != before["sha256"] + + +# ---------------------------------------------------------------------------------------------- +# Search over an indexed folder + + +def test_an_exact_copy_scores_one(folder, run): + run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + space = search.prepare(index) + square, copy = _clip_id(index, SQUARE), _clip_id(index, COPY) + [best] = search.nearest_clips(space, square, 1).hits + assert best.clip == copy and abs(best.score - 1.0) < 1e-6 + for w in range(5): + [hit] = search.nearest_moments(space, square, w, 1).hits + assert (hit.clip, hit.window) == (copy, w) and abs(hit.score - 1.0) < 1e-6 + + +# ---------------------------------------------------------------------------------------------- +# Walking the folder and bad files + + +def test_empty_folder_gives_a_valid_empty_index(tmp_path, run): + empty = tmp_path / "empty" + empty.mkdir() + (empty / "notes.txt").write_text("not a video") + report = run(empty) + index = index_file.load(empty / indexing.INDEX_NAME) + assert index.header["clips"] == [] and index.header["skipped"] == [] and index.pooled.shape == (0, 96) + assert report.clips == 0 + assert any(line.startswith("nothing was indexed (no videos were found)") for line in run.lines) + + +def test_all_skipped_folder_gives_a_valid_empty_index(tmp_path, template, run): + only = tmp_path / "only-short" + only.mkdir() + shutil.copyfile(template / SHORT, only / SHORT) + run(only) + index = index_file.load(only / indexing.INDEX_NAME) + assert index.header["clips"] == [] and [s["path"] for s in index.header["skipped"]] == [SHORT] + assert any(line.startswith("nothing was indexed (all 1 videos were skipped)") for line in run.lines) + + +def test_bad_files_are_skipped_and_the_run_continues(folder, run): + (folder / "garbage.mov").write_bytes(b"\0" * 1000) + (folder / "empty.MP4").write_bytes(b"") + write_video(folder / "slow.mp4", _square_frames()[:24], fps=4) # 4 fps, below the 8 fps asked for + (folder / "back\\slash.mp4").write_bytes((folder / GREY).read_bytes()) + (folder / "dangling.mp4").symlink_to(folder / "missing.mp4") + report = run(folder) + reasons = {s["path"]: s["reason"] for s in report.skipped} + assert set(reasons) == {SHORT, "garbage.mov", "empty.MP4", "slow.mp4", "back\\slash.mp4", "dangling.mp4"} + assert "fewer than the requested 8" in reasons["slow.mp4"] + assert reasons["garbage.mov"].startswith("the file cannot be decoded (") + assert str(folder) not in reasons["garbage.mov"] + assert "portably" in reasons["back\\slash.mp4"] + assert "cannot be read" in reasons["dangling.mp4"] + index = index_file.load(folder / indexing.INDEX_NAME) + assert len(index.header["clips"]) == 3 + assert index.header["skipped"] == report.skipped + + +def test_the_walk_rules(tmp_path, template, run): + top = tmp_path / "top" + (top / "sub" / "deeper").mkdir(parents=True) + (top / ".hidden").mkdir() + outside = tmp_path / "outside" + outside.mkdir() + for path in ( + top / "a.MOV", + top / "sub" / "b.mp4", + top / "sub" / "deeper" / "c.webm", + top / ".hidden" / "d.mp4", + top / ".e.mp4", + outside / "f.mp4", + ): + shutil.copyfile(template / GREY, path) + (top / "notes.txt").write_text("x") + (top / "link").symlink_to(outside, target_is_directory=True) + (top / "g.mkv").symlink_to(template / SQUARE) # a symlinked file is a video like any other + + assert [rel for rel, _ in indexing.find_videos(top)] == ["a.MOV", "g.mkv"] + assert [rel for rel, _ in indexing.find_videos(top, recursive=True)] == [ + "a.MOV", + "g.mkv", + "sub/b.mp4", + "sub/deeper/c.webm", + ] + run(top, recursive=True) + paths = [c["path"] for c in index_file.load(top / indexing.INDEX_NAME).header["clips"]] + assert paths == ["a.MOV", "g.mkv", "sub/b.mp4", "sub/deeper/c.webm"] + + +def test_out_outside_the_folder_records_an_absolute_root(folder, run, tmp_path): + out = tmp_path / "elsewhere" / "videos.index" + run(folder, out=out) + assert index_file.load(out).header["root"] == str(folder.resolve()) + assert not (folder / indexing.INDEX_NAME).exists() + + +def test_out_in_a_subfolder_records_an_absolute_root(folder, run): + out = folder / "indexes" / "x.index" + run(folder, out=out) + assert index_file.load(out).header["root"] == str(folder.resolve()) + + +# ---------------------------------------------------------------------------------------------- +# The command line + + +@pytest.fixture +def patched_cli(monkeypatch, tiny_encoder, cache_dir): + """Points the CLI at the tiny encoder and the test's cache; returns the providers it made.""" + made: list[TinyModel] = [] + + def provider(settings): + made.append(TinyModel(tiny_encoder, settings)) + return made[-1] + + monkeypatch.setattr(cli, "_model_provider", provider) + monkeypatch.setattr(cache, "CACHE_DIR", cache_dir) + return made + + +def test_cli_index_then_search_a_clip(folder, patched_cli, capsys, monkeypatch, cache_dir): + assert cli.main(["index", str(folder), *EMBED_ARGS]) == 0 + out = capsys.readouterr().out + assert "[4/4] square.mp4" in out and "indexed 3 clips (15 windows)" in out + assert any(cache_dir.iterdir()) + + assert cli.main(["search", str(folder), SQUARE, "--at", "1.0", "-k", "3"]) == 0 + out = capsys.readouterr().out + lines = out.splitlines() + assert lines[0].startswith("moments like square.mp4 at 1s (window") + first = lines[2].split() + assert first[0] == "1" and abs(float(first[1]) - 1.0) < 1e-4 and first[3] == COPY + assert len(lines) == 2 + 3 + + assert cli.main(["search", str(folder / indexing.INDEX_NAME), SQUARE, "--by", "clip"]) == 0 + lines = capsys.readouterr().out.splitlines() + assert lines[0] == "clips like square.mp4:" and lines[2].split()[3] == COPY and len(lines) == 2 + 2 + + # A path to the file itself resolves to the indexed clip. + monkeypatch.chdir(folder.parent) + assert cli.main(["search", "videos", "videos/square.mp4", "--same-clip", "-k", "1"]) == 0 + assert capsys.readouterr().out.startswith("moments like square.mp4") + + +def test_cli_search_with_an_outside_video(folder, patched_cli, capsys, tmp_path, tiny_encoder): + cli.main(["index", str(folder), *EMBED_ARGS]) + capsys.readouterr() + query = tmp_path / "query.mp4" + shutil.copyfile(folder / SQUARE, query) # same bytes, so the cache already holds it + assert cli.main(["search", str(folder), str(query), "--by", "clip", "-k", "2"]) == 0 + lines = capsys.readouterr().out.splitlines() + assert lines[0] == f"clips like {query}:" + assert {lines[2].split()[3], lines[3].split()[3]} == {SQUARE, COPY} + assert abs(float(lines[2].split()[1]) - 1.0) < 1e-4 + # The query was embedded with the index's own settings. + assert patched_cli[-1].settings == EmbedSettings( + model="test/tiny-vjepa2", fps=8.0, window=8, device="cpu", dtype="float32" + ) + + fresh = tmp_path / "fresh.mp4" + write_video(fresh, np.full((48, 64, 96, 3), 90, np.uint8)) + assert cli.main(["search", str(folder), str(fresh), "--at", "0.5"]) == 0 + assert capsys.readouterr().out.startswith(f"moments like {fresh} at 0.5s (window 0") + + +def test_cli_search_refuses_a_query_from_another_representation( + folder, patched_cli, capsys, monkeypatch, tmp_path, tiny_encoder +): + cli.main(["index", str(folder), *EMBED_ARGS]) + query = tmp_path / "query.mp4" + shutil.copyfile(folder / GREY, query) + monkeypatch.setattr(cli, "_model_provider", lambda settings: TinyModel(_other_weights(tiny_encoder), settings)) + with pytest.raises(SystemExit, match="cannot be embedded the way this index was"): + cli.main(["search", str(folder), str(query)]) + + +def test_cli_search_of_a_moved_external_index_needs_source_root(folder, patched_cli, capsys, tmp_path): + out = tmp_path / "videos.index" + cli.main(["index", str(folder), "--out", str(out), *EMBED_ARGS]) + moved = tmp_path / "moved" + folder.rename(moved) + with pytest.raises(SystemExit, match="--source-root"): + cli.main(["search", str(out), SQUARE]) + capsys.readouterr() + assert cli.main(["search", str(out), SQUARE, "--source-root", str(moved), "-k", "1"]) == 0 + assert capsys.readouterr().out.splitlines()[2].split()[3] == COPY + assert cli.main(["search", str(out), str(moved / "square.mp4"), "--source-root", str(moved), "-k", "1"]) == 0 + assert capsys.readouterr().out.startswith("moments like square.mp4") + + +def test_cli_search_reports_an_unavailable_query(folder, patched_cli, capsys): + """Ten identical clips centre to zero; here two identical clips alone do.""" + (folder / GREY).unlink() + (folder / SHORT).unlink() + cli.main(["index", str(folder), "--centre", *EMBED_ARGS]) + capsys.readouterr() + assert cli.main(["search", str(folder), SQUARE, "--by", "clip"]) == 0 + assert "no results: the query clip is unavailable" in capsys.readouterr().out + + +def test_cli_search_errors(folder, patched_cli, capsys, tmp_path): + with pytest.raises(SystemExit, match="jepa-scope index"): + cli.main(["search", str(folder), SQUARE]) + cli.main(["index", str(folder), *EMBED_ARGS]) + with pytest.raises(SystemExit, match="neither a clip in the index nor a video file"): + cli.main(["search", str(folder), "nope.mp4"]) + with pytest.raises(SystemExit, match="--by moment only"): + cli.main(["search", str(folder), SQUARE, "--by", "clip", "--at", "1"]) + with pytest.raises(SystemExit, match="--at must be a finite"): + cli.main(["search", str(folder), SQUARE, "--at", "nan"]) + with pytest.raises(SystemExit, match="-k must be between"): + cli.main(["search", str(folder), SQUARE, "-k", "0"]) + with pytest.raises(SystemExit, match="cannot be embedded: the video gives"): + cli.main(["search", str(folder), str(folder / SHORT)]) + + +def test_cli_index_of_an_empty_folder(tmp_path, patched_cli, capsys): + empty = tmp_path / "empty" + empty.mkdir() + assert cli.main(["index", str(empty), *EMBED_ARGS]) == 0 + assert "nothing was indexed" in capsys.readouterr().out + assert cli.main(["search", str(empty), "x.mp4"]) == 0 + assert "nothing was indexed" in capsys.readouterr().out diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 0eced4d..5166d2a 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,3 +1,5 @@ +import dataclasses + import numpy as np import pytest import torch @@ -202,29 +204,54 @@ def test_embed_without_patches(tiny_encoder, moving_square_video): def test_cache_key_depends_on_dtype_and_output_version(monkeypatch): - import argparse - - from jepa_scope import cli + from jepa_scope import cache - base = dict(model="m", fps=8.0, window=16, max_seconds=None, device="cpu", dtype=None, no_patches=False) - k1 = cli._cache_path("ab" * 32, argparse.Namespace(**base)) - k2 = cli._cache_path("ab" * 32, argparse.Namespace(**{**base, "dtype": "float16"})) - monkeypatch.setattr(cli, "OUTPUT_VERSION", 999) - k3 = cli._cache_path("ab" * 32, argparse.Namespace(**base)) + base = cache.EmbedSettings(model="m", fps=8.0, window=16, device="cpu") + k1 = cache.cache_path("ab" * 32, base, patches=True) + k2 = cache.cache_path("ab" * 32, dataclasses.replace(base, dtype="float16"), patches=True) + monkeypatch.setattr(cache, "OUTPUT_VERSION", 999) + k3 = cache.cache_path("ab" * 32, base, patches=True) assert len({k1, k2, k3}) == 3 +@pytest.mark.parametrize( + "settings, patches, name", + [ + (dict(model="m", fps=8.0, window=16, device="cpu"), True, "abababababababab-f685d2fd5882.latent"), + ( + dict(model="facebook/vjepa2-vitl-fpc64-256", fps=8.0, window=16, device="mps"), + False, + "abababababababab-fae44bcf9031.latent", + ), + ( + dict(model="m", fps=4.0, window=8, max_seconds=12.5, device="cpu", dtype="bfloat16"), + True, + "abababababababab-5c2e8661c4a5.latent", + ), + # An int fps or window from elsewhere (an index header) must land on the same file. + (dict(model="m", fps=8, window=16, device="cpu", batch=7), True, "abababababababab-f685d2fd5882.latent"), + ], +) +def test_cache_key_is_unchanged(settings, patches, name, tmp_path): + """Cache names recorded before the cache rule moved out of the CLI; existing cache files must + keep hitting. A change here needs an OUTPUT_VERSION decision, not a new expected value.""" + from jepa_scope import cache + + assert cache.cache_path("ab" * 32, cache.EmbedSettings(**settings), patches).name == name + assert cache.cache_path("ab" * 32, cache.EmbedSettings(**settings), patches, tmp_path).parent == tmp_path + + def test_truncated_cache_file_is_treated_as_missing(latent, tmp_path): - from jepa_scope import cli + from jepa_scope import cache broken = tmp_path / "broken.latent" broken.write_bytes(latent.read_bytes()[:-100]) with pytest.raises(latent_file.LatentFileError): latent_file.load(broken) expected = latent_file.read_header(latent)["representation"]["hash"] - assert cli._cache_is_valid(latent, expected) - assert not cli._cache_is_valid(broken, expected) - assert not cli._cache_is_valid(tmp_path / "missing.latent", expected) + assert cache.cache_is_valid(latent, expected) + assert not cache.cache_is_valid(broken, expected) + assert not cache.cache_is_valid(tmp_path / "missing.latent", expected) @pytest.mark.parametrize( @@ -267,7 +294,7 @@ def test_save_writes_the_current_format_version(latent, tmp_path): def test_an_older_format_is_refused_and_re_embedded(latent, tmp_path, capsys): - from jepa_scope import cli + from jepa_scope import cache tensors, header = latent_file.load(latent) old = tmp_path / "v1.latent" @@ -275,18 +302,18 @@ def test_an_older_format_is_refused_and_re_embedded(latent, tmp_path, capsys): for read in (latent_file.read_header, latent_file.load): with pytest.raises(latent_file.LatentFileError, match="re-embed the video"): read(old) - assert not cli._cache_is_valid(old, header["representation"]["hash"]) + assert not cache.cache_is_valid(old, header["representation"]["hash"]) assert "re-embedding" in capsys.readouterr().out def test_cache_is_refused_for_another_output_version_or_representation(latent, capsys, monkeypatch): - from jepa_scope import cli + from jepa_scope import cache expected = latent_file.read_header(latent)["representation"]["hash"] - assert not cli._cache_is_valid(latent, "0" * 16) + assert not cache.cache_is_valid(latent, "0" * 16) assert "different representation" in capsys.readouterr().out - monkeypatch.setattr(cli, "OUTPUT_VERSION", 999) - assert not cli._cache_is_valid(latent, expected) + monkeypatch.setattr(cache, "OUTPUT_VERSION", 999) + assert not cache.cache_is_valid(latent, expected) assert "another output version" in capsys.readouterr().out From f901f3ca127fe7c54c97332a6cf161cfea63bfcb Mon Sep 17 00:00:00 2001 From: RDCLder Date: Sun, 8 Nov 2026 17:00:05 -0600 Subject: [PATCH 4/6] Tighten the indexer's skip, search and cache rules after review Only skips caused by a file's content (too short, too few frames per second, no video stream, not decodable) are remembered across runs, together with the jepa-scope version; read errors, running out of memory and invalid model output are retried on the next run. Two identical failing files are tried once. --centre and --recursive are remembered by the index, and the index header records recursive. --no-centre and --no-recursive turn them off. A rebuild from scratch keeps them. search takes an existing file as a file: it is an indexed clip when its path lies under the folder, or when its contents match one, and only a name that is not a file is looked up as a clip path. A query that needs embedding on a device that is not available here is refused with a message. EmbedSettings.from_index and cache.find_cache_file give one cache lookup for the indexer and search, and a local --model folder is recorded by its absolute path. A symlinked index beside the videos keeps the root ".". .latent files are written through a unique temporary file. Signed-off-by: RDCLder --- README.md | 29 +++-- src/jepa_scope/cache.py | 54 +++++++- src/jepa_scope/cli.py | 115 +++++++++++------ src/jepa_scope/index_file.py | 3 + src/jepa_scope/indexing.py | 228 ++++++++++++++++++++++------------ src/jepa_scope/latent_file.py | 24 +++- src/jepa_scope/model.py | 11 ++ tests/test_index.py | 11 ++ tests/test_indexing.py | 208 +++++++++++++++++++++++++++++-- tests/test_pipeline.py | 31 +++++ tests/test_search.py | 1 + 11 files changed, 567 insertions(+), 148 deletions(-) diff --git a/README.md b/README.md index af08793..bcd7877 100644 --- a/README.md +++ b/README.md @@ -170,16 +170,25 @@ not load the model. `--verify` hashes every video whatever its size and time say file replaced by another of the same size and time. If the model, the embedding settings, the device or the dtype change, the index is rebuilt from scratch, and the program says why. -One bad file does not stop a run. A video that is too short for one window, has fewer frames per -second than `--fps`, has no video stream, cannot be decoded, or makes the model produce invalid -values is reported, skipped and listed in the index with the reason. It is not tried again until it -changes or `--verify` is given. - -`jepa-scope search` takes the folder (or the index file) and a query. A query that is one of the -indexed clips, given by its path inside the folder or by a path to the file, is searched with the -vectors already in the index. Any other video is embedded through the cache with the model and -settings the index was built with, and it is refused if this machine's model files would produce a -different representation. By default the results are the moments in other clips most like the +The index remembers two choices. `--recursive` and `--centre` (see below) stay in effect on later +runs until `--no-recursive` or `--no-centre` turns them off, even when the index is rebuilt. Turning +`--recursive` off drops the clips in subfolders from the index, and turning it on again adds them +back from the cache. + +One bad file does not stop a run. A video that cannot be read, is too short for one window, has +fewer frames per second than `--fps`, has no video stream, cannot be decoded, runs out of memory or +makes the model produce invalid values is reported, skipped and listed in the index with the +reason. When the cause is the file's content (too short, too few frames per second, no video +stream, not decodable), the file is not tried again until it changes, `--verify` is given or +jepa-scope is upgraded. Any other failure is tried again on the next run. + +`jepa-scope search` takes the folder (or the index file) and a query. A query that is a path to +one of the indexed clips is searched with the vectors already in the index, and so is a file with +the same contents as an indexed clip, which the output then says. If the query is not an existing +file, it is looked up as a clip's path inside the folder (`sub/clip.mp4`). Any other video is +embedded through the cache with the model and settings the index was built with. It is refused if +this machine's model files would produce a different representation, or if the device the index +was built on is not available here. By default the results are the moments in other clips most like the window at `--at` seconds, with overlapping windows of one clip collapsed to the best one; `--same-clip` also lists moments from the query's own clip, and `--by clip` compares whole clips. Each window is two seconds long at the defaults and a new one starts every second. Similarity is plain diff --git a/src/jepa_scope/cache.py b/src/jepa_scope/cache.py index 1178414..1dbadbc 100644 --- a/src/jepa_scope/cache.py +++ b/src/jepa_scope/cache.py @@ -57,9 +57,15 @@ def __post_init__(self) -> None: @classmethod def from_args(cls, args: Any) -> EmbedSettings: - """Settings from parsed command line arguments (``--model``, ``--fps`` and so on).""" + """Settings from parsed command line arguments (``--model``, ``--fps`` and so on). + + A ``--model`` that names an existing local folder becomes its absolute path, which is the + ``model_id`` the representation block records, so settings rebuilt from an index header + (``from_index``) give the same cache key. A Hugging Face id is kept as it is. + """ + model = str(Path(args.model).resolve()) if Path(args.model).is_dir() else args.model return cls( - model=args.model, + model=model, fps=args.fps, window=args.window, batch=args.batch, @@ -68,6 +74,21 @@ def from_args(cls, args: Any) -> EmbedSettings: dtype=args.dtype, ) + @classmethod + def from_index(cls, header: dict[str, Any]) -> EmbedSettings: + """The settings an index was built with, from its validated header: the model, window and + sampling rate of its representation block and its run settings. ``batch`` is the default, + since it does not change the output.""" + rep, run = header["representation"], header["run"] + return cls( + model=rep["model_id"], + fps=rep["window"]["sample_fps"], + window=rep["window"]["frames"], + max_seconds=run["max_seconds"], + device=run["device"], + dtype=run["dtype"], + ) + def _resolve(self) -> tuple[torch.device, torch.dtype]: from .model import pick_device, pick_dtype @@ -169,3 +190,32 @@ def cache_is_valid(path: Path, expected_representation: str, log: Callable[[str] log(f"re-embedding: the cache file {path.name} holds a different representation") return False return True + + +def find_cache_file( + video_sha256: str, + settings: EmbedSettings, + expected_hash: str, + cache_dir: Path | None = None, + log: Callable[[str], None] = _print, +) -> Path | None: + """A reusable cache file for this video and these settings, or None. + + The file with the ``patches`` tensor is tried first, then the one without. A file is reusable + when ``cache_is_valid`` accepts it and its header's ``source.sha256`` is the full + ``video_sha256`` (the cache name only carries a 16-character prefix of it). + """ + for patches in (True, False): + path = cache_path(video_sha256, settings, patches, cache_dir) + if not cache_is_valid(path, expected_hash, log): + continue + try: + source = latent_file.read_header(path)["source"] + except (latent_file.LatentFileError, OSError, ValueError) as exc: + log(f"re-embedding: the cache file {path.name} cannot be used ({exc})") + continue + if source.get("sha256") != video_sha256: + log(f"re-embedding: the cache file {path.name} belongs to another video") + continue + return path + return None diff --git a/src/jepa_scope/cli.py b/src/jepa_scope/cli.py index 541107f..a1baa96 100644 --- a/src/jepa_scope/cli.py +++ b/src/jepa_scope/cli.py @@ -11,6 +11,7 @@ import argparse import math +import os import sys import threading import webbrowser @@ -114,7 +115,7 @@ def _index(args) -> None: def _search(args) -> int: from . import index_file, search - from .indexing import INDEX_NAME, ClipSkipped, load_query + from .indexing import INDEX_NAME if args.k < 1 or args.k > search.MAX_K: raise SystemExit(f"error: -k must be between 1 and {search.MAX_K}") @@ -146,8 +147,12 @@ def _search(args) -> int: names = [c["path"] for c in clips] at = 0.0 if args.at is None else args.at - clip = _indexed_clip(args.query, names, root) + clip, query_sha = _indexed_clip(args.query, names, [c["sha256"] for c in clips], root) + if clip is None and query_sha is None: + raise SystemExit(f"error: {args.query} is neither a clip in the index nor a video file") if clip is not None: + if query_sha is not None: + print(f"{args.query} has the same contents as the indexed clip {names[clip]}; searching as that clip") if args.by == "clip": print(f"clips like {names[clip]}:") result = search.nearest_clips(space, clip, args.k) @@ -157,32 +162,10 @@ def _search(args) -> int: print(f"moments like {names[clip]} at {at:g}s (window {window}, from {start:.2f}s):") result = search.nearest_moments(space, clip, window, args.k, same_clip=args.same_clip) else: - if not args.query.is_file(): - raise SystemExit(f"error: {args.query} is neither a clip in the index nor a video file") + assert query_sha is not None if args.same_clip: raise SystemExit("error: --same-clip applies to a clip in the index only") - rep = index.header["representation"] - run = index.header["run"] - settings = EmbedSettings( - model=rep["model_id"], - fps=rep["window"]["sample_fps"], - window=rep["window"]["frames"], - max_seconds=run["max_seconds"], - device=run["device"], - dtype=run["dtype"], - ) - provider = _model_provider(settings) - got = provider.representation()["hash"] - if got != rep["hash"]: - raise SystemExit( - f"error: {args.query} cannot be embedded the way this index was: the model here gives " - f"representation {got}, the index holds {rep['hash']} (other weights or preprocessing); " - "rebuild the index with `jepa-scope index` or search with the model it was built with" - ) - try: - query = load_query(args.query, provider, settings, log=lambda m: print(m, flush=True)) - except ClipSkipped as exc: - raise SystemExit(f"error: {args.query} cannot be embedded: {exc}") from None + query = _outside_query(args.query, query_sha, index.header) if args.by == "clip": print(f"clips like {args.query}:") result = search.search_clips(space, search.query_vector(space, query.pooled), args.k) @@ -203,24 +186,66 @@ def _search(args) -> int: return 0 -def _indexed_clip(query: Path, names: list[str], root: Path) -> int | None: - """The position of ``query`` among the index's clips, or None when it is not one of them. +def _indexed_clip(query: Path, names: list[str], shas: list[str], root: Path) -> tuple[int | None, str | None]: + """Which indexed clip ``query`` names, as ``(clip, sha256)``. - ``query`` matches a clip when it is written exactly as the clip's path relative to the root - (``sub/clip.mp4``), or when it is a path to an existing file that lies under the root and whose - path relative to the root is a clip's. + An existing file is taken as a file, never as a name. It is the clip whose path it is, when + its absolute path (with the folders above it resolved, but not a symlink at the end) lies under + the root; then the result is ``(clip, None)``. Otherwise it is hashed, and the result is + ``(clip, sha256)`` for the first clip with the same contents, or ``(None, sha256)`` for a video + that is not in the index. Only an argument that is not an existing file is matched against the + clips' relative paths, as written (``sub/clip.mp4``); ``(None, None)`` means no match. """ - text = query.as_posix() - if text in names: - return names.index(text) if query.is_file(): + path = Path(os.path.abspath(query)) + path = path.parent.resolve() / path.name try: - rel = query.resolve().relative_to(root.resolve()).as_posix() + rel = path.relative_to(root.resolve()).as_posix() except ValueError: - return None + rel = None if rel in names: - return names.index(rel) - return None + return names.index(rel), None + try: + sha = latent_file.file_sha256(query) + except OSError as exc: + raise SystemExit(f"error: {query} cannot be read ({exc})") from None + return (shas.index(sha) if sha in shas else None), sha + text = query.as_posix() + if text in names: + return names.index(text), None + return None, None + + +def _outside_query(video: Path, video_sha256: str, header: dict): + """The window vectors of a video that is not in the index, embedded (through the cache) with + the index's model and settings. Exits with a message when that cannot be done here.""" + from .indexing import ClipSkipped, find_query, load_query + from .model import device_available + + settings = EmbedSettings.from_index(header) + rep = header["representation"] + if not device_available(settings.device): + raise SystemExit( + f"error: {video} is not in the index, and embedding it the way the index was built needs " + f"the {settings.device} device ({settings.dtype}), which is not available here; " + "a clip in the index can still be the query" + ) + provider = _model_provider(settings) + got = provider.representation()["hash"] + if got != rep["hash"]: + raise SystemExit( + f"error: {video} cannot be embedded the way this index was: the model here gives " + f"representation {got}, the index holds {rep['hash']} (other weights or preprocessing); " + "rebuild the index with `jepa-scope index` or search with the model it was built with" + ) + log = lambda m: print(m, flush=True) # noqa: E731 + found = find_query(video_sha256, provider, settings, log=log) + if found is not None: + return found + try: + return load_query(video, provider, settings, log=log, video_sha256=video_sha256) + except ClipSkipped as exc: + raise SystemExit(f"error: {video} cannot be embedded: {exc}") from None def main(argv: list[str] | None = None) -> int: @@ -270,8 +295,18 @@ def main(argv: list[str] | None = None) -> int: ) pi.add_argument("folder", type=Path) pi.add_argument("--out", type=Path, default=None, help="index path (default: /jepa-scope.index)") - pi.add_argument("--recursive", action="store_true", help="include subfolders (symlinked folders are never entered)") - pi.add_argument("--centre", action="store_true", help="search with cosine similarity centred on the folder mean") + pi.add_argument( + "--recursive", + action=argparse.BooleanOptionalAction, + default=None, + help="include subfolders, never entering symlinked ones (default: as the existing index, else no)", + ) + pi.add_argument( + "--centre", + action=argparse.BooleanOptionalAction, + default=None, + help="search with cosine similarity centred on the folder mean (default: as the existing index, else no)", + ) pi.add_argument("--verify", action="store_true", help="hash every file, even when its size and time are unchanged") _add_embed_args(pi, patches_flag="--patches") diff --git a/src/jepa_scope/index_file.py b/src/jepa_scope/index_file.py index c69df8d..3fb900b 100644 --- a/src/jepa_scope/index_file.py +++ b/src/jepa_scope/index_file.py @@ -35,6 +35,7 @@ "run", "embedding", "root", + "recursive", "skipped", "clips", ) @@ -246,6 +247,8 @@ def check_header(header: dict[str, Any]) -> None: root = header["root"] if not isinstance(root, str) or not (root == "." or Path(root).is_absolute()): raise IndexFileError("root must be '.' or an absolute path") + if not isinstance(header["recursive"], bool): + raise IndexFileError("recursive must be true or false") skipped = header["skipped"] if not isinstance(skipped, list) or not all( isinstance(e, dict) and isinstance(e.get("path"), str) and isinstance(e.get("reason"), str) for e in skipped diff --git a/src/jepa_scope/indexing.py b/src/jepa_scope/indexing.py index fd47d2a..5507718 100644 --- a/src/jepa_scope/indexing.py +++ b/src/jepa_scope/indexing.py @@ -9,11 +9,16 @@ Any other clip is hashed, and it is embedded only when no valid cache file exists for that hash. Byte-identical files share one cache file and are embedded at most once per run. An existing index made with other settings, from another folder, or in a format this version cannot read is rebuilt -from scratch, which re-embeds only what the cache does not already hold. +from scratch, which re-embeds only what the cache does not already hold. The similarity +(``centre``) and whether subfolders are included (``recursive``) are remembered from the existing +index unless the caller gives them. One bad file never stops the run: it is reported and listed under ``skipped`` in the index with -the reason. The model is reached through a provider object (``ModelProvider``) so that a refresh -with nothing new to embed never loads the weights. +the reason. A file whose content is the problem (too short, too few frames per second, no video +stream, not decodable) is remembered and not tried again until it changes; other failures (out of +memory, invalid model output, a read error) are tried again on the next run. The model is reached +through a provider object (``ModelProvider``) so that a refresh with nothing new to embed never +loads the weights. """ from __future__ import annotations @@ -29,8 +34,8 @@ import numpy as np from safetensors import SafetensorError, safe_open -from . import index_file, latent_file -from .cache import EmbedSettings, cache_is_valid, cache_path +from . import __version__, index_file, latent_file +from .cache import EmbedSettings, cache_path, find_cache_file from .embed import embed_video from .search import DEFAULT_SIMILARITY @@ -40,14 +45,18 @@ INDEX_NAME = "jepa-scope.index" VIDEO_EXTENSIONS = (".mp4", ".mov", ".m4v", ".mkv", ".webm", ".avi") # compared case-insensitively -# What reading or embedding one video can raise when the file itself is the problem. PyAV raises -# av.FFmpegError subclasses for files it cannot open or decode (InvalidDataError is also a -# ValueError, FileNotFoundError also an OSError, but others are neither). decode raises ValueError -# for no video stream, no frames, or a source below the sampling rate; embed_video raises ValueError -# for a clip shorter than one window and RuntimeError for NaN or inf model output (and -# Encoder.preprocess raises RuntimeError when preprocessing gives an unexpected shape); OSError -# covers a file that cannot be read. -CLIP_ERRORS = (av.FFmpegError, ValueError, RuntimeError, OSError) +# What embedding one video can raise when that video, or this run on it, is the problem. PyAV +# raises av.FFmpegError subclasses for files it cannot open or decode (InvalidDataError is also a +# ValueError, FileNotFoundError also an OSError, MemoryError also a MemoryError, but some are +# neither). decode raises ValueError for no video stream, no frames, or a source below the sampling +# rate; embed_video raises ValueError for a clip shorter than one window and RuntimeError for NaN or +# inf model output (Encoder.preprocess raises RuntimeError when preprocessing gives an unexpected +# shape, and torch raises RuntimeError subclasses when a device runs out of memory). OSError covers +# a file that cannot be read, MemoryError a clip too long for this machine. +CLIP_ERRORS = (av.FFmpegError, ValueError, RuntimeError, OSError, MemoryError) +# Of those, the ones that depend only on the file's content and so recur until it changes. +CONTENT_ERRORS = (ValueError, av.FFmpegError) +TRANSIENT_ERRORS = (OSError, MemoryError, RuntimeError) # retried next run, even from PyAV class ModelProvider(Protocol): @@ -78,10 +87,15 @@ class ClipVectors: class ClipSkipped(Exception): """A clip cannot be indexed; the message is the reason recorded under ``skipped``. - ``sha256`` is set when the file could be hashed, so that the skip can be remembered. + ``remember`` is True when the cause is the file's content, so the same bytes would fail the + same way again; ``sha256`` is set once the file has been hashed. Only a skip with both is + remembered across runs. """ - sha256: str | None = None + def __init__(self, reason: str, remember: bool = False): + super().__init__(reason) + self.remember = remember + self.sha256: str | None = None @dataclass @@ -91,10 +105,14 @@ class Report: path: Path # the index file written clips: int = 0 windows: int = 0 - skipped: list[dict[str, str]] = field(default_factory=list) + # The header's skipped entries: path and reason, plus sha256, size, mtime_ns and jepa_scope + # (the version) for a skip that is remembered. + skipped: list[dict[str, Any]] = field(default_factory=list) embedded: int = 0 # clips embedded in this run hashed: int = 0 # files whose SHA-256 was computed in this run rebuilt: str | None = None # why an existing index was rebuilt from scratch, if it was + centre: bool = False # the choices this run used, given or remembered + recursive: bool = False def find_videos(folder: Path, recursive: bool = False) -> list[tuple[str, Path]]: @@ -119,23 +137,19 @@ def find_videos(folder: Path, recursive: bool = False) -> list[tuple[str, Path]] return sorted(found, key=lambda item: item[0]) -def read_clip_vectors(path: Path, video_sha256: str, hidden_size: int) -> ClipVectors: +def read_clip_vectors(path: Path, hidden_size: int) -> ClipVectors: """Read the window vectors and the clip facts from a ``.latent`` file, without its patches. - Raises ``ClipSkipped`` when the file belongs to another video (its ``source.sha256`` is not - ``video_sha256``; the cache name only carries a 16-character prefix) or its vectors would not - make a valid index row set. + The caller has already matched the file to its video (``cache.find_cache_file``). Raises + ``ClipSkipped`` when the vectors would not make a valid set of index rows. """ try: - header = latent_file.read_header(path) + source = latent_file.read_header(path)["source"] with safe_open(str(path), framework="numpy") as f: pooled = f.get_tensor("pooled") starts = f.get_tensor("window_start_time") except (latent_file.LatentFileError, SafetensorError, OSError, ValueError, KeyError) as exc: raise ClipSkipped(f"the cache file {path.name} cannot be read ({exc})") from exc - source = header["source"] - if source.get("sha256") != video_sha256: - raise ClipSkipped(f"the cache file {path.name} belongs to another video") if pooled.dtype != np.float32 or pooled.ndim != 2 or pooled.shape[1] != hidden_size or len(pooled) == 0: raise ClipSkipped(f"the cache file {path.name} has pooled of shape {pooled.shape} ({pooled.dtype})") if starts.dtype != np.float64 or starts.shape != (len(pooled),): @@ -174,34 +188,30 @@ def __init__( self._done: dict[str, ClipVectors | ClipSkipped] = {} self._missing: set[str] = set() # hashes already looked up with no valid cache file - def _expected(self) -> dict[str, Any]: - return self.model.representation() - def find(self, video_sha256: str) -> ClipVectors | None: - """A valid cache file for this hash, with patches first, then without; None if neither.""" + """The vectors from a valid cache file for this hash, or None if there is none.""" done = self._done.get(video_sha256) if isinstance(done, ClipVectors): return done if done is not None or video_sha256 in self._missing: return None - expected = self._expected() - for patches in (True, False): - path = cache_path(video_sha256, self.settings, patches, self.cache_dir) - if not cache_is_valid(path, expected["hash"], self.log): - continue + expected = self.model.representation() + path = find_cache_file(video_sha256, self.settings, expected["hash"], self.cache_dir, self.log) + if path is not None: try: - vectors = read_clip_vectors(path, video_sha256, expected["hidden_size"]) + vectors = read_clip_vectors(path, expected["hidden_size"]) except ClipSkipped as exc: self.log(f"re-embedding: {exc}") - continue - self._done[video_sha256] = vectors - return vectors + else: + self._done[video_sha256] = vectors + return vectors self._missing.add(video_sha256) return None def get(self, video: Path, video_sha256: str) -> ClipVectors: """The vectors for ``video`` (whose hash is ``video_sha256``), embedding it when no valid - cache file exists. Raises ``ClipSkipped`` when the video cannot be embedded.""" + cache file exists. Raises ``ClipSkipped`` when the video cannot be embedded; a second file + with the same hash gets the same ``ClipSkipped`` without another attempt.""" done = self._done.get(video_sha256) if isinstance(done, ClipSkipped): raise done @@ -211,6 +221,7 @@ def get(self, video: Path, video_sha256: str) -> ClipVectors: try: vectors = self._embed(video, video_sha256) except ClipSkipped as exc: + exc.sha256 = video_sha256 self._done[video_sha256] = exc raise self._done[video_sha256] = vectors @@ -218,7 +229,7 @@ def get(self, video: Path, video_sha256: str) -> ClipVectors: def _embed(self, video: Path, video_sha256: str) -> ClipVectors: s = self.settings - expected = self._expected() + expected = self.model.representation() encoder = self.model.encoder() # outside the try: a model that cannot load stops the run try: tensors, header = embed_video( @@ -231,11 +242,8 @@ def _embed(self, video: Path, video_sha256: str) -> ClipVectors: keep_patches=self.patches, log=lambda _: None, ) - except av.FFmpegError as exc: - # PyAV's message repeats the absolute path; the index entry already names the file. - raise ClipSkipped(f"the file cannot be decoded ({exc.strerror or type(exc).__name__})") from exc except CLIP_ERRORS as exc: - raise ClipSkipped(str(exc) or type(exc).__name__) from exc + raise _skipped(exc) from exc if header["representation"]["hash"] != expected["hash"]: # Every clip would disagree in the same way, so this is not one bad file. raise RuntimeError( @@ -248,44 +256,78 @@ def _embed(self, video: Path, video_sha256: str) -> ClipVectors: out.parent.mkdir(parents=True, exist_ok=True) latent_file.save(out, tensors, header) # a cache that cannot be written stops the run self.embedded += 1 - return read_clip_vectors(out, video_sha256, expected["hidden_size"]) + return read_clip_vectors(out, expected["hidden_size"]) + + +def _skipped(exc: BaseException) -> ClipSkipped: + """The skip for an error from embedding one clip, remembered only for a content cause.""" + remember = isinstance(exc, CONTENT_ERRORS) and not isinstance(exc, TRANSIENT_ERRORS) + if isinstance(exc, av.FFmpegError): + # PyAV's message repeats the absolute path; the index entry already names the file. + return ClipSkipped(f"the file cannot be decoded ({exc.strerror or type(exc).__name__})", remember) + return ClipSkipped(str(exc) or type(exc).__name__, remember) + + +def _index_dir(out: Path) -> Path: + """The directory an index file at ``out`` sits in: the link's own directory when ``out`` is a + symlink, with the directories above it resolved.""" + return Path(os.path.abspath(out)).parent.resolve() def _root_for(folder: Path, out: Path) -> str: """``"."`` when the index sits directly in the folder, otherwise the folder's absolute path.""" folder = folder.resolve() - return "." if out.resolve().parent == folder else str(folder) + return "." if _index_dir(out) == folder else str(folder) + + +@dataclass +class _Previous: + """What the existing index says: its entries by path (empty when it is rebuilt), why it is + rebuilt (None when it is not, or there was none), and its remembered choices, if readable.""" + + clips: dict[str, dict[str, Any]] = field(default_factory=dict) + skipped: dict[str, dict[str, Any]] = field(default_factory=dict) + rebuilt: str | None = None + centre: bool | None = None + recursive: bool | None = None def _previous( out: Path, folder: Path, representation: dict[str, Any], run: dict[str, Any], log: Callable[[str], None] -) -> tuple[dict[str, dict[str, Any]], dict[str, dict[str, Any]], str | None]: - """The clip entries and the skipped entries of the existing index at ``out``, each by path, - and the reason the index is rebuilt from scratch instead (then both are empty; the reason is - None when there was no index).""" +) -> _Previous: if not out.exists(): - return {}, {}, None + return _Previous() try: header = index_file.load(out).header except (index_file.IndexFileError, OSError) as exc: - reason = f"the existing index cannot be read ({exc})" + prev = _Previous(rebuilt=f"the existing index cannot be read ({exc})") + try: # the tensors may be damaged while the choices in the header are still readable + header = index_file.read_header(out) + except (index_file.IndexFileError, OSError): + header = None else: - old_root = Path(header["root"]) if header["root"] != "." else out.resolve().parent + prev = _Previous() + old_root = Path(header["root"]) if header["root"] != "." else _index_dir(out) old_run = {k: header["run"][k] for k in index_file.RUN_KEYS} if header["representation"]["hash"] != representation["hash"]: - reason = ( + prev.rebuilt = ( f"the representation changed ({header['representation']['hash']} in the index, " f"{representation['hash']} now)" ) elif old_run != run: changed = ", ".join(f"{k} {old_run[k]} -> {run[k]}" for k in index_file.RUN_KEYS if old_run[k] != run[k]) - reason = f"the run settings changed ({changed})" + prev.rebuilt = f"the run settings changed ({changed})" elif old_root.resolve() != folder.resolve(): - reason = f"the existing index was built from another folder ({old_root})" + prev.rebuilt = f"the existing index was built from another folder ({old_root})" else: - return {e["path"]: e for e in header["clips"]}, {e["path"]: e for e in header["skipped"]}, None - log(f"rebuilding the index from scratch: {reason}") - return {}, {}, reason + prev.clips = {e["path"]: e for e in header["clips"]} + prev.skipped = {e["path"]: e for e in header["skipped"]} + if header is not None: + prev.centre = header["embedding"]["similarity"] == "centred-cosine" + prev.recursive = header["recursive"] + if prev.rebuilt is not None: + log(f"rebuilding the index from scratch: {prev.rebuilt}") + return prev def build_or_refresh( @@ -293,9 +335,9 @@ def build_or_refresh( model: ModelProvider, settings: EmbedSettings, out: str | Path | None = None, - recursive: bool = False, + recursive: bool | None = None, patches: bool = False, - centre: bool = False, + centre: bool | None = None, verify: bool = False, log: Callable[[str], None] = print, cache_dir: Path | None = None, @@ -303,10 +345,13 @@ def build_or_refresh( """Index the videos in ``folder`` and write the index to ``out`` (default ``/jepa-scope.index``); return what was done. - ``patches`` keeps the ``patches`` tensor in newly embedded cache files. ``centre`` records - ``centred-cosine`` as the index's similarity (otherwise the default, plain cosine); changing - only that rewrites the index without embedding. ``verify`` hashes every file whatever its size - and modification time say. ``cache_dir`` overrides the result cache directory. + ``recursive`` includes subfolders, and ``centre`` records ``centred-cosine`` as the index's + similarity (otherwise the default, plain cosine). Either one left as None keeps what the + existing index recorded, even when it is rebuilt from scratch, and is False for a new index. + Changing only ``centre`` rewrites the index without embedding. ``patches`` keeps the + ``patches`` tensor in newly embedded cache files. ``verify`` hashes every file whatever its + size and modification time say, and tries remembered skips again. ``cache_dir`` overrides the + result cache directory. """ folder = Path(folder) if not folder.is_dir(): @@ -315,9 +360,13 @@ def build_or_refresh( t0 = time.perf_counter() representation = model.representation() run = settings.run() - previous, previous_skipped, rebuilt = _previous(out, folder, representation, run, log) + previous = _previous(out, folder, representation, run, log) + centre = bool(previous.centre) if centre is None else centre + recursive = bool(previous.recursive) if recursive is None else recursive + similarity = "centred-cosine" if centre else DEFAULT_SIMILARITY + log(f"similarity: {similarity}; subfolders: {'included' if recursive else 'not included'}") clips = ClipCache(model, settings, patches, cache_dir, log) - report = Report(path=out, rebuilt=rebuilt) + report = Report(path=out, rebuilt=previous.rebuilt, centre=centre, recursive=recursive) videos = find_videos(folder, recursive) entries: list[dict[str, Any]] = [] @@ -334,18 +383,18 @@ def build_or_refresh( embedded_before = clips.embedded try: st = _stat(path) - was = previous_skipped.get(rel) - if not verify and was is not None and _unchanged(was, st) and "sha256" in was: - # A file that failed before and has not changed since fails again; trying would - # load the model for nothing. --verify tries every file again. + was = previous.skipped.get(rel) + if not verify and was is not None and _still_fails(was, st): + # A file whose content failed before and has not changed since fails again; trying + # would load the model for nothing. --verify tries every file again. report.skipped.append(was) log(f"{prefix}: skipped as before (unchanged), {was['reason']}") continue - vectors, sha, how = _clip(path, st, previous.get(rel), clips, verify, report) + vectors, sha, how = _clip(path, st, previous.clips.get(rel), clips, verify, report) except ClipSkipped as exc: - entry = {"path": rel, "reason": str(exc)} - if exc.sha256 is not None: - entry.update(sha256=exc.sha256, size=st.st_size, mtime_ns=st.st_mtime_ns) + entry: dict[str, Any] = {"path": rel, "reason": str(exc)} + if exc.remember and exc.sha256 is not None: + entry.update(sha256=exc.sha256, size=st.st_size, mtime_ns=st.st_mtime_ns, jepa_scope=__version__) report.skipped.append(entry) log(f"{prefix}: skipped, {exc}") continue @@ -379,8 +428,9 @@ def build_or_refresh( header = { "representation": representation, "run": run, - "embedding": {"kind": "window_vector", "similarity": "centred-cosine" if centre else DEFAULT_SIMILARITY}, + "embedding": {"kind": "window_vector", "similarity": similarity}, "root": _root_for(folder, out), + "recursive": recursive, "skipped": report.skipped, "clips": entries, } @@ -414,12 +464,7 @@ def _clip( return found, prev["sha256"], "unchanged" sha = _sha256(path) report.hashed += 1 - try: - vectors = clips.get(path, sha) - except ClipSkipped as exc: - exc.sha256 = sha - raise - return vectors, sha, "reused the cache" + return clips.get(path, sha), sha, "reused the cache" def _unchanged(entry: dict[str, Any], st: os.stat_result) -> bool: @@ -427,6 +472,12 @@ def _unchanged(entry: dict[str, Any], st: os.stat_result) -> bool: return entry.get("size") == st.st_size and entry.get("mtime_ns") == st.st_mtime_ns +def _still_fails(skipped: dict[str, Any], st: os.stat_result) -> bool: + """Whether a skipped entry is a remembered one that still applies: made by this jepa-scope + version, for a file whose size and modification time have not changed since.""" + return "sha256" in skipped and skipped.get("jepa_scope") == __version__ and _unchanged(skipped, st) + + def _stat(path: Path) -> os.stat_result: try: return path.stat() @@ -441,14 +492,27 @@ def _sha256(path: Path) -> str: raise ClipSkipped(f"the file cannot be read ({exc})") from exc +def find_query( + video_sha256: str, + model: ModelProvider, + settings: EmbedSettings, + log: Callable[[str], None] = print, + cache_dir: Path | None = None, +) -> ClipVectors | None: + """The window vectors of a video from the cache only, or None when it holds none.""" + return ClipCache(model, settings, False, cache_dir, log).find(video_sha256) + + def load_query( video: Path, model: ModelProvider, settings: EmbedSettings, log: Callable[[str], None] = print, cache_dir: Path | None = None, + video_sha256: str | None = None, ) -> ClipVectors: """The window vectors of a video that is not in an index, through the cache (embedding it, without patches, only when no valid cache file exists). Raises ``ClipSkipped`` with the reason when the video cannot be read or embedded.""" - return ClipCache(model, settings, False, cache_dir, log).get(video, _sha256(video)) + sha = _sha256(video) if video_sha256 is None else video_sha256 + return ClipCache(model, settings, False, cache_dir, log).get(video, sha) diff --git a/src/jepa_scope/latent_file.py b/src/jepa_scope/latent_file.py index 9d97d2a..362c1f8 100644 --- a/src/jepa_scope/latent_file.py +++ b/src/jepa_scope/latent_file.py @@ -10,6 +10,7 @@ import hashlib import json import os +import tempfile from pathlib import Path from typing import Any @@ -105,10 +106,25 @@ def save(path: str | Path, tensors: dict[str, np.ndarray], header: dict[str, Any header = {**header, "format": FORMAT, "format_version": FORMAT_VERSION} _check_header(header) contiguous = {k: np.ascontiguousarray(v) for k, v in tensors.items()} - # Write next to the target and rename, so an interrupted write never leaves a half file behind. - tmp = Path(path).with_name(Path(path).name + ".tmp") - save_file(contiguous, str(tmp), metadata={META_KEY: json.dumps(header, sort_keys=True)}) - os.replace(tmp, path) + # Write to a uniquely named file in the target's directory and rename it, so an interrupted + # write never leaves a half file behind and two processes embedding the same video never share + # a temporary name. The temporary file gets the permissions an ordinary new file would + # (mkstemp creates it private). + fd, tmp = tempfile.mkstemp(dir=Path(path).absolute().parent, prefix=Path(path).name + ".", suffix=".tmp") + os.close(fd) + try: + save_file(contiguous, tmp, metadata={META_KEY: json.dumps(header, sort_keys=True)}) + os.chmod(tmp, 0o666 & ~_umask()) + os.replace(tmp, path) + except BaseException: + Path(tmp).unlink(missing_ok=True) + raise + + +def _umask() -> int: + mask = os.umask(0) + os.umask(mask) + return mask def read_header(path: str | Path) -> dict[str, Any]: diff --git a/src/jepa_scope/model.py b/src/jepa_scope/model.py index c2dc4ec..d0b548d 100644 --- a/src/jepa_scope/model.py +++ b/src/jepa_scope/model.py @@ -37,6 +37,17 @@ def pick_device(requested: str | None = None) -> torch.device: return torch.device("cpu") +def device_available(device: str) -> bool: + """Whether ``device`` (``cpu``, ``mps``, ``cuda`` or ``cuda:1`` and so on) can run here.""" + kind = torch.device(device).type + if kind == "cpu": + return True + if kind == "mps": + return torch.backends.mps.is_available() + is_available = getattr(getattr(torch, kind, None), "is_available", None) + return bool(is_available()) if callable(is_available) else False + + def pick_dtype(device: torch.device, requested: str | None = None) -> torch.dtype: if requested: return {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[requested] diff --git a/tests/test_index.py b/tests/test_index.py index f46c71a..c966a30 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -70,6 +70,7 @@ def make_index(counts=WINDOWS_PER_CLIP, dim: int = DIM, similarity: str = "cosin "run": {"max_seconds": 30.0, "device": "cpu", "dtype": "float32"}, "embedding": {"kind": "window_vector", "similarity": similarity}, "root": ".", + "recursive": True, "skipped": [{"path": "bad.mp4", "reason": "could not decode"}], "clips": clips, } @@ -310,6 +311,14 @@ def test_rejects_missing_header_key(key): index_file.check(tensors, header) +def test_rejects_a_header_without_recursive(): + """Named on its own, since the parametrised test above follows REQUIRED_HEADER_KEYS.""" + tensors, header = make_index() + del header["recursive"] + with pytest.raises(IndexFileError, match="missing"): + index_file.check(tensors, header) + + def test_rejects_representation_with_wrong_hash(): tensors, header = make_index() header["representation"]["pooling"] = "max" @@ -439,6 +448,8 @@ def duplicate_paths(header): pytest.param(set_top("root", ""), "root must be", id="root-empty"), pytest.param(set_top("root", 5), "root must be", id="root-not-a-string"), pytest.param(set_top("format_version", True), "format_version", id="format-version-bool"), + pytest.param(set_top("recursive", 1), "recursive must be", id="recursive-int"), + pytest.param(set_top("recursive", None), "recursive must be", id="recursive-null"), pytest.param(set_top("skipped", "x"), "skipped must be", id="skipped-not-a-list"), pytest.param(set_top("skipped", {}), "skipped must be", id="skipped-empty-object"), pytest.param(set_top("skipped", [1, "x"]), "skipped must be", id="skipped-not-objects"), diff --git a/tests/test_indexing.py b/tests/test_indexing.py index bbb0237..9b7f393 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -386,7 +386,7 @@ def test_changing_only_centre_rewrites_the_similarity(folder, run, embeds): index = index_file.load(folder / indexing.INDEX_NAME) assert index.header["embedding"]["similarity"] == "centred-cosine" assert embeds == [] and report.rebuilt is None and report.hashed == 0 - run(folder) + run(folder, centre=False) assert index_file.load(folder / indexing.INDEX_NAME).header["embedding"]["similarity"] == "cosine" @@ -516,6 +516,138 @@ def test_out_in_a_subfolder_records_an_absolute_root(folder, run): assert index_file.load(out).header["root"] == str(folder.resolve()) +@pytest.mark.parametrize( + "error", + [RuntimeError("MPS backend out of memory"), MemoryError(), OSError(5, "Input/output error")], + ids=["runtime", "memory", "os"], +) +def test_a_transient_failure_is_retried_next_run(folder, run, embeds, monkeypatch, error): + real = indexing.embed_video + + def failing_grey(video, *args, **kwargs): + if Path(video).name == GREY: + raise error + return real(video, *args, **kwargs) + + monkeypatch.setattr(indexing, "embed_video", failing_grey) + report = run(folder) + grey = next(s for s in report.skipped if s["path"] == GREY) + assert set(grey) == {"path", "reason"} # listed, not remembered + monkeypatch.setattr(indexing, "embed_video", real) + embeds.clear() + report = run(folder) + assert GREY in [c["path"] for c in index_file.load(folder / indexing.INDEX_NAME).header["clips"]] + assert report.embedded == 1 + + +def test_a_content_failure_is_remembered_with_the_version(folder, run, embeds, monkeypatch): + report = run(folder) + short = next(s for s in report.skipped if s["path"] == SHORT) + assert {"sha256", "size", "mtime_ns", "jepa_scope"} <= set(short) and short["jepa_scope"] == indexing.__version__ + embeds.clear() + run(folder) + assert embeds == [] + monkeypatch.setattr(indexing, "__version__", "99.0.0") + run(folder) + assert embeds == [SHORT] # another jepa-scope version tries again + + +def test_two_identical_failing_files_are_tried_once(folder, run, embeds): + shutil.copyfile(folder / SHORT, folder / "short-copy.mp4") + report = run(folder) + assert embeds.count(SHORT) + embeds.count("short-copy.mp4") == 1 + reasons = {s["path"]: s["reason"] for s in report.skipped} + assert reasons[SHORT] == reasons["short-copy.mp4"] + + +def test_a_provider_that_misdescribes_the_encoder_stops_the_run(folder, run, tiny_encoder, cache_dir): + class Misdescribed(TinyModel): + def representation(self): + return _other_weights(self.enc).representation(self.settings.window, self.settings.fps) + + with pytest.raises(RuntimeError, match="does not match the one its description gave"): + run(folder, model=Misdescribed(tiny_encoder)) + assert not cache_dir.exists() or not list(cache_dir.glob("*.latent")) + assert not (folder / indexing.INDEX_NAME).exists() + + +def test_a_file_that_changes_while_embedding_is_skipped(folder, run, monkeypatch): + real = indexing.embed_video + + def changed(video, *args, **kwargs): + tensors, header = real(video, *args, **kwargs) + if Path(video).name == GREY: + header = {**header, "source": {**header["source"], "sha256": "f" * 64}} + return tensors, header + + monkeypatch.setattr(indexing, "embed_video", changed) + report = run(folder) + grey = next(s for s in report.skipped if s["path"] == GREY) + assert grey == {"path": GREY, "reason": "the file changed while it was being indexed"} + + +def test_centre_and_recursive_are_remembered(folder, run, embeds, template): + (folder / "sub").mkdir() + shutil.copyfile(template / GREY, folder / "sub" / "deep.mp4") + report = run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + assert index.header["recursive"] is False and "sub/deep.mp4" not in [c["path"] for c in index.header["clips"]] + assert sum(line.startswith("similarity: ") for line in run.lines) == 1 + + run(folder, centre=True, recursive=True) + run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + assert index.header["embedding"]["similarity"] == "centred-cosine" and index.header["recursive"] is True + assert "sub/deep.mp4" in [c["path"] for c in index.header["clips"]] + assert [line for line in run.lines if line.startswith("similarity: ")][-1] == ( + "similarity: centred-cosine; subfolders: included" + ) + + # A rebuild from scratch keeps them too. + report = run(folder, settings=dataclasses.replace(SETTINGS, max_seconds=2.5)) + index = index_file.load(folder / indexing.INDEX_NAME) + assert report.rebuilt is not None and report.centre and report.recursive + assert index.header["embedding"]["similarity"] == "centred-cosine" and index.header["recursive"] is True + + embeds.clear() + run(folder, settings=dataclasses.replace(SETTINGS, max_seconds=2.5), recursive=False) + assert "sub/deep.mp4" not in [c["path"] for c in index_file.load(folder / indexing.INDEX_NAME).header["clips"]] + run(folder, settings=dataclasses.replace(SETTINGS, max_seconds=2.5), recursive=True) + assert "sub/deep.mp4" in [c["path"] for c in index_file.load(folder / indexing.INDEX_NAME).header["clips"]] + assert embeds == [] # the cache holds the subfolder clip, and the short clip is remembered + + +def test_settings_from_an_index_give_the_indexer_s_cache_path(folder, run, cache_dir): + run(folder) + index = index_file.load(folder / indexing.INDEX_NAME) + sha = _entry(index, SQUARE)["sha256"] + from_index = EmbedSettings.from_index(index.header) + assert cache.cache_path(sha, from_index, False, cache_dir) == cache.cache_path(sha, SETTINGS, False, cache_dir) + assert cache.cache_path(sha, from_index, False, cache_dir).exists() + assert cache.find_cache_file(sha, from_index, index.header["representation"]["hash"], cache_dir) is not None + + +def test_from_args_makes_a_local_model_folder_absolute(tmp_path, monkeypatch): + import argparse + + monkeypatch.chdir(tmp_path) + (tmp_path / "my-model").mkdir() + base = dict(fps=8.0, window=16, batch=2, max_seconds=None, device=None, dtype=None) + local = EmbedSettings.from_args(argparse.Namespace(model="my-model", **base)) + assert local.model == str((tmp_path / "my-model").resolve()) + hub = EmbedSettings.from_args(argparse.Namespace(model="facebook/vjepa2-vitl-fpc64-256", **base)) + assert hub.model == "facebook/vjepa2-vitl-fpc64-256" + + +def test_a_symlinked_index_beside_the_videos_records_dot(folder, run, tmp_path): + real = tmp_path / "real.index" + run(folder, out=real) + (folder / indexing.INDEX_NAME).symlink_to(real) + report = run(folder) + assert report.rebuilt is None + assert index_file.load(folder / indexing.INDEX_NAME).header["root"] == "." + + # ---------------------------------------------------------------------------------------------- # The command line @@ -561,22 +693,68 @@ def test_cli_index_then_search_a_clip(folder, patched_cli, capsys, monkeypatch, def test_cli_search_with_an_outside_video(folder, patched_cli, capsys, tmp_path, tiny_encoder): cli.main(["index", str(folder), *EMBED_ARGS]) capsys.readouterr() - query = tmp_path / "query.mp4" - shutil.copyfile(folder / SQUARE, query) # same bytes, so the cache already holds it - assert cli.main(["search", str(folder), str(query), "--by", "clip", "-k", "2"]) == 0 + fresh = tmp_path / "fresh.mp4" + write_video(fresh, np.full((48, 64, 96, 3), 90, np.uint8)) + assert cli.main(["search", str(folder), str(fresh), "--at", "0.5"]) == 0 lines = capsys.readouterr().out.splitlines() - assert lines[0] == f"clips like {query}:" - assert {lines[2].split()[3], lines[3].split()[3]} == {SQUARE, COPY} - assert abs(float(lines[2].split()[1]) - 1.0) < 1e-4 + assert lines[0].startswith(f"moments like {fresh} at 0.5s (window 0") + assert len(lines) == 2 + 9 # windows 0, 2 and 4 of each clip are the non-overlapping ones # The query was embedded with the index's own settings. assert patched_cli[-1].settings == EmbedSettings( model="test/tiny-vjepa2", fps=8.0, window=8, device="cpu", dtype="float32" ) + # A second search finds it in the cache. + assert cli.main(["search", str(folder), str(fresh), "--by", "clip"]) == 0 + assert capsys.readouterr().out.startswith(f"clips like {fresh}:") + +def test_cli_search_with_a_file_holding_an_indexed_clip(folder, patched_cli, capsys, tmp_path): + """A file outside the folder with the same bytes as an indexed clip is searched as that clip.""" + cli.main(["index", str(folder), *EMBED_ARGS]) + capsys.readouterr() + query = tmp_path / "query.mp4" + shutil.copyfile(folder / SQUARE, query) + assert cli.main(["search", str(folder), str(query), "--by", "clip", "-k", "2"]) == 0 + lines = capsys.readouterr().out.splitlines() + assert lines[0] == f"{query} has the same contents as the indexed clip {COPY}; searching as that clip" + assert lines[1] == f"clips like {COPY}:" + assert lines[3].split()[3] == SQUARE and abs(float(lines[3].split()[1]) - 1.0) < 1e-4 + assert len(lines) == 3 + 2 # the clip it stands for is excluded + + +def test_cli_search_prefers_an_existing_file_to_a_clip_name(folder, patched_cli, capsys, tmp_path, monkeypatch): + """A file in the working directory named like an indexed clip is that file, not the clip.""" + cli.main(["index", str(folder), *EMBED_ARGS]) + capsys.readouterr() + elsewhere = tmp_path / "cwd" + elsewhere.mkdir() + shutil.copyfile(folder / GREY, elsewhere / SQUARE) + monkeypatch.chdir(elsewhere) + assert cli.main(["search", str(folder), SQUARE, "--by", "clip"]) == 0 + out = capsys.readouterr().out + assert out.startswith(f"{SQUARE} has the same contents as the indexed clip {GREY}") + + +def test_cli_search_by_a_symlinked_indexed_clip(folder, patched_cli, capsys, template): + """A symlink in the folder is the clip of its own name, not of the file it points to.""" + (folder / "link.mp4").symlink_to(template / SQUARE) + cli.main(["index", str(folder), *EMBED_ARGS]) + capsys.readouterr() + assert cli.main(["search", str(folder), str(folder / "link.mp4"), "-k", "1"]) == 0 + out = capsys.readouterr().out + assert out.startswith("moments like link.mp4 at 0s"), out + + +def test_cli_search_needs_the_index_device_to_embed_a_query(folder, patched_cli, monkeypatch, tmp_path): + import torch + + cli.main(["index", str(folder), *EMBED_ARGS[:-4], "--device", "mps", "--dtype", "float16"]) + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False) fresh = tmp_path / "fresh.mp4" write_video(fresh, np.full((48, 64, 96, 3), 90, np.uint8)) - assert cli.main(["search", str(folder), str(fresh), "--at", "0.5"]) == 0 - assert capsys.readouterr().out.startswith(f"moments like {fresh} at 0.5s (window 0") + with pytest.raises(SystemExit, match=r"needs the mps device \(float16\), which is not available here"): + cli.main(["search", str(folder), str(fresh)]) + assert cli.main(["search", str(folder), SQUARE, "-k", "1"]) == 0 # an indexed clip still works def test_cli_search_refuses_a_query_from_another_representation( @@ -584,7 +762,7 @@ def test_cli_search_refuses_a_query_from_another_representation( ): cli.main(["index", str(folder), *EMBED_ARGS]) query = tmp_path / "query.mp4" - shutil.copyfile(folder / GREY, query) + write_video(query, np.full((48, 64, 96, 3), 90, np.uint8)) monkeypatch.setattr(cli, "_model_provider", lambda settings: TinyModel(_other_weights(tiny_encoder), settings)) with pytest.raises(SystemExit, match="cannot be embedded the way this index was"): cli.main(["search", str(folder), str(query)]) @@ -630,6 +808,16 @@ def test_cli_search_errors(folder, patched_cli, capsys, tmp_path): cli.main(["search", str(folder), str(folder / SHORT)]) +def test_cli_centre_and_recursive_flags_are_remembered(folder, patched_cli, capsys): + cli.main(["index", str(folder), "--centre", "--recursive", *EMBED_ARGS]) + cli.main(["index", str(folder), *EMBED_ARGS]) + header = index_file.read_header(folder / indexing.INDEX_NAME) + assert header["embedding"]["similarity"] == "centred-cosine" and header["recursive"] is True + cli.main(["index", str(folder), "--no-centre", "--no-recursive", *EMBED_ARGS]) + header = index_file.read_header(folder / indexing.INDEX_NAME) + assert header["embedding"]["similarity"] == "cosine" and header["recursive"] is False + + def test_cli_index_of_an_empty_folder(tmp_path, patched_cli, capsys): empty = tmp_path / "empty" empty.mkdir() diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 5166d2a..3945213 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -293,6 +293,37 @@ def test_save_writes_the_current_format_version(latent, tmp_path): assert latent_file.read_header(out)["format_version"] == latent_file.FORMAT_VERSION == 2 +def test_save_uses_a_unique_temporary_file(latent, tmp_path): + """Two processes saving the same cache file must not share a temporary name: a directory in + the way of the old fixed name ("x.latent.tmp") does not stop a save.""" + import os + + tensors, header = latent_file.load(latent) + out = tmp_path / "out" / "x.latent" + out.parent.mkdir() + (out.parent / "x.latent.tmp").mkdir() + latent_file.save(out, tensors, header) + assert latent_file.read_header(out)["source"] == header["source"] + mask = os.umask(0) + os.umask(mask) + assert out.stat().st_mode & 0o777 == 0o666 & ~mask + assert sorted(p.name for p in out.parent.iterdir()) == ["x.latent", "x.latent.tmp"] + + +def test_a_failed_save_leaves_no_temporary_file(latent, tmp_path, monkeypatch): + tensors, header = latent_file.load(latent) + + def boom(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(latent_file, "save_file", boom) + out = tmp_path / "out" + out.mkdir() + with pytest.raises(OSError, match="disk full"): + latent_file.save(out / "x.latent", tensors, header) + assert list(out.iterdir()) == [] + + def test_an_older_format_is_refused_and_re_embedded(latent, tmp_path, capsys): from jepa_scope import cache diff --git a/tests/test_search.py b/tests/test_search.py index 9602bd2..f8eed9e 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -54,6 +54,7 @@ def build(clips, similarity: str = "cosine", dim: int | None = None) -> Index: "run": {"max_seconds": None, "device": "cpu", "dtype": "float32"}, "embedding": {"kind": "window_vector", "similarity": similarity}, "root": ".", + "recursive": False, "skipped": [], "clips": entries, } From d3e17172999cf7b12166dd50bd46d377c90129bf Mon Sep 17 00:00:00 2001 From: RDCLder Date: Sun, 8 Nov 2026 17:06:33 -0600 Subject: [PATCH 5/6] Embed an outside search query only when the cache lacks it, and log T4 The device check for a query video that is not in the index now runs only when the query must be embedded, so a query the cache already holds works on any machine. Two tests cover the remaining skip and sticky-choice rules, and the plan records T4 as complete with its measurements. Signed-off-by: RDCLder --- README.md | 2 +- ...6-weeks3-4-folder-search-trajectory-log.md | 48 +++++++++++++++++++ ...11-06-weeks3-4-folder-search-trajectory.md | 2 +- src/jepa_scope/cli.py | 13 ++--- tests/test_indexing.py | 29 ++++++++++- 5 files changed, 84 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index bcd7877..677cc6d 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ device or the dtype change, the index is rebuilt from scratch, and the program s The index remembers two choices. `--recursive` and `--centre` (see below) stay in effect on later runs until `--no-recursive` or `--no-centre` turns them off, even when the index is rebuilt. Turning `--recursive` off drops the clips in subfolders from the index, and turning it on again adds them -back from the cache. +back from the cache without embedding them again, though each one is hashed again. One bad file does not stop a run. A video that cannot be read, is too short for one window, has fewer frames per second than `--fps`, has no video stream, cannot be decoded, runs out of memory or diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md index 4c61078..a78f550 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md @@ -148,3 +148,51 @@ The plan is `2026-11-06-weeks3-4-folder-search-trajectory.md`. Entries are newes `nearest_window`; T6 maps `k > MAX_K` to a 400. T4 should repeat the window-weighted-mean check against the indexer's output. - **Permanent decision record needed:** Yes, at T10 (search as plain cosine over window means). + +- **2026-11-08 — T4 complete: the indexer and the `index` and `search` commands.** + - **Evidence:** an executor (Opus) wrote `cache.py`, `indexing.py`, the two commands, 35 tests + in `tests/test_indexing.py` and a README section, and break-checked 30 rules. The cache + file names are byte-identical to the old `cli._cache_path` (pinned by a test, and checked by + the reviewer against the pre-T4 code), so no `OUTPUT_VERSION` question arose. The tiny + encoder's exact copy scores within `1e-6` of 1.0 by clip and at every window. A fresh + reviewer (Opus) found no critical defects but required three changes. Skips caused by the + environment (out of memory, `OSError`, any `RuntimeError`) were remembered forever; now only + content causes are remembered, keyed on size, mtime and the jepa-scope version. Two hard + checks (a representation mismatch, a file changed while it was indexed) had no test. And + `--centre` was silently dropped by the next refresh. The fix round also made query + resolution in `search` prefer an existing file and match by hash, refused a query that + needs an unavailable device with a message, added the shared `find_cache_file` and + `EmbedSettings.from_index` for T6, stopped a symlinked index file from forcing rebuilds, and + gave `latent_file.save` a unique temporary file so two processes embedding one video cannot + collide. Round 2 passed. The lead added its optional follow-ups: the device check for an + outside query now runs only when the query must be embedded, so a cached query works on any + machine; a test that PyAV's own `OSError` subclasses are not remembered; a test that the + choices survive a rebuild of an index with damaged tensors; and a README line. Each new rule + was break-checked. Deferred: reading the umask once at import (round 2 noted that `_umask()` + sets it to 0 briefly, which matters only if files are saved from threads; T6 does not, and + the umask tests need the live value). The fast suite passes. + - **Lead decision (mid-task):** `--centre` and `--recursive` are sticky. A refresh keeps the + choices the existing index recorded unless `--centre/--no-centre` or + `--recursive/--no-recursive` is given, and `recursive` is a new required index header key. + The index contract is not merged, so `INDEX_FORMAT_VERSION` stays 1. Without this, T6's + `jepa-scope ` would silently switch a centred, recursive index back to plain cosine + and drop every subfolder clip. + - **Deviations:** the tests are in `tests/test_indexing.py`, not `tests/test_index.py`. The + model parameter of `build_or_refresh` is a `ModelProvider` named `model` (it describes the + representation without weights and loads the encoder lazily) instead of `encoder_factory`. + Files skipped for a content reason are remembered in `skipped` with `sha256`, `size`, + `mtime_ns` and the version, so a folder holding one short clip does not load the model on + every refresh. An index whose recorded root is another folder is rebuilt. + - **Measured (this laptop, MPS, 10 Kinetics-mini val clips of 10 s, 21 MB):** the first build + took 34 s (about 0.32 s of embedding per second of video, against the 0.35 s estimate). A + refresh with nothing new takes 3.2 to 3.8 s with no model load, against the "under two + seconds" estimate; almost all of it is importing transformers and torchvision inside + `Encoder.describe`, which loads the video processor for the preprocessing hash. Memoising + the description is possible later and not needed now. `search` on an indexed clip takes + 0.7 s end to end, and the query itself is far below a millisecond at 86 windows (T3 measured + 1.5 ms at 30,000). + - **Impact:** T5 builds its index with `build_or_refresh`. T6 uses `find_cache_file` and + `EmbedSettings.from_index` to find a clip's `.latent`, reads only `path` and `reason` from + `skipped`, and passes `centre=None, recursive=None` from the default command. + - **Permanent decision record needed:** Yes, at T10 (folded into the index contract entry: the + sticky choices and the remembered skips). diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md index 121e361..69c3901 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md @@ -637,7 +637,7 @@ flowchart LR ### T4 — Week 3. The indexer and the CLI -- **Status:** Pending +- **Status:** Complete - **Dependencies:** T3 - **Parallel group:** Sequential (optional overlap with T6, see Dependency summary) - **Complexity:** STRONG diff --git a/src/jepa_scope/cli.py b/src/jepa_scope/cli.py index a1baa96..eec7d8d 100644 --- a/src/jepa_scope/cli.py +++ b/src/jepa_scope/cli.py @@ -224,12 +224,6 @@ def _outside_query(video: Path, video_sha256: str, header: dict): settings = EmbedSettings.from_index(header) rep = header["representation"] - if not device_available(settings.device): - raise SystemExit( - f"error: {video} is not in the index, and embedding it the way the index was built needs " - f"the {settings.device} device ({settings.dtype}), which is not available here; " - "a clip in the index can still be the query" - ) provider = _model_provider(settings) got = provider.representation()["hash"] if got != rep["hash"]: @@ -242,6 +236,13 @@ def _outside_query(video: Path, video_sha256: str, header: dict): found = find_query(video_sha256, provider, settings, log=log) if found is not None: return found + # Only embedding needs the device; a query the cache already holds works anywhere. + if not device_available(settings.device): + raise SystemExit( + f"error: {video} is not in the index, and embedding it the way the index was built needs " + f"the {settings.device} device ({settings.dtype}), which is not available here; " + "a clip in the index can still be the query" + ) try: return load_query(video, provider, settings, log=log, video_sha256=video_sha256) except ClipSkipped as exc: diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 9b7f393..2cad3f1 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -15,6 +15,7 @@ import struct from pathlib import Path +import av import numpy as np import pytest @@ -518,8 +519,14 @@ def test_out_in_a_subfolder_records_an_absolute_root(folder, run): @pytest.mark.parametrize( "error", - [RuntimeError("MPS backend out of memory"), MemoryError(), OSError(5, "Input/output error")], - ids=["runtime", "memory", "os"], + [ + RuntimeError("MPS backend out of memory"), + MemoryError(), + OSError(5, "Input/output error"), + # PyAV's own OSError subclasses are FFmpegErrors too, and still transient. + av.error.PermissionError(13, "Permission denied"), + ], + ids=["runtime", "memory", "os", "pyav-permission"], ) def test_a_transient_failure_is_retried_next_run(folder, run, embeds, monkeypatch, error): real = indexing.embed_video @@ -617,6 +624,19 @@ def test_centre_and_recursive_are_remembered(folder, run, embeds, template): assert embeds == [] # the cache holds the subfolder clip, and the short clip is remembered +def test_choices_survive_a_rebuild_of_an_index_with_damaged_tensors(folder, run): + out = folder / indexing.INDEX_NAME + run(folder, centre=True) + with open(out, "r+b") as f: # the last bytes are tensor data; all-ones bytes make a NaN or a bad id + f.seek(-4, 2) + f.write(b"\xff" * 4) + with pytest.raises(index_file.IndexFileError): + index_file.load(out) + report = run(folder) + assert report.rebuilt is not None and report.centre + assert index_file.load(out).header["embedding"]["similarity"] == "centred-cosine" + + def test_settings_from_an_index_give_the_indexer_s_cache_path(folder, run, cache_dir): run(folder) index = index_file.load(folder / indexing.INDEX_NAME) @@ -755,6 +775,11 @@ def test_cli_search_needs_the_index_device_to_embed_a_query(folder, patched_cli, with pytest.raises(SystemExit, match=r"needs the mps device \(float16\), which is not available here"): cli.main(["search", str(folder), str(fresh)]) assert cli.main(["search", str(folder), SQUARE, "-k", "1"]) == 0 # an indexed clip still works + # A query the cache already holds needs no device. + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + assert cli.main(["search", str(folder), str(fresh), "-k", "1"]) == 0 + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False) + assert cli.main(["search", str(folder), str(fresh), "-k", "1"]) == 0 def test_cli_search_refuses_a_query_from_another_representation( From ddadee82897ef6556e109aad0788bb4a1713b3bd Mon Sep 17 00:00:00 2001 From: RDCLder Date: Sun, 8 Nov 2026 18:13:52 -0600 Subject: [PATCH 6/6] Check that search through the folder index reproduces the retrieval experiment An opt-in real-model test indexes the 50 frozen Kinetics-mini val clips with the folder indexer and scores them with nearest_clips and nearest_moments. It reproduces the experiment exactly: 42/50 clip-level top-1, a window-level rate of 0.7427, and 16/16 localisation queries. The composite builder is shared with the experiment's own test. Signed-off-by: RDCLder --- ...6-weeks3-4-folder-search-trajectory-log.md | 32 +++ ...11-06-weeks3-4-folder-search-trajectory.md | 2 +- tests/retrieval_eval.py | 10 + tests/test_real_model.py | 200 ++++++++++++++++-- 4 files changed, 228 insertions(+), 16 deletions(-) diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md index a78f550..e09726c 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory-log.md @@ -196,3 +196,35 @@ The plan is `2026-11-06-weeks3-4-folder-search-trajectory.md`. Entries are newes `skipped`, and passes `centre=None, recursive=None` from the default command. - **Permanent decision record needed:** Yes, at T10 (folded into the index contract entry: the sticky choices and the remembered skips). + +- **2026-11-08 — T5 complete: the product code path reproduces T0 exactly.** + - **Evidence:** `test_search_reproduces_experiment` (opt-in) copies the 50 frozen val clips + into a folder with one subfolder per class, indexes it with `build_or_refresh` at the CLI + defaults (recursive), and scores retrieval with `nearest_clips` and `nearest_moments`. One + run, 136 s on this laptop (MPS), of which 124 s was the first index (40 clips embedded, 10 + already cached by T4's real-folder run) and 10 s the composite index (one clip embedded). No + download: the clips come from the cached dataset snapshot (95 MB, fetched in T0). + - **Numbers (val, default (a), plain cosine):** clip-level top-1 42/50 (0.84), T0 42/50, with + identical per-class hits (archery 6, bowling 10, flying kite 7, high jump 10, marching 9). + Window-level top-1 0.7427, identical to T0. Both bars are met (0.60 rate; margin 27 clips + over T0's colour baseline of 15, against 8 needed). Localisation through + `nearest_moments(same_clip=False)`, the viewer's default, passes 16/16, with no top hit on a + transition window. Descriptively, with `same_clip=True` only 7 of the 16 queries pass: 9 top + hits (every bowling and flying-kite query, and archery at 3 s) are a non-overlapping later + window of the query's own clip, which is nearer than the composite's copy of the same + moment. The README (T10) should say that within one clip, "similar moments" mostly finds + other moments of the same scene. + - **Deviations:** the clips are copied from the cached snapshot instead of + `snapshot_download(local_dir=...)`, which would download 95 MB again; the index is still + never written beside the Hugging Face cache. T0 removed composite windows that cross a + segment boundary from the candidates, while `nearest_moments` keeps them (the plan called the + two candidate sets identical, which was not quite true), so a top hit on such a window is + scored as a miss and counted; there were none. The margin bar uses T0's recorded colour + baseline (`retrieval_eval.T0_VAL`) instead of recomputing it. The composite builder moved into + `_composite_sources` and `_composite`, shared with T0's test. + - **Review:** a fresh reviewer (Sonnet) found that only the clip-level result was asserted, so + a broken `nearest_moments` could pass; the test now also holds the window-level rate to + within 0.02 of T0's and localisation to at least 15 of 16, and was rerun from the cache + (passed). It confirmed the composite refactor is behaviour-preserving. + - **Impact:** T10 takes these numbers into the README. PR 2 (T2 to T5) is ready to open. + - **Permanent decision record needed:** No (results; T10 records the search decision). diff --git a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md index 69c3901..bf7150b 100644 --- a/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md +++ b/docs/plans/2026-11-06-weeks3-4-folder-search-trajectory.md @@ -693,7 +693,7 @@ flowchart LR ### T5 — Week 3. The real code path reproduces T0 -- **Status:** Pending +- **Status:** Complete - **Dependencies:** T4 - **Parallel group:** Sequential (optional overlap with T7) - **Complexity:** ROUTINE (the bars and the recipe are T0's) diff --git a/tests/retrieval_eval.py b/tests/retrieval_eval.py index b36f842..80d2bca 100644 --- a/tests/retrieval_eval.py +++ b/tests/retrieval_eval.py @@ -47,6 +47,16 @@ QUERY_WINDOWS = 4 # windows starting at 0, 1, 2 and 3 s lie fully inside a 5 s segment TOLERANCE_SECONDS = 1.0 +# T0's recorded result on the val split, for the default (a). T5 reproduces it through the product +# code; the numbers are in the plan's execution log (2026-11-06-weeks3-4-folder-search-trajectory). +T0_VAL = { + "default": "a", + "clip_top1_count": 42, + "window_top1_rate": 0.7427184466019418, + "colour_baseline": 15, # the stronger of the two colour baselines (whole clip) + "localisation_passed": 16, +} + @dataclass class Clip: diff --git a/tests/test_real_model.py b/tests/test_real_model.py index ad6723d..199f23a 100644 --- a/tests/test_real_model.py +++ b/tests/test_real_model.py @@ -95,6 +95,32 @@ def _write_once(path: Path, frames: np.ndarray, fps: int) -> Path: return path +def _composite_sources(clips): + """T0's composite sources, in order: the first clip of archery, bowling, flying kite and + marching (the four queried segments), then the second archery and bowling clips.""" + from . import retrieval_eval as rv + + by_class = {k: [c for c in clips if c.label == k] for k in rv.CLASSES} + return [by_class[k][0] for k in ("archery", "bowling", "flying_kite", "marching")] + [ + by_class["archery"][1], + by_class["bowling"][1], + ] + + +def _composite(sources, out_dir: Path) -> tuple[Path, list[int]]: + """The 30-second localisation composite: five seconds of each source, written once (so its + cache file is reused), and the frame count of each segment.""" + from . import retrieval_eval as rv + + seconds = rv.SEGMENT_SECONDS + parts = [ + _fit(decode(c.path, sample_fps=rv.COMPOSITE_FPS, max_seconds=seconds).frames[: rv.COMPOSITE_FPS * seconds]) + for c in sources + ] + path = _write_once(out_dir / "composite.mp4", np.concatenate(parts), fps=rv.COMPOSITE_FPS) + return path, [len(p) for p in parts] + + def _rate_line(name: str, count: int, n: int, extra: str = "") -> str: return f" {name:<28} {count:>3}/{n} {count / n:.2f}{extra}" @@ -213,21 +239,11 @@ def cen(c, value=centre): # Temporal localisation and its mutation (selection set only). localised = {} if selecting: - by_class = {k: [e for e in kept if e.clip.label == k] for k in rv.CLASSES} - sources = [by_class[k][0] for k in ("archery", "bowling", "flying_kite", "marching")] - sources += [by_class["archery"][1], by_class["bowling"][1]] - seconds = rv.SEGMENT_SECONDS - parts = [ - _fit( - decode(s.clip.path, sample_fps=rv.COMPOSITE_FPS, max_seconds=seconds).frames[ - : rv.COMPOSITE_FPS * seconds - ] - ) - for s in sources - ] - if not all(len(p) == rv.COMPOSITE_FPS * seconds for p in parts): - failures.append(f"a composite segment is short: {[len(p) for p in parts]} frames") - comp_path = _write_once(out_dir / "composite.mp4", np.concatenate(parts), fps=rv.COMPOSITE_FPS) + order = _composite_sources([e.clip for e in kept]) + sources = [next(e for e in kept if e.clip is c) for c in order] + comp_path, lengths = _composite(order, out_dir) + if not all(n_ == rv.COMPOSITE_FPS * rv.SEGMENT_SECONDS for n_ in lengths): + failures.append(f"a composite segment is short: {lengths} frames") comp = rv.embed_clip(emb, rv.Clip(name="composite.mp4", label="composite", path=comp_path)) report.put("composite_window_starts", comp.starts.tolist()) if not (len(comp.starts) == 29 and np.allclose(comp.starts, np.arange(29), atol=0.05)): @@ -332,3 +348,157 @@ def cen(c, value=centre): report.put("failures", failures) print(f"wrote {report.write(out_dir)} in {report.data['seconds']} s") assert not failures, "; ".join(failures) + + +def test_search_reproduces_experiment(tmp_path_factory): + """T0's result through the product code: the folder indexer, the index file and ``search``. + + The 50 val clips are copied out of the Hugging Face cache into a folder with one subfolder per + class (so the index is never written beside the cache), indexed recursively with the CLI + defaults, and scored with ``nearest_clips``. The vectors come from a fresh embedding run, so + the clip-level count may differ from T0's by one clip. Localisation runs on a second index + that also holds T0's composite, through ``nearest_moments`` as the viewer calls it. The + clip-level count and the window-level rate must match T0's within one clip (0.02), and + localisation may miss at most one of T0's 16 passes. + """ + import json + import shutil + + from huggingface_hub import snapshot_download + + from jepa_scope import index_file, search + from jepa_scope.cache import EmbedSettings, PretrainedModel + from jepa_scope.indexing import INDEX_NAME, build_or_refresh + from jepa_scope.model import DEFAULT_MODEL + + from . import retrieval_eval as rv + + t_start = time.perf_counter() + model_dir = Path(snapshot_download(DEFAULT_MODEL, local_files_only=True)) + assert model_dir.name == rv.MODEL_REVISION, f"model snapshot is {model_dir.name}, not the frozen one" + root = Path( + snapshot_download(rv.DATASET, repo_type="dataset", revision=rv.DATASET_REVISION, allow_patterns="val/*") + ) + clips = rv.manifest(root, "val") + assert len(clips) == len(rv.CLASSES) * rv.CLIPS_PER_CLASS + out_dir = _eval_dir(tmp_path_factory) + settings = EmbedSettings(model=DEFAULT_MODEL, fps=rv.SAMPLE_FPS, window=rv.WINDOW) + model = PretrainedModel(settings) + result = {"t0": rv.T0_VAL} + + def folder_of(name: str, extra: list[Path] = ()) -> Path: + folder = tmp_path_factory.mktemp(name) + for c in clips: + (folder / c.label).mkdir(exist_ok=True) + shutil.copyfile(c.path, folder / c.label / c.path.name) + for p in extra: + shutil.copyfile(p, folder / p.name) + return folder + + def build(folder: Path): + t0 = time.perf_counter() + report = build_or_refresh(folder, model, settings, recursive=True, log=lambda m: print(m, flush=True)) + assert report.skipped == [], report.skipped + index = index_file.load(folder / INDEX_NAME) + return index, search.prepare(index), report, round(time.perf_counter() - t0, 1) + + # Clip-level and window-level retrieval. + index, space, report, seconds = build(folder_of("val")) + header = index.header + assert header["embedding"]["similarity"] == search.DEFAULT_SIMILARITY == "cosine" + assert header["representation"]["pooling"] == "mean_all" and header["representation"]["token_norm"] == "none" + paths = [c["path"] for c in header["clips"]] + labels = [p.split("/")[0] for p in paths] + n = len(paths) + assert n == len(clips) + clip_hits = [] + # T0 excluded the query's YouTube source group; nearest_clips excludes only the query clip. + # These coincide because every clip in the frozen manifest has its own id (T0 asserts it). + for i in range(n): + hits = search.nearest_clips(space, i, 1).hits + clip_hits.append(bool(hits) and labels[hits[0].clip] == labels[i]) + count = sum(clip_hits) + per_class = {k: sum(h for h, lab in zip(clip_hits, labels, strict=True) if lab == k) for k in rv.CLASSES} + window_hits = [] + for i in range(n): + rows = index.clip_rows(i) + for w in range(rows.stop - rows.start): + hits = search.nearest_moments(space, i, w, 1).hits + window_hits.append(bool(hits) and labels[hits[0].clip] == labels[i]) + result.update( + clips=n, + windows=int(len(index.pooled)), + embedded=report.embedded, + index_seconds=seconds, + clip_top1_count=count, + clip_top1_rate=count / n, + per_class_hits=per_class, + window_top1_rate=sum(window_hits) / len(window_hits), + per_clip=[{"clip": p, "hit": h} for p, h in zip(paths, clip_hits, strict=True)], + ) + + # Localisation on a second index that also holds the composite. + comp_path, lengths = _composite(_composite_sources(clips), out_dir) + assert all(n_ == rv.COMPOSITE_FPS * rv.SEGMENT_SECONDS for n_ in lengths), lengths + index2, space2, report2, seconds2 = build(folder_of("val-composite", [comp_path])) + paths2 = [c["path"] for c in index2.header["clips"]] + comp = paths2.index("composite.mp4") + comp_starts = index2.window_start_time[index2.clip_rows(comp)] + assert len(comp_starts) == 29 and np.allclose(comp_starts, np.arange(29), atol=0.05), comp_starts + window_seconds = rv.WINDOW / rv.SAMPLE_FPS + boundaries = [rv.SEGMENT_SECONDS * k for k in range(1, len(comp_starts) + 1)] + queries = [] + for i, source in enumerate(_composite_sources(clips)[:4]): + own = paths2.index(f"{source.label}/{source.path.name}") + starts = index2.window_start_time[index2.clip_rows(own)] + for w in range(rv.QUERY_WINDOWS): + target = rv.SEGMENT_SECONDS * i + float(starts[w]) + row = {"segment": i, "query_start": float(starts[w]), "target": target} + for key, same in (("cross_clip", False), ("same_clip", True)): + top = search.nearest_moments(space2, own, w, 1, same_clip=same).hits[0] + on_comp = top.clip == comp + crossing = on_comp and rv.transition(top.start_time, boundaries, window_seconds) + row[key] = { + "hit": paths2[top.clip], + "hit_start": top.start_time, + "transition": crossing, + # A transition window was not a candidate in T0, so a hit on one counts as a miss. + "pass": on_comp and not crossing and abs(top.start_time - target) <= rv.TOLERANCE_SECONDS + 1e-6, + "own_clip": top.clip == own, + } + queries.append(row) + loc = { + key: { + "passed": sum(q[key]["pass"] for q in queries), + "transition_hits": sum(q[key]["transition"] for q in queries), + "own_clip_hits": sum(q[key]["own_clip"] for q in queries), + "of": len(queries), + } + for key in ("cross_clip", "same_clip") + } + result.update(localisation=loc, localisation_queries=queries, composite_index_seconds=seconds2) + result.update(embedded_composite_index=report2.embedded, seconds=round(time.perf_counter() - t_start, 1)) + bars = rv.at_least(count, n, rv.BAR_RATE) and rv.at_least(count - rv.T0_VAL["colour_baseline"], n, rv.BAR_MARGIN) + result["meets_bars"] = bars + + print(f"\n=== search reproduces T0 (val): {n} clips, {result['windows']} windows") + print(f"index built in {seconds} s, embedded {report.embedded}; composite index {seconds2} s") + print(_rate_line("clip-level top-1", count, n, f" (T0 {rv.T0_VAL['clip_top1_count']}/{n}) {per_class}")) + print(f" window-level top-1 rate {result['window_top1_rate']:.4f} (T0 {rv.T0_VAL['window_top1_rate']:.4f})") + for key, r in loc.items(): + print( + f"localisation {key}: {r['passed']}/{r['of']} pass; {r['transition_hits']} on a transition window; " + f"{r['own_clip_hits']} in the query's own clip" + ) + path = out_dir / "search-reproduction-val.json" + path.write_text(json.dumps(result, indent=2, sort_keys=True)) + print(f"wrote {path} in {result['seconds']} s") + + assert abs(count - rv.T0_VAL["clip_top1_count"]) <= 1, f"{count}/{n} against T0's {rv.T0_VAL['clip_top1_count']}" + assert bars, f"{count}/{n} misses the bars (rate {rv.BAR_RATE}, margin {rv.BAR_MARGIN} over the colour baseline)" + # nearest_moments carries the exclusion, overlap and tie rules that nearest_clips does not, so + # its results are held to T0's too (same_clip=True is descriptive only). + rate, t0_rate = result["window_top1_rate"], rv.T0_VAL["window_top1_rate"] + assert abs(rate - t0_rate) <= 0.02, f"window-level top-1 {rate:.4f} against T0's {t0_rate:.4f}" + passed, t0_passed = loc["cross_clip"]["passed"], rv.T0_VAL["localisation_passed"] + assert passed >= t0_passed - 1, f"localisation {passed}/{len(queries)} against T0's {t0_passed}"