blinklinmult.data.embeddings

Precomputed eye embeddings, for frozen-encoder runs.

A frozen encoder computes the identical 256-d vector for a given crop on every epoch: the weights are fixed and so is the input, so every epoch after the first recomputes a constant. Caching that turns the expensive part of a batch into a lookup.

Measured on this benchmark: a batch is batch_size x 45 crops through ConvNeXt, costing 0.74 s with a frozen encoder against milliseconds from cache, and 3.4 GB of embeddings replaces 51 GB of crops – small enough to sit in RAM, which removes the I/O question along with the compute one.

Two preconditions, both refused in config rather than warned about (DataConfig):

  • model.encoder_freeze – an encoder that is still learning produces different embeddings after every step, so a cache is stale from step one.

  • data.augment must be None – augmentation perturbs the crop before the encoder, so caching would freeze one augmented view per window for the whole run, keeping the distortion and losing the variety.

Either mistake trains against embeddings from a model that no longer exists and reports a plausible, wrong number. That is the same failure class as the carrier-corpus contamination and the write-only quality signals found earlier in this project, so the checks here are hard errors that name the mismatched field.

The cache is keyed by the encoder’s own weights. encoder_digest hashes the state dict, and the digest is part of the directory path, so pointing model.encoder_weights at a different arm misses the cache and rebuilds rather than silently reusing the wrong vectors.

Caches are not deleted after a run. One cache serves every frozen arm that shares a seed checkpoint – deleting it between arms would rebuild it each time, at a full forward pass over the corpus per rebuild, to reclaim 3.4 GB against the corpora’s 51 GB. make clean-cache removes them when that is actually wanted.

Attributes

logger

Module-level logger.

MANIFEST_NAME

Per-shard sidecar describing what the cache holds and what built it.

EMBEDDING_DTYPE

Stored precision.

Exceptions

EmbeddingCacheError

Raised when a cache cannot be used for the run that asked for it.

Classes

CacheKey

Everything that determines a cache's contents.

Functions

encoder_digest(→ str)

Hash an encoder's parameters, to key a cache by the weights that built it.

shard_path(→ pathlib.Path)

Where one corpus-split's embeddings live.

write_shard(→ None)

Write one corpus-split's embeddings and its manifest.

read_shard(→ dict[str, numpy.ndarray])

Load a shard, refusing one built by a different encoder or config.

describe(→ list[dict[str, object]])

Every shard under a cache root, with what it holds.

Module Contents

blinklinmult.data.embeddings.logger

Module-level logger.

blinklinmult.data.embeddings.MANIFEST_NAME = 'manifest.json'

Per-shard sidecar describing what the cache holds and what built it.

blinklinmult.data.embeddings.EMBEDDING_DTYPE

Stored precision.

Half the size at no measurable cost: these are inputs to a d_model=32 transformer, far above float16’s resolution. The equivalence test in tests/data/test_embeddings.py asserts the round trip stays within tolerance rather than assuming it.

exception blinklinmult.data.embeddings.EmbeddingCacheError

Bases: RuntimeError

Raised when a cache cannot be used for the run that asked for it.

Initialize self. See help(type(self)) for accurate signature.

class blinklinmult.data.embeddings.CacheKey

Everything that determines a cache’s contents.

Compared field by field on load, so a mismatch names the field that changed rather than failing obscurely later. A cache that silently disagreed with its run would train on the wrong vectors and report a believable number.

encoder_digest

Hash of the encoder’s state_dict.

Type:

str

backbone

Backbone name, for a readable error.

Type:

str

output_dim

Embedding width.

Type:

int

image_size

Crop size the encoder was fed.

Type:

int

as_dict() → dict[str, object]

Serialise for the manifest.

Returns:

JSON-safe fields.

Return type:

dict[str, object]

blinklinmult.data.embeddings.encoder_digest(encoder: torch.nn.Module) → str

Hash an encoder’s parameters, to key a cache by the weights that built it.

Parameters are hashed in sorted name order so the digest is stable across runs, and on the CPU in float32 so it does not depend on the device or the autocast state a run happened to use.

Parameters:

encoder (torch.nn.Module) – The eye encoder.

Returns:

Hex SHA-256 of the parameter tensors.

Return type:

str

blinklinmult.data.embeddings.shard_path(root: pathlib.Path | str, key: CacheKey, corpus: str, subset: str) → pathlib.Path

Where one corpus-split’s embeddings live.

The digest is a path component, not merely a manifest field: a different encoder lands in a different directory, so the wrong cache cannot be picked up even if a manifest check were somehow skipped.

Parameters:

  • root (Path | str) – Cache root from data.embedding_cache.

  • key (CacheKey) – Identity of the encoder that built it.

  • corpus (str) – Corpus name.

  • subset (str) – Split name.

Returns:

The shard’s .h5 path.

Return type:

Path

blinklinmult.data.embeddings.write_shard(path: pathlib.Path, key: CacheKey, keys: list[str], embeddings: numpy.ndarray, masks: numpy.ndarray) → None

Write one corpus-split’s embeddings and its manifest.

Written to a temporary file and renamed, so an interrupted build leaves no half-written shard that a later run would trust.

Parameters:

  • path (Path) – Destination, from shard_path().

  • key (CacheKey) – Identity to record.

  • keys (list[str]) – Sample ids, in row order.

  • embeddings (np.ndarray) – (N, T, D).

  • masks (np.ndarray) – (N, T) bool.

blinklinmult.data.embeddings.read_shard(path: pathlib.Path, key: CacheKey) → dict[str, numpy.ndarray]

Load a shard, refusing one built by a different encoder or config.

Parameters:

  • path (Path) – The shard.

  • key (CacheKey) – What this run requires.

Returns:

keys, embedding and mask.

Return type:

dict[str, np.ndarray]

Raises:

EmbeddingCacheError – If the shard or its manifest is missing, or any identity field disagrees with key.

blinklinmult.data.embeddings.describe(root: pathlib.Path | str) → list[dict[str, object]]

Every shard under a cache root, with what it holds.

Answers “which corpora and splits are cached, and which encoder built them?” without reading the embedding arrays – the manifests alone carry it. Sorted by encoder digest, then corpus, then split, so several encoders’ caches under one root stay visually grouped.

Parameters:

root (Path | str) – Cache root, as passed to data.embedding_cache.

Returns:

One manifest per shard, each with an added path. Empty when the root does not exist or holds no shard.

Return type:

list[dict[str, object]]