DCEE

API reference

Public symbols exported from the dcee package (see dcee/__init__.py). Types use numpy.ndarray float32 rows unless noted.

Recommended import

from dcee import (
    DCEEConfig,
    DCEEEngine,
    DCEEIndex,
    ClusterBlockV2,
    is_gpu_available,
    load_index,
    save_index,
)
import dcee

print(dcee.__version__)

is_gpu_available()

Returns: boolTrue if CuPy is installed and a CUDA device is visible (search/reconstruction use the GPU). Otherwise computations use NumPy on CPU.

Example
from dcee import is_gpu_available

if is_gpu_available():
    print("Using CuPy / CUDA")
else:
    print("Using NumPy on CPU")

DCEEConfig

A @dataclass holding all build and query hyperparameters. Construct with explicit fields or use tuned_for (below).

Fields

FieldType / defaultDescription
dimint, 128Embedding dimension (must match columns of emb in build).
n_clustersint, 64MiniBatchKMeans cluster count for routing.
keyframe_everyint, 16Reset interval inside each ordered cluster (keyframes vs deltas).
quantizationstr, "int8""int8" | "float16" | "float32" for stored deltas.
top_k_refineint, 20Max coarse candidates before full-precision cosine refinement.
n_probeint, 8Minimum number of top keyframe clusters to score per query.
n_probe_maxint, 32Upper cap when Adaptive Margin Probing expands the probe set.
batch_sizeint, 4096MiniBatchKMeans batch size during clustering.
adaptive_probebool, TrueEnable AMP (extra clusters when keyframe scores are flat).
adaptive_probe_marginfloat, 0.028Score gap threshold between sorted keyframes to keep adding clusters.
verbosebool, TrueLogging and tqdm during build / tensor preparation.

DCEEConfig.tuned_for(n_vectors: int, dim: int) -> DCEEConfig

Static helper returning a new config with scale-aware defaults (roughly √N clusters, probe budgets, refinement pool, keyframe spacing, batch size).

Example
import numpy as np
from dcee import DCEEConfig

emb = np.random.randn(50_000, 384).astype(np.float32)

cfg = DCEEConfig.tuned_for(len(emb), emb.shape[1])
assert cfg.dim == 384

# Override after creation
cfg.verbose = False
cfg.n_probe = 24

DCEEEngine

High-level facade around clustering, encoding, GPU/CPU query engine, and I/O.

DCEEEngine(cfg: DCEEConfig)

Construct an engine. Does not build until you call build() or load state via from_file / load.

Attributes (read-only use): cfg, index (DCEEIndex | None), internal preprocessor/encoder used after build.

Example
from dcee import DCEEConfig, DCEEEngine

cfg = DCEEConfig(dim=128, n_clusters=32, verbose=False)
engine = DCEEEngine(cfg)

build(self, emb: np.ndarray) -> None

Cluster emb (N × dim float32), delta-encode each non-empty cluster, precompute reconstructed tensors for query. Raises nothing on success; wrong dim leads to shape errors downstream.

Example
import numpy as np
from dcee import DCEEConfig, DCEEEngine

emb = np.random.randn(1000, 64).astype(np.float32)
emb /= np.linalg.norm(emb, axis=1, keepdims=True)

cfg = DCEEConfig(dim=64, n_clusters=16, verbose=False)
engine = DCEEEngine(cfg)
engine.build(emb)

search(self, query: np.ndarray, top_k: int = 5) -> list[tuple[int, float]]

Returns: list of (global_index, cosine_similarity) with length ≤ top_k. Query is any length-dim vector; it is normalized internally. Requires build() or from_file() first.

Example
query = emb[42]  # or any (dim,) float32 vector
hits = engine.search(query, top_k=10)
for idx, score in hits:
    print(idx, score)

save(self, path: str | pathlib.Path) -> None

Write binary index (.dce2). Raises if no index was built. Uses save_index internally; format version is written in the file header.

Example
engine.save("artifacts/index.dce2")

load(self, path: str | pathlib.Path) -> None

Replace this engine's state from disk (same process). Equivalent to reconstructing from from_file but mutates self.

Example
engine.load("artifacts/index.dce2")

from_file(path, *, verbose: bool | None = None) -> DCEEEngine (classmethod)

Load index from path; config is read from the file (no placeholder DCEEConfig needed). Returns a new engine ready for search.

Example
from pathlib import Path
from dcee import DCEEEngine

engine = DCEEEngine.from_file(Path("artifacts/index.dce2"), verbose=False)
hits = engine.search(query_vector, top_k=10)

load_index / save_index

Low-level I/O for DCEEIndex. Most applications should use DCEEEngine.save / from_file. Use these when you manage DCEEIndex yourself.

  • save_index(index: DCEEIndex, path, *, verbose: bool = True) -> None
  • load_index(path, *, verbose: bool = True) -> DCEEIndex

load_index raises ValueError for bad magic bytes or unsupported format version.

Example
from dcee import load_index, save_index, DCEEEngine, DCEEConfig
import numpy as np

# After you already have an engine with a built index:
engine = DCEEEngine(DCEEConfig(dim=32, n_clusters=4, verbose=False))
engine.build(np.random.randn(100, 32).astype(np.float32))

save_index(engine.index, "out.dce2", verbose=False)

idx = load_index("out.dce2", verbose=False)
assert idx.config.dim == 32

DCEEIndex & ClusterBlockV2

DCEEIndex: config, clusters (list of blocks), keyframe_matrix (num_nonempty_clusters × dim float32).

ClusterBlockV2: one cluster's packed deltas, scales, keyframe masks, and global row indices. Useful for advanced inspection or custom tooling — not required for normal search.


See also Configuration for tuning guidance and Examples for end-to-end snippets.