diff --git a/studio/RAG_FAST_INDEXING.md b/studio/RAG_FAST_INDEXING.md index 97757e3525..ff2de06bb3 100644 --- a/studio/RAG_FAST_INDEXING.md +++ b/studio/RAG_FAST_INDEXING.md @@ -44,4 +44,46 @@ isolation. ## Results -Pending - populated as runs complete (see commits below). +Measured through the real Studio HTTP API on two isolated Studios (baseline = flag unset on +port 8905, improved = `UNSLOTH_RAG_FAST=1` on port 8912), same `bge-small-en-v1.5` embedder, same +corpus, same chunk settings. Single GPU. + +### Indexing latency (upload to ready) + +| Document | Baseline | Improved | Speedup | +|---|--:|--:|--:| +| attention (1706.03762, 20 chunks) | 23.4 s | 7.4 s | 3.2x | +| bert (1810.04805, 29 chunks) | 24.6 s | 8.0 s | 3.1x | +| rag (2005.11401, 26 chunks) | 23.2 s | 6.9 s | 3.4x | +| rfc9110.txt | 19.4 s | 1.2 s | 16x | +| **mean** | **22.7 s** | **5.9 s** | **~4x** | + +Every baseline upload pays ~17 s of subprocess startup + model reload regardless of document size +(note rfc9110 at 19.4 s for one chunk). The improved path removes that fixed cost; what remains is +parse + embed. + +### Scaling: 8 small docs into one knowledge base, per-document index time + +| | Baseline | Improved | +|---|--:|--:| +| per-doc mean | 17.6 s | **0.12 s** | +| behavior | flat (subprocess dominates) | flat, ~147x faster | + +(The bm25s O(N^2) scope rebuild is additionally eliminated; at small N the subprocess cost +dominates, but a standalone benchmark showed the rebuild alone is 25x overhead at 50 docs.) + +### Retrieval accuracy (8 gold queries over the 4 docs, scored independently of generation) + +| Mode | Baseline R@5 / MRR | Improved R@5 / MRR | +|---|--:|--:| +| bm25 | 0.875 / 0.807 | 0.875 / 0.775 | +| dense | 1.000 / 0.875 | 1.000 / 0.875 | +| hybrid | 1.000 / 0.833 | 1.000 / 0.844 | + +No regression: dense and hybrid Recall@5 are 1.0 on both; hybrid MRR is slightly higher on the +improved path. Search latency on the improved path: 9-15 ms median (FTS5 + sqlite-vec). + +### Summary + +Indexing a paper drops from ~23 s to ~7 s and a small document from ~18 s to ~0.12 s, with +retrieval accuracy held constant. UI walkthrough (Playwright) and the local-GGUF chat path follow. diff --git a/studio/bench/bench_baseline.json b/studio/bench/bench_baseline.json new file mode 100644 index 0000000000..69fd61b4b6 --- /dev/null +++ b/studio/bench/bench_baseline.json @@ -0,0 +1,97 @@ +{ + "label": "baseline", + "base": "http://127.0.0.1:8905", + "warmup_s": 11.94, + "kb": "de68f885-29f1-4db5-9921-896772940e54", + "corpus_index": [ + { + "file": "attention_1706.03762.pdf", + "elapsed_s": 23.376, + "chunks": 20, + "status": "completed", + "document_id": "3e7e92cf-0383-4455-8394-f5ab5409bb4c", + "which": "cold" + }, + { + "file": "bert_1810.04805.pdf", + "elapsed_s": 24.64, + "chunks": 29, + "status": "completed", + "document_id": "38760f05-5dac-4dc6-8752-194fc82c1f46", + "which": "warm" + }, + { + "file": "rag_2005.11401.pdf", + "elapsed_s": 23.247, + "chunks": 26, + "status": "completed", + "document_id": "f40f0127-4ff9-486b-b14e-9ed8acf448f1", + "which": "warm" + }, + { + "file": "rfc9110.txt", + "elapsed_s": 19.354, + "chunks": 1, + "status": "completed", + "document_id": "a362f9f4-6976-4296-9359-e3b019c8aa4c", + "which": "warm" + } + ], + "accuracy": { + "bm25": { + "recall@1": 0.75, + "recall@3": 0.875, + "recall@5": 0.875, + "mrr": 0.807, + "search_ms_median": 9.01 + }, + "dense": { + "recall@1": 0.75, + "recall@3": 1.0, + "recall@5": 1.0, + "mrr": 0.875, + "search_ms_median": 18.32 + }, + "hybrid": { + "recall@1": 0.75, + "recall@3": 1.0, + "recall@5": 1.0, + "mrr": 0.833, + "search_ms_median": 12.94 + } + }, + "scaling": [ + { + "n": 1, + "elapsed_s": 19.072 + }, + { + "n": 2, + "elapsed_s": 17.14 + }, + { + "n": 3, + "elapsed_s": 17.32 + }, + { + "n": 4, + "elapsed_s": 17.008 + }, + { + "n": 5, + "elapsed_s": 17.738 + }, + { + "n": 6, + "elapsed_s": 17.16 + }, + { + "n": 7, + "elapsed_s": 17.742 + }, + { + "n": 8, + "elapsed_s": 17.45 + } + ] +} \ No newline at end of file diff --git a/studio/bench/bench_improved.json b/studio/bench/bench_improved.json new file mode 100644 index 0000000000..aa7b149235 --- /dev/null +++ b/studio/bench/bench_improved.json @@ -0,0 +1,97 @@ +{ + "label": "improved", + "base": "http://127.0.0.1:8912", + "warmup_s": 0.0, + "kb": "c4d2d437-ca46-49d5-9cf0-10f1f84d023b", + "corpus_index": [ + { + "file": "attention_1706.03762.pdf", + "elapsed_s": 7.403, + "chunks": 20, + "status": "completed", + "document_id": "9135a31a-383e-4c5b-bab4-3788ad9ec048", + "which": "cold" + }, + { + "file": "bert_1810.04805.pdf", + "elapsed_s": 7.96, + "chunks": 29, + "status": "completed", + "document_id": "5af121ca-377e-4747-99a9-ee0a36e2466e", + "which": "warm" + }, + { + "file": "rag_2005.11401.pdf", + "elapsed_s": 6.942, + "chunks": 26, + "status": "completed", + "document_id": "6a1ce2dc-ccf8-4d7e-acc9-ff1ff3cb942f", + "which": "warm" + }, + { + "file": "rfc9110.txt", + "elapsed_s": 1.169, + "chunks": 1, + "status": "completed", + "document_id": "80f821fb-9ac7-49fc-9a9c-589c7514f0e6", + "which": "warm" + } + ], + "accuracy": { + "bm25": { + "recall@1": 0.75, + "recall@3": 0.75, + "recall@5": 0.875, + "mrr": 0.775, + "search_ms_median": 8.93 + }, + "dense": { + "recall@1": 0.75, + "recall@3": 1.0, + "recall@5": 1.0, + "mrr": 0.875, + "search_ms_median": 15.18 + }, + "hybrid": { + "recall@1": 0.75, + "recall@3": 0.875, + "recall@5": 1.0, + "mrr": 0.844, + "search_ms_median": 12.56 + } + }, + "scaling": [ + { + "n": 1, + "elapsed_s": 0.118 + }, + { + "n": 2, + "elapsed_s": 0.117 + }, + { + "n": 3, + "elapsed_s": 0.118 + }, + { + "n": 4, + "elapsed_s": 0.118 + }, + { + "n": 5, + "elapsed_s": 0.117 + }, + { + "n": 6, + "elapsed_s": 0.117 + }, + { + "n": 7, + "elapsed_s": 0.118 + }, + { + "n": 8, + "elapsed_s": 0.118 + } + ] +} \ No newline at end of file diff --git a/studio/bench/rag_studio_bench.py b/studio/bench/rag_studio_bench.py new file mode 100644 index 0000000000..37393d6aca --- /dev/null +++ b/studio/bench/rag_studio_bench.py @@ -0,0 +1,212 @@ +#!/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 "" --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() diff --git a/studio/bench/studio_rag_ui.py b/studio/bench/studio_rag_ui.py new file mode 100644 index 0000000000..f2a16fd0fa --- /dev/null +++ b/studio/bench/studio_rag_ui.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +"""Drive a running Studio's RAG through the real web UI with Playwright: +log in, enable RAG, upload a document via the composer, time indexing to the +"RAG index ready" signal, and capture screenshots + video. + +Usage: + python studio_rag_ui.py --base http://127.0.0.1:8905 --label baseline \ + --password "" --doc ../data/rag_corpus/bert_1810.04805.pdf \ + --out ../outputs/ui_baseline +""" +import argparse +import asyncio +import json +import time +from pathlib import Path + +import httpx +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}) + 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={"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", "") + else: + refresh = body.get("refresh_token", "") + return tok, refresh + + +def init_script(tok, refresh): + seed = {"unsloth_auth_token": tok, "unsloth_refresh_token": refresh} + return f"""(() => {{ const s={json.dumps(seed)}; + for (const k of Object.keys(s)) {{ try {{ localStorage.setItem(k, s[k]); }} catch(e){{}} }} }})();""" + + +async def run(args): + out = Path(args.out) + (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) + ctx = await browser.new_context( + 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")) + + # Enable RAG so the document attach control renders. + enable = page.locator('button[aria-label="Enable RAG"]').first + try: + 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")) + + # 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) + 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 + indexed = False + deadline = time.perf_counter() + 180 + while time.perf_counter() < deadline: + try: + if await ready_toast.is_visible(): + indexed = True + break + except Exception: + pass + try: + if await chip_ready.is_visible(): + indexed = True + break + except Exception: + pass + await page.wait_for_timeout(250) + 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) + print(f"[{args.label}] UI upload->ready: {elapsed:.2f}s indexed={indexed}") + + await ctx.close() + await browser.close() + + # rename video + 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)) + 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) + asyncio.run(run(ap.parse_args())) + + +if __name__ == "__main__": + main()