diff --git a/README.md b/README.md index 4fb836e..677cc6d 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,65 @@ 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. + +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 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 +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 +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/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..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 @@ -106,3 +106,125 @@ 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). + +- **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). + +- **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). + +- **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 2da4383..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 @@ -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 @@ -593,7 +594,7 @@ flowchart LR ### T3 — Week 3. Search on hand-built vectors -- **Status:** Pending +- **Status:** Complete - **Dependencies:** T2 - **Parallel group:** Sequential - **Complexity:** ROUTINE @@ -636,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 @@ -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 @@ -689,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/src/jepa_scope/cache.py b/src/jepa_scope/cache.py new file mode 100644 index 0000000..1dbadbc --- /dev/null +++ b/src/jepa_scope/cache.py @@ -0,0 +1,221 @@ +"""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). + + 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=model, + fps=args.fps, + window=args.window, + batch=args.batch, + max_seconds=args.max_seconds, + device=args.device, + 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 + + 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 + + +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 37c3ec7..eec7d8d 100644 --- a/src/jepa_scope/cli.py +++ b/src/jepa_scope/cli.py @@ -3,25 +3,27 @@ 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 os 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 +33,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 +89,164 @@ 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) + + +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 + + 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: + 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, 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) + 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: + assert query_sha is not None + if args.same_clip: + raise SystemExit("error: --same-clip applies to a clip in the index only") + 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) + 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], shas: list[str], root: Path) -> tuple[int | None, str | None]: + """Which indexed clip ``query`` names, as ``(clip, sha256)``. - 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. + 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. """ - 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 + if query.is_file(): + path = Path(os.path.abspath(query)) + path = path.parent.resolve() / path.name + try: + rel = path.relative_to(root.resolve()).as_posix() + except ValueError: + rel = None + if rel in names: + 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"] + 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 + # 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: - 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 + 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: @@ -149,9 +262,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 +288,43 @@ 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=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") + + 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 +332,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/index_file.py b/src/jepa_scope/index_file.py new file mode 100644 index 0000000..3fb900b --- /dev/null +++ b/src/jepa_scope/index_file.py @@ -0,0 +1,337 @@ +"""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", + "recursive", + "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") + 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 + ): + 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/src/jepa_scope/indexing.py b/src/jepa_scope/indexing.py new file mode 100644 index 0000000..5507718 --- /dev/null +++ b/src/jepa_scope/indexing.py @@ -0,0 +1,518 @@ +"""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. 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. 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 + +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 __version__, index_file, latent_file +from .cache import EmbedSettings, cache_path, find_cache_file +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 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): + """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``. + + ``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. + """ + + def __init__(self, reason: str, remember: bool = False): + super().__init__(reason) + self.remember = remember + self.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 + # 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]]: + """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, hidden_size: int) -> ClipVectors: + """Read the window vectors and the clip facts from a ``.latent`` file, without its patches. + + 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: + 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 + 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 find(self, video_sha256: str) -> ClipVectors | None: + """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.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, expected["hidden_size"]) + except ClipSkipped as exc: + self.log(f"re-embedding: {exc}") + 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; 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 + found = self.find(video_sha256) + if found is not None: + return found + 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 + return vectors + + def _embed(self, video: Path, video_sha256: str) -> ClipVectors: + s = self.settings + expected = self.model.representation() + 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 CLIP_ERRORS as 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( + "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, 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 _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] +) -> _Previous: + if not out.exists(): + return _Previous() + try: + header = index_file.load(out).header + except (index_file.IndexFileError, OSError) as 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: + 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"]: + 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]) + prev.rebuilt = f"the run settings changed ({changed})" + elif old_root.resolve() != folder.resolve(): + prev.rebuilt = f"the existing index was built from another folder ({old_root})" + else: + 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( + folder: str | Path, + model: ModelProvider, + settings: EmbedSettings, + out: str | Path | None = None, + recursive: bool | None = None, + patches: bool = False, + centre: bool | None = None, + 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. + + ``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(): + 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(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=previous.rebuilt, centre=centre, recursive=recursive) + + 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 _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.clips.get(rel), clips, verify, report) + except ClipSkipped as exc: + 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 + 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": similarity}, + "root": _root_for(folder, out), + "recursive": recursive, + "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 + return clips.get(path, sha), 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 _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() + 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 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.""" + 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/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/retrieval_eval.py b/tests/retrieval_eval.py index cdde3dc..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: @@ -119,8 +129,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 +165,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_index.py b/tests/test_index.py new file mode 100644 index 0000000..c966a30 --- /dev/null +++ b/tests/test_index.py @@ -0,0 +1,649 @@ +"""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": ".", + "recursive": True, + "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_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" + 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("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"), + 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] diff --git a/tests/test_indexing.py b/tests/test_indexing.py new file mode 100644 index 0000000..2cad3f1 --- /dev/null +++ b/tests/test_indexing.py @@ -0,0 +1,852 @@ +"""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 av +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, centre=False) + 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()) + + +@pytest.mark.parametrize( + "error", + [ + 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 + + 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_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) + 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 + + +@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() + 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].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)) + 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( + folder, patched_cli, capsys, monkeypatch, tmp_path, tiny_encoder +): + cli.main(["index", str(folder), *EMBED_ARGS]) + query = tmp_path / "query.mp4" + 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)]) + + +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_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() + 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..3945213 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( @@ -266,8 +293,39 @@ 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 cli + from jepa_scope import cache tensors, header = latent_file.load(latent) old = tmp_path / "v1.latent" @@ -275,18 +333,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 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}" diff --git a/tests/test_search.py b/tests/test_search.py new file mode 100644 index 0000000..f8eed9e --- /dev/null +++ b/tests/test_search.py @@ -0,0 +1,530 @@ +"""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": ".", + "recursive": False, + "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)