[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
4f20be0d2f
commit
479e7d0303
2 changed files with 121 additions and 58 deletions
|
|
@ -3,6 +3,7 @@
|
|||
Asks gold questions grounded in an indexed KB, with rag_scope + tools enabled,
|
||||
and reports whether the model called search_knowledge_base and answered correctly.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
|
|
@ -12,21 +13,29 @@ import httpx
|
|||
QUESTIONS = [
|
||||
("What are BERT's two pre-training objectives?", ["masked", "next sentence"]),
|
||||
("How many attention heads does the base Transformer use?", ["8", "eight"]),
|
||||
("What retriever does the RAG paper use to fetch passages?", ["dpr", "dense passage", "bi-encoder"]),
|
||||
(
|
||||
"What retriever does the RAG paper use to fetch passages?",
|
||||
["dpr", "dense passage", "bi-encoder"],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def token(c, base, pw):
|
||||
new = pw + "Aa1!"
|
||||
r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": pw})
|
||||
r = c.post(f"{base}/api/auth/login", json = {"username": "unsloth", "password": pw})
|
||||
if r.status_code == 401:
|
||||
r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": new})
|
||||
r = c.post(
|
||||
f"{base}/api/auth/login", json = {"username": "unsloth", "password": new}
|
||||
)
|
||||
r.raise_for_status()
|
||||
b = r.json()
|
||||
t = b["access_token"]
|
||||
if b.get("must_change_password"):
|
||||
r2 = c.post(f"{base}/api/auth/change-password", headers={"Authorization": f"Bearer {t}"},
|
||||
json={"current_password": pw, "new_password": new})
|
||||
r2 = c.post(
|
||||
f"{base}/api/auth/change-password",
|
||||
headers = {"Authorization": f"Bearer {t}"},
|
||||
json = {"current_password": pw, "new_password": new},
|
||||
)
|
||||
r2.raise_for_status()
|
||||
t = r2.json()["access_token"]
|
||||
return t
|
||||
|
|
@ -34,18 +43,23 @@ def token(c, base, pw):
|
|||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--base", required=True)
|
||||
ap.add_argument("--password", required=True)
|
||||
ap.add_argument("--kb-name-contains", default="corpus")
|
||||
ap.add_argument("--out", default=None)
|
||||
ap.add_argument("--base", required = True)
|
||||
ap.add_argument("--password", required = True)
|
||||
ap.add_argument("--kb-name-contains", default = "corpus")
|
||||
ap.add_argument("--out", default = None)
|
||||
args = ap.parse_args()
|
||||
|
||||
results = {"base": args.base, "qa": []}
|
||||
with httpx.Client(timeout=300) as c:
|
||||
with httpx.Client(timeout = 300) as c:
|
||||
tok = token(c, args.base, args.password)
|
||||
H = {"Authorization": f"Bearer {tok}"}
|
||||
kbs = c.get(f"{args.base}/api/rag/knowledge-bases", headers=H).json()["knowledge_bases"]
|
||||
kb = next((k for k in kbs if args.kb_name_contains in k["name"]), kbs[0] if kbs else None)
|
||||
kbs = c.get(f"{args.base}/api/rag/knowledge-bases", headers = H).json()[
|
||||
"knowledge_bases"
|
||||
]
|
||||
kb = next(
|
||||
(k for k in kbs if args.kb_name_contains in k["name"]),
|
||||
kbs[0] if kbs else None,
|
||||
)
|
||||
if not kb:
|
||||
sys.exit("no KB found")
|
||||
kb_id = kb["id"]
|
||||
|
|
@ -54,32 +68,48 @@ def main():
|
|||
for q, phrases in QUESTIONS:
|
||||
payload = {
|
||||
"model": "local",
|
||||
"messages": [{"role": "user", "content": q +
|
||||
" Use the knowledge base; cite the source."}],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": q + " Use the knowledge base; cite the source.",
|
||||
}
|
||||
],
|
||||
"stream": False,
|
||||
"enable_tools": True,
|
||||
"enable_thinking": False,
|
||||
"max_tokens": 400,
|
||||
"temperature": 0.3,
|
||||
"rag_scope": {"kb_id": kb_id, "default_top_k": 5, "min_score": 0.0, "mode": "hybrid"},
|
||||
"rag_scope": {
|
||||
"kb_id": kb_id,
|
||||
"default_top_k": 5,
|
||||
"min_score": 0.0,
|
||||
"mode": "hybrid",
|
||||
},
|
||||
}
|
||||
# The endpoint streams SSE regardless of stream flag; parse it.
|
||||
answer, called_tool, tool_query, src = "", False, "", ""
|
||||
status = 0
|
||||
with c.stream("POST", f"{args.base}/api/inference/chat/completions",
|
||||
headers=H, json=payload) as r:
|
||||
with c.stream(
|
||||
"POST",
|
||||
f"{args.base}/api/inference/chat/completions",
|
||||
headers = H,
|
||||
json = payload,
|
||||
) as r:
|
||||
status = r.status_code
|
||||
for line in r.iter_lines():
|
||||
if not line or not line.startswith("data:"):
|
||||
continue
|
||||
data = line[len("data:"):].strip()
|
||||
data = line[len("data:") :].strip()
|
||||
if data == "[DONE]":
|
||||
break
|
||||
try:
|
||||
ev = json.loads(data)
|
||||
except Exception:
|
||||
continue
|
||||
if ev.get("type") == "tool_start" and ev.get("tool_name") == "search_knowledge_base":
|
||||
if (
|
||||
ev.get("type") == "tool_start"
|
||||
and ev.get("tool_name") == "search_knowledge_base"
|
||||
):
|
||||
called_tool = True
|
||||
tool_query = (ev.get("arguments") or {}).get("query", "")
|
||||
elif ev.get("type") == "tool_end":
|
||||
|
|
@ -92,20 +122,32 @@ def main():
|
|||
if delta.get("content"):
|
||||
answer += delta["content"]
|
||||
hit = any(p in answer.lower() for p in phrases)
|
||||
print(f"\nQ: {q}\n status={status} tool_called={called_tool} src={src} answer_has_fact={hit}")
|
||||
print(
|
||||
f"\nQ: {q}\n status={status} tool_called={called_tool} src={src} answer_has_fact={hit}"
|
||||
)
|
||||
print(f" tool_query={tool_query!r}")
|
||||
print(f" A: {answer[:300].replace(chr(10),' ')}")
|
||||
results["qa"].append({"q": q, "status": status, "tool_called": called_tool,
|
||||
"retrieved_source": src, "tool_query": tool_query,
|
||||
"answer_has_fact": hit, "answer": answer[:600]})
|
||||
results["qa"].append(
|
||||
{
|
||||
"q": q,
|
||||
"status": status,
|
||||
"tool_called": called_tool,
|
||||
"retrieved_source": src,
|
||||
"tool_query": tool_query,
|
||||
"answer_has_fact": hit,
|
||||
"answer": answer[:600],
|
||||
}
|
||||
)
|
||||
|
||||
if args.out:
|
||||
with open(args.out, "w") as f:
|
||||
json.dump(results, f, indent=2)
|
||||
json.dump(results, f, indent = 2)
|
||||
print(f"\nwrote {args.out}")
|
||||
n = len(results["qa"])
|
||||
print(f"\nSUMMARY: {sum(r['answer_has_fact'] for r in results['qa'])}/{n} answers contain the fact, "
|
||||
f"{sum(r['tool_called'] for r in results['qa'])}/{n} called the RAG tool")
|
||||
print(
|
||||
f"\nSUMMARY: {sum(r['answer_has_fact'] for r in results['qa'])}/{n} answers contain the fact, "
|
||||
f"{sum(r['tool_called'] for r in results['qa'])}/{n} called the RAG tool"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Usage:
|
|||
--password StudioBench2026! --doc ../data/rag_corpus/bert_1810.04805.pdf \
|
||||
--question "What are BERT's two pre-training objectives?" --out ../outputs/composer_improved
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
|
|
@ -20,16 +21,23 @@ from playwright.async_api import async_playwright
|
|||
|
||||
def get_token(base, pw):
|
||||
new = pw + "Aa1!"
|
||||
with httpx.Client(timeout=30) as c:
|
||||
r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": pw})
|
||||
with httpx.Client(timeout = 30) as c:
|
||||
r = c.post(
|
||||
f"{base}/api/auth/login", json = {"username": "unsloth", "password": pw}
|
||||
)
|
||||
if r.status_code == 401:
|
||||
r = c.post(f"{base}/api/auth/login", json={"username": "unsloth", "password": new})
|
||||
r = c.post(
|
||||
f"{base}/api/auth/login", json = {"username": "unsloth", "password": new}
|
||||
)
|
||||
r.raise_for_status()
|
||||
b = r.json()
|
||||
t = b["access_token"]
|
||||
if b.get("must_change_password"):
|
||||
r2 = c.post(f"{base}/api/auth/change-password", headers={"Authorization": f"Bearer {t}"},
|
||||
json={"current_password": pw, "new_password": new})
|
||||
r2 = c.post(
|
||||
f"{base}/api/auth/change-password",
|
||||
headers = {"Authorization": f"Bearer {t}"},
|
||||
json = {"current_password": pw, "new_password": new},
|
||||
)
|
||||
r2.raise_for_status()
|
||||
t, b = r2.json()["access_token"], r2.json()
|
||||
return t, b.get("refresh_token", "")
|
||||
|
|
@ -43,21 +51,29 @@ 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.password)
|
||||
result = {"label": args.label, "doc": Path(args.doc).name, "question": args.question}
|
||||
result = {
|
||||
"label": args.label,
|
||||
"doc": Path(args.doc).name,
|
||||
"question": args.question,
|
||||
}
|
||||
|
||||
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})
|
||||
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.goto(f"{args.base}/chat", wait_until = "domcontentloaded")
|
||||
await page.locator("form:has(textarea) textarea").first.wait_for(
|
||||
state = "visible", timeout = 30000
|
||||
)
|
||||
await page.wait_for_timeout(1500)
|
||||
await page.screenshot(path=str(out / "01_chat.png"))
|
||||
await page.screenshot(path = str(out / "01_chat.png"))
|
||||
|
||||
# Enable the RAG toggle (the pill is active once a model is loaded).
|
||||
for lbl in ("Enable RAG", "Disable RAG"):
|
||||
|
|
@ -67,15 +83,15 @@ async def run(args):
|
|||
await btn.click()
|
||||
break
|
||||
await page.wait_for_timeout(700)
|
||||
await page.screenshot(path=str(out / "02_rag_on.png"))
|
||||
await page.screenshot(path = str(out / "02_rag_on.png"))
|
||||
|
||||
# Upload a document through the composer attach control.
|
||||
file_input = page.locator('input[type="file"][accept*=".pdf"]').first
|
||||
await file_input.wait_for(state="attached", timeout=10000)
|
||||
await file_input.wait_for(state = "attached", timeout = 10000)
|
||||
t0 = time.perf_counter()
|
||||
await file_input.set_input_files(args.doc)
|
||||
ready = page.get_by_text("Ready", exact=True).first
|
||||
toast = page.get_by_text("RAG index ready", exact=False).first
|
||||
ready = page.get_by_text("Ready", exact = True).first
|
||||
toast = page.get_by_text("RAG index ready", exact = False).first
|
||||
indexed = False
|
||||
while time.perf_counter() - t0 < 180:
|
||||
for sig in (toast, ready):
|
||||
|
|
@ -90,30 +106,35 @@ async def run(args):
|
|||
await page.wait_for_timeout(250)
|
||||
result["index_seconds"] = round(time.perf_counter() - t0, 2)
|
||||
result["indexed"] = indexed
|
||||
await page.screenshot(path=str(out / "03_indexed.png"))
|
||||
await page.screenshot(path = str(out / "03_indexed.png"))
|
||||
|
||||
# Ask a grounded question.
|
||||
box = page.locator("form:has(textarea) textarea").first
|
||||
await box.click()
|
||||
await box.fill(args.question)
|
||||
await box.press("Enter")
|
||||
stop = page.locator('button[aria-label="Stop generating"], button:has-text("Stop")').first
|
||||
stop = page.locator(
|
||||
'button[aria-label="Stop generating"], button:has-text("Stop")'
|
||||
).first
|
||||
try:
|
||||
await stop.wait_for(state="visible", timeout=30000)
|
||||
await stop.wait_for(state = "visible", timeout = 30000)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await stop.wait_for(state="hidden", timeout=180000)
|
||||
await stop.wait_for(state = "hidden", timeout = 180000)
|
||||
except Exception:
|
||||
pass
|
||||
await page.wait_for_timeout(1000)
|
||||
await page.screenshot(path=str(out / "04_answer.png"), full_page=True)
|
||||
await page.screenshot(path = str(out / "04_answer.png"), full_page = True)
|
||||
|
||||
answer = await page.evaluate(
|
||||
"""() => Array.from(document.querySelectorAll('[data-role="assistant"], [data-message-role="assistant"]'))
|
||||
.map(n => n.textContent || '').join('\\n')""")
|
||||
.map(n => n.textContent || '').join('\\n')"""
|
||||
)
|
||||
result["answer_excerpt"] = (answer or "")[-600:]
|
||||
print(f"[{args.label}] indexed={indexed} in {result['index_seconds']}s; answer chars={len(answer or '')}")
|
||||
print(
|
||||
f"[{args.label}] indexed={indexed} in {result['index_seconds']}s; answer chars={len(answer or '')}"
|
||||
)
|
||||
|
||||
await ctx.close()
|
||||
await browser.close()
|
||||
|
|
@ -121,18 +142,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("--doc", required=True)
|
||||
ap.add_argument("--question", 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("--doc", required = True)
|
||||
ap.add_argument("--question", 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