[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
parent
4bed2ce28c
commit
6e40e49eea
2 changed files with 144 additions and 85 deletions
|
|
@ -6,6 +6,7 @@ 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
|
||||
|
|
@ -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}")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Usage:
|
|||
--password "<bootstrap-or-changed>" --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()))
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue