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)
|
--top-k N Number of results (default: 5)
|
||||||
--complexity N Search complexity (default: 64)
|
--complexity N Search complexity (default: 64)
|
||||||
--recompute Use recomputation for highest accuracy
|
--recompute Use recomputation for highest accuracy
|
||||||
|
--no-recompute Disable recomputation (requires non-compact HNSW index)
|
||||||
--pruning-strategy {global,local,proportional}
|
--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
|
from . import faiss # type: ignore
|
||||||
|
|
||||||
if not recompute_embeddings:
|
if not recompute_embeddings and self.is_pruned:
|
||||||
if self.is_pruned:
|
raise RuntimeError(
|
||||||
raise RuntimeError("Recompute is required for pruned index.")
|
"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 recompute_embeddings:
|
||||||
if zmq_port is None:
|
if zmq_port is None:
|
||||||
raise ValueError("zmq_port must be provided if recompute_embeddings is True")
|
raise ValueError("zmq_port must be provided if recompute_embeddings is True")
|
||||||
|
|||||||
Reference in New Issue
Block a user