fecee94af1
* feat: finance bench * docs: results * chore: ignroe data README * feat: fix financebench * feat: laion, also required idmaps support * style: format * style: format * fix: resolve ruff linting errors - Remove unused variables in benchmark scripts - Rename unused loop variables to follow convention * feat: enron email bench * experiments for running DiskANN & BM25 on Arch 4090 * style: format * chore(ci): remove paru-bin submodule and config to fix checkout --recurse-submodules * docs: data * docs: data updated * fix: as package * fix(ci): only run pre-commit * chore: use http url of astchunk; use group for some dev deps * fix(ci): should checkout modules as well since `uv sync` checks * fix(ci): run with lint only * fix: find links to install wheels available * CI: force local wheels in uv install step * CI: install local wheels via file paths * CI: pick wheels matching current Python tag * CI: handle python tag mismatches for local wheels * CI: use matrix python venv and set macOS deployment target * CI: revert install step to match main * CI: use uv group install with local wheel selection * CI: rely on setup-uv for Python and tighten group install * CI: install build deps with uv python interpreter * CI: use temporary uv venv for build deps * CI: add build venv scripts path for wheel repair
125 lines
3.7 KiB
Python
125 lines
3.7 KiB
Python
# /// script
|
|
# dependencies = [
|
|
# "leann-backend-diskann"
|
|
# ]
|
|
# ///
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def load_queries(path: Path, limit: int | None) -> list[str]:
|
|
out: list[str] = []
|
|
with open(path, encoding="utf-8") as f:
|
|
for line in f:
|
|
obj = json.loads(line)
|
|
out.append(obj["query"])
|
|
if limit and len(out) >= limit:
|
|
break
|
|
return out
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(
|
|
description="DiskANN baseline on real NQ queries (search-only timing)"
|
|
)
|
|
ap.add_argument(
|
|
"--index-dir",
|
|
default="benchmarks/data/indices/diskann_rpj_wiki",
|
|
help="Directory containing DiskANN files",
|
|
)
|
|
ap.add_argument("--index-prefix", default="ann")
|
|
ap.add_argument("--queries-file", default="benchmarks/data/queries/nq_open.jsonl")
|
|
ap.add_argument("--num-queries", type=int, default=200)
|
|
ap.add_argument("--top-k", type=int, default=10)
|
|
ap.add_argument("--complexity", type=int, default=62)
|
|
ap.add_argument("--threads", type=int, default=1)
|
|
ap.add_argument("--beam-width", type=int, default=1)
|
|
ap.add_argument("--cache-mechanism", type=int, default=2)
|
|
ap.add_argument("--num-nodes-to-cache", type=int, default=0)
|
|
args = ap.parse_args()
|
|
|
|
index_dir = Path(args.index_dir).resolve()
|
|
if not index_dir.is_dir():
|
|
raise SystemExit(f"Index dir not found: {index_dir}")
|
|
|
|
qpath = Path(args.queries_file).resolve()
|
|
if not qpath.exists():
|
|
raise SystemExit(f"Queries file not found: {qpath}")
|
|
|
|
queries = load_queries(qpath, args.num_queries)
|
|
print(f"Loaded {len(queries)} queries from {qpath}")
|
|
|
|
# Compute embeddings once (exclude from timing)
|
|
from leann.api import compute_embeddings as _compute
|
|
|
|
embs = _compute(
|
|
queries,
|
|
model_name="facebook/contriever-msmarco",
|
|
mode="sentence-transformers",
|
|
use_server=False,
|
|
).astype(np.float32)
|
|
if embs.ndim != 2:
|
|
raise SystemExit("Embedding compute failed or returned wrong shape")
|
|
|
|
# Build searcher
|
|
from leann_backend_diskann.diskann_backend import DiskannSearcher as _DiskannSearcher
|
|
|
|
index_prefix_path = str(index_dir / args.index_prefix)
|
|
searcher = _DiskannSearcher(
|
|
index_prefix_path,
|
|
num_threads=int(args.threads),
|
|
cache_mechanism=int(args.cache_mechanism),
|
|
num_nodes_to_cache=int(args.num_nodes_to_cache),
|
|
)
|
|
|
|
# Warmup (not timed)
|
|
_ = searcher.search(
|
|
embs[0:1],
|
|
top_k=args.top_k,
|
|
complexity=args.complexity,
|
|
beam_width=args.beam_width,
|
|
prune_ratio=0.0,
|
|
recompute_embeddings=False,
|
|
batch_recompute=False,
|
|
dedup_node_dis=False,
|
|
)
|
|
|
|
# Timed loop
|
|
times: list[float] = []
|
|
for i in range(embs.shape[0]):
|
|
t0 = time.time()
|
|
_ = searcher.search(
|
|
embs[i : i + 1],
|
|
top_k=args.top_k,
|
|
complexity=args.complexity,
|
|
beam_width=args.beam_width,
|
|
prune_ratio=0.0,
|
|
recompute_embeddings=False,
|
|
batch_recompute=False,
|
|
dedup_node_dis=False,
|
|
)
|
|
times.append(time.time() - t0)
|
|
|
|
times_sorted = sorted(times)
|
|
avg = float(sum(times) / len(times))
|
|
p50 = times_sorted[len(times) // 2]
|
|
p95 = times_sorted[max(0, int(len(times) * 0.95) - 1)]
|
|
|
|
print("\nDiskANN (NQ, search-only) Report")
|
|
print(f" queries: {len(times)}")
|
|
print(
|
|
f" k: {args.top_k}, complexity: {args.complexity}, beam_width: {args.beam_width}, threads: {args.threads}"
|
|
)
|
|
print(f" avg per query: {avg:.6f} s")
|
|
print(f" p50/p95: {p50:.6f}/{p95:.6f} s")
|
|
print(f" QPS: {1.0 / avg:.2f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|