File size: 2,691 Bytes
668f59a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """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 # must be set before numpy/torch/faiss load
import time # noqa: E402
import faiss # noqa: E402
import numpy as np # noqa: E402
import pandas as pd # noqa: E402
import torch # noqa: E402
from gatestats import data, retrieve # noqa: E402
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()
|