From 6e40e49eeab4cf18ba32e202a2e6040ee9484874 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 31 May 2026 13:41:45 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/bench/rag_studio_bench.py | 164 ++++++++++++++++++++----------- studio/bench/studio_rag_ui.py | 61 +++++++----- 2 files changed, 142 insertions(+), 83 deletions(-) diff --git a/studio/bench/rag_studio_bench.py b/studio/bench/rag_studio_bench.py index 37393d6aca..bd36ddebf8 100644 --- a/studio/bench/rag_studio_bench.py +++ b/studio/bench/rag_studio_bench.py @@ -6,6 +6,7 @@ Usage: python rag_studio_bench.py --base http://127.0.0.1:8901 --label baseline \ --password "" --corpus ../data/rag_corpus --out ../logs/bench_baseline.json """ + import argparse import json import time @@ -19,22 +20,42 @@ API = "/api/rag" # 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"]), + ( + "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"]), ] @@ -42,15 +63,18 @@ 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}) + 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 = 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 = 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 @@ -63,32 +87,38 @@ def H(tok): def warmup(c, base, tok): t = time.perf_counter() try: - c.post(f"{base}{API}/warmup", headers=H(tok), timeout=600) + 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 = 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 = 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): +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 = 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: @@ -107,13 +137,21 @@ def index_doc(c, base, tok, kb_id, path): 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"]} + 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}) +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"] @@ -124,7 +162,7 @@ def score_mode(c, base, tok, kb_id, mode): lat = [] for query, fsub, phrases in GOLD: t = time.perf_counter() - hits = search(c, base, tok, kb_id, query, mode, top_k=10) + 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): @@ -143,19 +181,25 @@ def score_mode(c, base, tok, kb_id, mode): 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)} + 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) + 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) + 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 @@ -163,24 +207,27 @@ def make_synthetic(dirpath: Path, n: int): 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) + 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: + 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")] + 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"] = [] @@ -188,12 +235,17 @@ def main(): 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']}") + 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")} + 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']}") + 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") @@ -204,7 +256,7 @@ def main(): 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)) + Path(args.out).write_text(json.dumps(res, indent = 2)) print(f"[{args.label}] wrote {args.out}") diff --git a/studio/bench/studio_rag_ui.py b/studio/bench/studio_rag_ui.py index f2a16fd0fa..c81c86c09c 100644 --- a/studio/bench/studio_rag_ui.py +++ b/studio/bench/studio_rag_ui.py @@ -8,6 +8,7 @@ Usage: --password "" --doc ../data/rag_corpus/bert_1810.04805.pdf \ --out ../outputs/ui_baseline """ + import argparse import asyncio import json @@ -20,17 +21,21 @@ from playwright.async_api import async_playwright def get_token(base, user, pw): new = pw + "Aa1!" - with httpx.Client(timeout=30) as c: - r = c.post(f"{base}/api/auth/login", json={"username": user, "password": pw}) + with httpx.Client(timeout = 30) as c: + 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 = 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={"Authorization": f"Bearer {tok}"}, - json={"current_password": pw, "new_password": new}) + r2 = c.post( + f"{base}/api/auth/change-password", + headers = {"Authorization": f"Bearer {tok}"}, + json = {"current_password": pw, "new_password": new}, + ) r2.raise_for_status() tok = r2.json()["access_token"] refresh = r2.json().get("refresh_token", "") @@ -47,41 +52,43 @@ def init_script(tok, refresh): async def run(args): out = Path(args.out) - (out / "video").mkdir(parents=True, exist_ok=True) + (out / "video").mkdir(parents = True, exist_ok = True) tok, refresh = get_token(args.base, args.username, args.password) result = {"label": args.label, "doc": Path(args.doc).name} async with async_playwright() as p: - browser = await p.chromium.launch(headless=True) + browser = await p.chromium.launch(headless = True) ctx = await browser.new_context( - viewport={"width": 1440, "height": 900}, - record_video_dir=str(out / "video"), - record_video_size={"width": 1440, "height": 900}, + viewport = {"width": 1440, "height": 900}, + record_video_dir = str(out / "video"), + record_video_size = {"width": 1440, "height": 900}, ) await ctx.add_init_script(init_script(tok, refresh)) page = await ctx.new_page() - await page.goto(f"{args.base}/chat", wait_until="domcontentloaded") - await page.locator("form:has(textarea) textarea").first.wait_for(state="visible", timeout=30000) - await page.screenshot(path=str(out / "01_chat.png")) + await page.goto(f"{args.base}/chat", wait_until = "domcontentloaded") + await page.locator("form:has(textarea) textarea").first.wait_for( + state = "visible", timeout = 30000 + ) + await page.screenshot(path = str(out / "01_chat.png")) # Enable RAG so the document attach control renders. enable = page.locator('button[aria-label="Enable RAG"]').first try: - await enable.click(timeout=8000) + await enable.click(timeout = 8000) except Exception: pass # already enabled await page.wait_for_timeout(500) - await page.screenshot(path=str(out / "02_rag_on.png")) + await page.screenshot(path = str(out / "02_rag_on.png")) # Upload via the hidden file input (native picker can't be driven). file_input = page.locator('input[type="file"][accept*=".pdf"]').first - await file_input.wait_for(state="attached", timeout=8000) + await file_input.wait_for(state = "attached", timeout = 8000) t0 = time.perf_counter() await file_input.set_input_files(args.doc) # Wait for the global "RAG index ready" toast (fallback: chip "Ready"). - ready_toast = page.get_by_text("RAG index ready", exact=False).first - chip_ready = page.get_by_text("Ready", exact=True).first + ready_toast = page.get_by_text("RAG index ready", exact = False).first + chip_ready = page.get_by_text("Ready", exact = True).first indexed = False deadline = time.perf_counter() + 180 while time.perf_counter() < deadline: @@ -101,7 +108,7 @@ async def run(args): elapsed = time.perf_counter() - t0 result["index_seconds"] = round(elapsed, 2) result["indexed"] = indexed - await page.screenshot(path=str(out / "03_indexed.png"), full_page=True) + await page.screenshot(path = str(out / "03_indexed.png"), full_page = True) print(f"[{args.label}] UI upload->ready: {elapsed:.2f}s indexed={indexed}") await ctx.close() @@ -111,18 +118,18 @@ async def run(args): webms = sorted((out / "video").glob("*.webm")) if webms: webms[-1].rename(out / "video" / f"{args.label}.webm") - Path(out / "result.json").write_text(json.dumps(result, indent=2)) + Path(out / "result.json").write_text(json.dumps(result, indent = 2)) print(f"[{args.label}] wrote {out/'result.json'}") 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("--doc", required=True) - ap.add_argument("--out", required=True) + 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("--doc", required = True) + ap.add_argument("--out", required = True) asyncio.run(run(ap.parse_args()))