hnsw: move pruned/no-recompute assertion into backend; api: drop global assertion; docs: will adjust after benchmarking

This commit is contained in:
Andy Lee
2025-08-14 11:08:34 -07:00
parent 3e162fb177
commit 2bd557d1cf
3 changed files with 88 additions and 3 deletions

View File

@@ -515,6 +515,7 @@ Options:
--top-k N Number of results (default: 5)
--complexity N Search complexity (default: 64)
--recompute Use recomputation for highest accuracy
--no-recompute Disable recomputation (requires non-compact HNSW index)
--pruning-strategy {global,local,proportional}
```

View File

@@ -0,0 +1,82 @@
import os
import time
from pathlib import Path
from leann import LeannBuilder, LeannSearcher
def ensure_index(
index_path: str, num_docs: int = 5000, is_recompute: bool = True, is_compact: bool = True
):
path = Path(index_path)
if (path.parent / f"{path.stem}.meta.json").exists():
return
builder = LeannBuilder(
backend_name="hnsw",
embedding_model=os.getenv("LEANN_EMBED_MODEL", "facebook/contriever"),
embedding_mode=os.getenv("LEANN_EMBED_MODE", "sentence-transformers"),
graph_degree=32,
complexity=64,
is_compact=is_compact,
is_recompute=is_recompute,
num_threads=4,
)
for i in range(num_docs):
builder.add_text(
f"This is a test document number {i}. It contains some repeated text for benchmarking."
)
builder.build_index(index_path)
def bench_once(index_path: str, recompute: bool, top_k: int = 10) -> float:
searcher = LeannSearcher(index_path=index_path)
t0 = time.time()
_ = searcher.search(
"test document number 42",
top_k=top_k,
complexity=64,
prune_ratio=0.0,
recompute_embeddings=recompute,
)
return time.time() - t0
def main():
base = Path.cwd() / ".leann" / "indexes" / "bench"
base.parent.mkdir(parents=True, exist_ok=True)
index_path_recompute = str(base / "recompute.leann")
index_path_norecompute = str(base / "norecompute.leann")
# Build two variants: pruned (recompute) and non-compact (no-recompute)
ensure_index(index_path_recompute, is_recompute=True, is_compact=True)
ensure_index(index_path_norecompute, is_recompute=False, is_compact=False)
# Warm up
bench_once(index_path_recompute, recompute=True)
bench_once(index_path_norecompute, recompute=False)
t_recompute = bench_once(index_path_recompute, recompute=True)
t_norecompute = bench_once(index_path_norecompute, recompute=False)
size_recompute = sum(
f.stat().st_size for f in Path(index_path_recompute).parent.iterdir() if f.is_file()
)
size_norecompute = sum(
f.stat().st_size for f in Path(index_path_norecompute).parent.iterdir() if f.is_file()
)
print("Benchmark results (HNSW):")
print(
f" recompute=True: search_time={t_recompute:.3f}s, size={size_recompute / 1024 / 1024:.1f}MB"
)
print(
f" recompute=False: search_time={t_norecompute:.3f}s, size={size_norecompute / 1024 / 1024:.1f}MB"
)
print("Expectation: no-recompute should be faster but larger on disk.")
if __name__ == "__main__":
main()

View File

@@ -185,9 +185,11 @@ class HNSWSearcher(BaseSearcher):
"""
from . import faiss # type: ignore
if not recompute_embeddings:
if self.is_pruned:
raise RuntimeError("Recompute is required for pruned index.")
if not recompute_embeddings and self.is_pruned:
raise RuntimeError(
"Recompute is required for pruned/compact HNSW index. "
"Re-run search with --recompute, or rebuild with --no-recompute and --no-compact."
)
if recompute_embeddings:
if zmq_port is None:
raise ValueError("zmq_port must be provided if recompute_embeddings is True")