212 lines
8.1 KiB
Python
212 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Benchmark a running Studio's RAG over its real HTTP API: indexing latency,
|
|
scaling, and retrieval accuracy. Same script drives baseline and improved.
|
|
|
|
Usage:
|
|
python rag_studio_bench.py --base http://127.0.0.1:8901 --label baseline \
|
|
--password "<bootstrap>" --corpus ../data/rag_corpus --out ../logs/bench_baseline.json
|
|
"""
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
|
|
API = "/api/rag"
|
|
|
|
# Gold queries: (query, filename_substring, [accepted answer phrases any-of]).
|
|
# A hit counts when a retrieved chunk is FROM the right document AND contains an
|
|
# accepted phrase, so retrieval is scored independently of the chat model.
|
|
GOLD = [
|
|
("What are the sinusoidal positional encodings based on?", "attention",
|
|
["sine", "sinusoid", "wavelength"]),
|
|
("How many attention heads does the base Transformer use?", "attention",
|
|
["eight", "h = 8", "8 parallel", "8 attention", "h=8"]),
|
|
("What are BERT's two pre-training objectives?", "bert",
|
|
["masked", "next sentence"]),
|
|
("How many Transformer layers does BERT-large have?", "bert",
|
|
["24", "l = 24", "l=24"]),
|
|
("Difference between RAG-Sequence and RAG-Token models?", "rag",
|
|
["rag-token", "rag-sequence"]),
|
|
("Which retriever does RAG use to fetch passages?", "rag",
|
|
["dpr", "dense passage", "bi-encoder", "mips"]),
|
|
("What does the HTTP GET method do?", "rfc9110",
|
|
["transfer a current representation", "retrieve", "selector"]),
|
|
("Which HTTP status code means the resource was not found?", "rfc9110",
|
|
["404"]),
|
|
]
|
|
|
|
|
|
def login(c, base, user, pw):
|
|
new = pw + "Aa1!"
|
|
# Re-run safe: bootstrap pw works first time; after we change it, the changed
|
|
# pw works on later runs.
|
|
r = c.post(f"{base}/api/auth/login", json={"username": user, "password": pw})
|
|
if r.status_code == 401:
|
|
r = c.post(f"{base}/api/auth/login", json={"username": user, "password": new})
|
|
r.raise_for_status()
|
|
body = r.json()
|
|
tok = body["access_token"]
|
|
if body.get("must_change_password"):
|
|
r2 = c.post(f"{base}/api/auth/change-password", headers=H(tok),
|
|
json={"current_password": pw, "new_password": new})
|
|
r2.raise_for_status()
|
|
tok = r2.json()["access_token"]
|
|
return tok
|
|
|
|
|
|
def H(tok):
|
|
return {"Authorization": f"Bearer {tok}"}
|
|
|
|
|
|
def warmup(c, base, tok):
|
|
t = time.perf_counter()
|
|
try:
|
|
c.post(f"{base}{API}/warmup", headers=H(tok), timeout=600)
|
|
except Exception as e:
|
|
print("warmup err:", e)
|
|
return time.perf_counter() - t
|
|
|
|
|
|
def create_kb(c, base, tok, name):
|
|
r = c.post(f"{base}{API}/knowledge-bases", headers=H(tok),
|
|
json={"name": name, "mode": "text", "chunking_strategy": "standard"})
|
|
r.raise_for_status()
|
|
return r.json()["kb_id" if "kb_id" in r.json() else "id"]
|
|
|
|
|
|
def upload(c, base, tok, kb_id, path: Path):
|
|
with open(path, "rb") as f:
|
|
r = c.post(f"{base}{API}/knowledge-bases/{kb_id}/documents",
|
|
headers=H(tok), files={"file": (path.name, f, "application/octet-stream")},
|
|
timeout=600)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
|
|
|
|
def wait_indexed(c, base, tok, kb_id, document_id, timeout=600):
|
|
t0 = time.perf_counter()
|
|
while time.perf_counter() - t0 < timeout:
|
|
r = c.get(f"{base}{API}/knowledge-bases/{kb_id}/documents", headers=H(tok))
|
|
r.raise_for_status()
|
|
for d in r.json()["documents"]:
|
|
if d["id"] == document_id:
|
|
if d["status"] == "completed":
|
|
return time.perf_counter() - t0, d["num_chunks"], "completed"
|
|
if d["status"] == "failed":
|
|
return time.perf_counter() - t0, 0, "failed"
|
|
time.sleep(0.1)
|
|
return timeout, 0, "timeout"
|
|
|
|
|
|
def index_doc(c, base, tok, kb_id, path):
|
|
t0 = time.perf_counter()
|
|
up = upload(c, base, tok, kb_id, path)
|
|
if up.get("already_indexed"):
|
|
return {"file": path.name, "elapsed_s": 0.0, "chunks": 0, "status": "dup"}
|
|
elapsed, chunks, status = wait_indexed(c, base, tok, kb_id, up["document_id"])
|
|
total = time.perf_counter() - t0
|
|
return {"file": path.name, "elapsed_s": round(total, 3), "chunks": chunks,
|
|
"status": status, "document_id": up["document_id"]}
|
|
|
|
|
|
def search(c, base, tok, kb_id, query, mode, top_k=10):
|
|
r = c.post(f"{base}{API}/search", headers=H(tok),
|
|
json={"query": query, "kb_id": kb_id, "mode": mode, "top_k": top_k})
|
|
r.raise_for_status()
|
|
return r.json()["hits"]
|
|
|
|
|
|
def score_mode(c, base, tok, kb_id, mode):
|
|
r1 = r3 = r5 = 0
|
|
mrr = 0.0
|
|
lat = []
|
|
for query, fsub, phrases in GOLD:
|
|
t = time.perf_counter()
|
|
hits = search(c, base, tok, kb_id, query, mode, top_k=10)
|
|
lat.append((time.perf_counter() - t) * 1000)
|
|
rank = None
|
|
for i, h in enumerate(hits):
|
|
fn = (h.get("filename") or "").lower()
|
|
txt = (h.get("text") or "").lower()
|
|
if fsub in fn and any(p in txt for p in phrases):
|
|
rank = i
|
|
break
|
|
if rank is not None:
|
|
if rank == 0:
|
|
r1 += 1
|
|
if rank < 3:
|
|
r3 += 1
|
|
if rank < 5:
|
|
r5 += 1
|
|
mrr += 1.0 / (rank + 1)
|
|
n = len(GOLD)
|
|
lat.sort()
|
|
return {"recall@1": round(r1 / n, 3), "recall@3": round(r3 / n, 3),
|
|
"recall@5": round(r5 / n, 3), "mrr": round(mrr / n, 3),
|
|
"search_ms_median": round(lat[len(lat) // 2], 2)}
|
|
|
|
|
|
def make_synthetic(dirpath: Path, n: int):
|
|
dirpath.mkdir(parents=True, exist_ok=True)
|
|
paths = []
|
|
for i in range(n):
|
|
p = dirpath / f"syn_{i:02d}.txt"
|
|
body = (f"Synthetic document number {i}. Project codename Orbit-{i} concerns "
|
|
f"widget {i} calibration at {100 + i} hertz. Unique token zglyph{i} marks "
|
|
f"this file. " * 8)
|
|
p.write_text(body)
|
|
paths.append(p)
|
|
return paths
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--base", required=True)
|
|
ap.add_argument("--label", required=True)
|
|
ap.add_argument("--password", required=True)
|
|
ap.add_argument("--username", default="unsloth")
|
|
ap.add_argument("--corpus", required=True)
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--scaling-n", type=int, default=8)
|
|
args = ap.parse_args()
|
|
|
|
res = {"label": args.label, "base": args.base}
|
|
with httpx.Client(timeout=120) as c:
|
|
tok = login(c, args.base, args.username, args.password)
|
|
res["warmup_s"] = round(warmup(c, args.base, tok), 2)
|
|
|
|
# --- Real corpus: index timing (first = cold, rest = warm) + accuracy ---
|
|
corpus = sorted(Path(args.corpus).glob("*"))
|
|
corpus = [p for p in corpus if p.suffix.lower() in
|
|
(".pdf", ".txt", ".md", ".html", ".htm", ".docx")]
|
|
kb = create_kb(c, args.base, tok, f"{args.label}-corpus")
|
|
res["kb"] = kb
|
|
res["corpus_index"] = []
|
|
for i, p in enumerate(corpus):
|
|
r = index_doc(c, args.base, tok, kb, p)
|
|
r["which"] = "cold" if i == 0 else "warm"
|
|
res["corpus_index"].append(r)
|
|
print(f"[{args.label}] index {p.name}: {r['elapsed_s']}s ({r['chunks']} chunks) {r['which']}")
|
|
|
|
res["accuracy"] = {m: score_mode(c, args.base, tok, kb, m)
|
|
for m in ("bm25", "dense", "hybrid")}
|
|
for m, s in res["accuracy"].items():
|
|
print(f"[{args.label}] {m}: R@1={s['recall@1']} R@5={s['recall@5']} MRR={s['mrr']}")
|
|
|
|
# --- Scaling: N small docs into one fresh KB, per-doc index time ---
|
|
kb2 = create_kb(c, args.base, tok, f"{args.label}-scaling")
|
|
syn = make_synthetic(Path(args.corpus).parent / "rag_synthetic", args.scaling_n)
|
|
res["scaling"] = []
|
|
for i, p in enumerate(syn):
|
|
r = index_doc(c, args.base, tok, kb2, p)
|
|
res["scaling"].append({"n": i + 1, "elapsed_s": r["elapsed_s"]})
|
|
print(f"[{args.label}] scaling doc {i+1}/{len(syn)}: {r['elapsed_s']}s")
|
|
|
|
Path(args.out).write_text(json.dumps(res, indent=2))
|
|
print(f"[{args.label}] wrote {args.out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|