| """Per-query retrieval latency with a fixed thread count (HF CPU Basic profile: 2 vCPU).
|
|
|
| Run on an otherwise idle machine after scripts/run_all.py: uv run python scripts/latency.py
|
| """
|
| import os
|
| import sys
|
| from pathlib import Path
|
|
|
| import yaml
|
|
|
| ROOT = Path(__file__).resolve().parents[1]
|
| CFG = yaml.safe_load((ROOT / "configs" / "default.yaml").read_text())
|
| THREADS = str(CFG["latency"]["threads"])
|
| for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
|
| os.environ[var] = THREADS
|
|
|
| import time
|
|
|
| import faiss
|
| import numpy as np
|
| import pandas as pd
|
| import torch
|
|
|
| from gatestats import data, retrieve
|
|
|
|
|
| def main():
|
| lat = CFG["latency"]
|
| torch.set_num_threads(lat["threads"])
|
| faiss.omp_set_num_threads(lat["threads"])
|
| name, cache = lat["dataset"], ROOT / "cache" / lat["dataset"]
|
| ds = data.load(name)
|
|
|
| fns = {}
|
| for system in lat["systems"]:
|
| if system == "bm25":
|
| index = retrieve.bm25_index(ds.doc_texts)
|
| fns[system] = lambda q, index=index: retrieve.bm25_run(index, [q], 10)
|
| continue
|
| key = CFG["hnsw"]["model"] if system.startswith("hnsw") else system
|
| mcfg = CFG["models"][key]
|
| model = retrieve.load_model(mcfg, CFG["max_seq_length"])
|
| if system.startswith("hnsw"):
|
| index = faiss.read_index(str(cache / "hnsw_s0.faiss"))
|
| index.hnsw.efSearch = int(system.split("_ef")[1].split("_")[0])
|
| else:
|
| index = retrieve.exact_index(np.load(cache / f"emb_{key}_docs.npy"))
|
| fns[system] = lambda q, m=model, p=mcfg["query_prefix"], ix=index: ix.search(
|
| retrieve.embed(m, [q], p, progress=False), 10)
|
|
|
| rows, rng = [], np.random.default_rng([CFG["seed"], 5])
|
| for run in range(lat["runs"]):
|
| for system, fn in fns.items():
|
| order = rng.permutation(len(ds.query_ids))
|
| for i in order[: lat["warmup"]]:
|
| fn(ds.query_texts[i])
|
| for i in order:
|
| t0 = time.perf_counter()
|
| fn(ds.query_texts[i])
|
| rows.append((system, run, ds.query_ids[i], (time.perf_counter() - t0) * 1000))
|
| print(f"run {run + 1}/{lat['runs']} done", file=sys.stderr, flush=True)
|
| out = pd.DataFrame(rows, columns=["system", "run", "qid", "ms"])
|
| out.to_parquet(ROOT / "results" / "latency.parquet", index=False)
|
| print(out.groupby("system")["ms"].describe(percentiles=[0.5, 0.95, 0.99]).round(2))
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
