hnsw: move pruned/no-recompute assertion into backend; api: drop global assertion; docs: will adjust after benchmarking
This commit is contained in:
@@ -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}
|
||||
```
|
||||
|
||||
|
||||
82
benchmarks/benchmark_no_recompute.py
Normal file
82
benchmarks/benchmark_no_recompute.py
Normal 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()
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user