Public symbols exported from the dcee package (see dcee/__init__.py). Types use numpy.ndarray float32 rows unless noted.
from dcee import (
DCEEConfig,
DCEEEngine,
DCEEIndex,
ClusterBlockV2,
is_gpu_available,
load_index,
save_index,
)
import dcee
print(dcee.__version__)is_gpu_available()Returns: bool — True if CuPy is installed and a CUDA device is visible (search/reconstruction use the GPU). Otherwise computations use NumPy on CPU.
from dcee import is_gpu_available
if is_gpu_available():
print("Using CuPy / CUDA")
else:
print("Using NumPy on CPU")DCEEConfigA @dataclass holding all build and query hyperparameters. Construct with explicit fields or use tuned_for (below).
| Field | Type / default | Description |
|---|---|---|
dim | int, 128 | Embedding dimension (must match columns of emb in build). |
n_clusters | int, 64 | MiniBatchKMeans cluster count for routing. |
keyframe_every | int, 16 | Reset interval inside each ordered cluster (keyframes vs deltas). |
quantization | str, "int8" | "int8" | "float16" | "float32" for stored deltas. |
top_k_refine | int, 20 | Max coarse candidates before full-precision cosine refinement. |
n_probe | int, 8 | Minimum number of top keyframe clusters to score per query. |
n_probe_max | int, 32 | Upper cap when Adaptive Margin Probing expands the probe set. |
batch_size | int, 4096 | MiniBatchKMeans batch size during clustering. |
adaptive_probe | bool, True | Enable AMP (extra clusters when keyframe scores are flat). |
adaptive_probe_margin | float, 0.028 | Score gap threshold between sorted keyframes to keep adding clusters. |
verbose | bool, True | Logging and tqdm during build / tensor preparation. |
DCEEConfig.tuned_for(n_vectors: int, dim: int) -> DCEEConfigStatic helper returning a new config with scale-aware defaults (roughly √N clusters, probe budgets, refinement pool, keyframe spacing, batch size).
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 = 24DCEEEngineHigh-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.
from dcee import DCEEConfig, DCEEEngine
cfg = DCEEConfig(dim=128, n_clusters=32, verbose=False)
engine = DCEEEngine(cfg)build(self, emb: np.ndarray) -> NoneCluster 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.
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.
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) -> NoneWrite binary index (.dce2). Raises if no index was built. Uses save_index internally; format version is written in the file header.
engine.save("artifacts/index.dce2")load(self, path: str | pathlib.Path) -> NoneReplace this engine's state from disk (same process). Equivalent to reconstructing from from_file but mutates self.
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.
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_indexLow-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) -> Noneload_index(path, *, verbose: bool = True) -> DCEEIndexload_index raises ValueError for bad magic bytes or unsupported format version.
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 == 32DCEEIndex & ClusterBlockV2DCEEIndex: 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.