chore(scripts): add neighbor expansion benchmark sweep
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
374
scripts/run_neighbor_sweep.py
Normal file
374
scripts/run_neighbor_sweep.py
Normal file
@@ -0,0 +1,374 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run a Benchmark Sweep: fixed_size × text-embedding-3-large × Neighbor ±0..±3.
|
||||
|
||||
Calls the running FastAPI app over HTTP (same path as the Dashboard).
|
||||
Doc-major order; gap-fills Process (fixed_size only) when SQLite provenance
|
||||
says the Model Corpus is not ready; retries each unit twice on failure.
|
||||
|
||||
Usage:
|
||||
.venv/bin/python scripts/run_neighbor_sweep.py
|
||||
.venv/bin/python scripts/run_neighbor_sweep.py --base-url http://127.0.0.1:8000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
CORPUS_MODEL_ID = "text-embedding-3-large"
|
||||
STRATEGY = "fixed_size"
|
||||
TOP_K = 5
|
||||
NEIGHBOR_LEVELS = (0, 1, 2, 3)
|
||||
MAX_ATTEMPTS = 3 # 1 try + 2 retries
|
||||
DEFAULT_BASE_URL = "http://127.0.0.1:8000"
|
||||
# Benchmarks with many questions can take a long time (LLM + eval per Q).
|
||||
REQUEST_TIMEOUT_S = 60 * 60 * 3 # 3 hours
|
||||
|
||||
# App-truth document ↔ questions map (filenames as stored in the app).
|
||||
SWEEP_PAIRS: list[tuple[str, str]] = [
|
||||
("bazresi.docx", "files/bazresi.json"),
|
||||
("customer1.docx", "files/customer1.json"),
|
||||
("fire.docx", "files/fire.json"),
|
||||
("general-havades-individuals.doc", "files/general-havades-individuals.json"),
|
||||
("havades.docx", "files/havades.json"),
|
||||
("life-time-individual.docx", "files/life-time-individual.json"),
|
||||
("moavenin.docx", "files/moavenin.json"),
|
||||
("Refah.docx", "files/refah.json"),
|
||||
("website.docx", "files/website.json"),
|
||||
("lifetime-compensation.docx", "files/lifetime-compensation.json"),
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnitResult:
|
||||
document: str
|
||||
neighbor: int
|
||||
kind: str # "process" | "experiment"
|
||||
ok: bool
|
||||
detail: str
|
||||
experiment_id: str | None = None
|
||||
attempts: int = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class SweepState:
|
||||
results: list[UnitResult] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def failed(self) -> list[UnitResult]:
|
||||
return [r for r in self.results if not r.ok]
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
print(msg, flush=True)
|
||||
|
||||
|
||||
def is_ready(doc: dict[str, Any]) -> bool:
|
||||
corpus = doc.get("last_corpus_embedding_model_id")
|
||||
counts = doc.get("chunk_counts") or {}
|
||||
fixed = int(counts.get("fixed_size") or 0)
|
||||
return corpus == CORPUS_MODEL_ID and fixed > 0
|
||||
|
||||
|
||||
def fetch_documents(client: httpx.Client) -> dict[str, dict[str, Any]]:
|
||||
"""Return filename → document dict (paginated)."""
|
||||
by_name: dict[str, dict[str, Any]] = {}
|
||||
offset = 0
|
||||
limit = 100
|
||||
while True:
|
||||
r = client.get("/documents", params={"offset": offset, "limit": limit})
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
items = data.get("items") or []
|
||||
for doc in items:
|
||||
by_name[doc["filename"]] = doc
|
||||
total = int(data.get("total") or 0)
|
||||
offset += len(items)
|
||||
if offset >= total or not items:
|
||||
break
|
||||
return by_name
|
||||
|
||||
|
||||
def error_detail(exc: BaseException) -> str:
|
||||
if isinstance(exc, httpx.HTTPStatusError):
|
||||
body = ""
|
||||
try:
|
||||
body = exc.response.text[:500]
|
||||
except Exception:
|
||||
pass
|
||||
return f"HTTP {exc.response.status_code}: {body or exc}"
|
||||
return str(exc)
|
||||
|
||||
|
||||
def with_retries(
|
||||
label: str,
|
||||
fn,
|
||||
*,
|
||||
state: SweepState,
|
||||
document: str,
|
||||
neighbor: int,
|
||||
kind: str,
|
||||
) -> Any | None:
|
||||
last_err = ""
|
||||
for attempt in range(1, MAX_ATTEMPTS + 1):
|
||||
try:
|
||||
log(f" [{kind}] {label} (attempt {attempt}/{MAX_ATTEMPTS})")
|
||||
value = fn()
|
||||
return value
|
||||
except Exception as exc:
|
||||
last_err = error_detail(exc)
|
||||
log(f" ! failed attempt {attempt}/{MAX_ATTEMPTS}: {last_err}")
|
||||
if attempt < MAX_ATTEMPTS:
|
||||
time.sleep(min(5 * attempt, 15))
|
||||
state.results.append(
|
||||
UnitResult(
|
||||
document=document,
|
||||
neighbor=neighbor,
|
||||
kind=kind,
|
||||
ok=False,
|
||||
detail=last_err,
|
||||
attempts=MAX_ATTEMPTS,
|
||||
)
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def ensure_corpus(
|
||||
client: httpx.Client,
|
||||
doc: dict[str, Any],
|
||||
*,
|
||||
state: SweepState,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Process fixed_size under CORPUS_MODEL_ID if not ready. Returns updated doc or None."""
|
||||
filename = doc["filename"]
|
||||
if is_ready(doc):
|
||||
log(f" corpus ready ({CORPUS_MODEL_ID}, fixed_size={doc['chunk_counts'].get('fixed_size')})")
|
||||
return doc
|
||||
|
||||
log(
|
||||
f" corpus not ready "
|
||||
f"(last_corpus={doc.get('last_corpus_embedding_model_id')!r}, "
|
||||
f"fixed_size={ (doc.get('chunk_counts') or {}).get('fixed_size') }) — processing"
|
||||
)
|
||||
|
||||
def _process() -> dict[str, Any]:
|
||||
r = client.post(
|
||||
f"/documents/{doc['id']}/process",
|
||||
json={
|
||||
"strategies": [STRATEGY],
|
||||
"corpus_model_id": CORPUS_MODEL_ID,
|
||||
},
|
||||
)
|
||||
r.raise_for_status()
|
||||
body = r.json()
|
||||
failed = body.get("strategies_failed") or []
|
||||
if failed:
|
||||
raise RuntimeError(f"process strategies_failed: {json.dumps(failed)}")
|
||||
# Refresh document record
|
||||
docs = fetch_documents(client)
|
||||
updated = docs.get(filename)
|
||||
if updated is None:
|
||||
raise RuntimeError(f"document disappeared after process: {filename}")
|
||||
if not is_ready(updated):
|
||||
raise RuntimeError(
|
||||
"process finished but provenance still not ready: "
|
||||
f"last_corpus={updated.get('last_corpus_embedding_model_id')!r}, "
|
||||
f"chunk_counts={updated.get('chunk_counts')}"
|
||||
)
|
||||
return updated
|
||||
|
||||
result = with_retries(
|
||||
f"process {filename}",
|
||||
_process,
|
||||
state=state,
|
||||
document=filename,
|
||||
neighbor=-1,
|
||||
kind="process",
|
||||
)
|
||||
if result is None:
|
||||
return None
|
||||
state.results.append(
|
||||
UnitResult(
|
||||
document=filename,
|
||||
neighbor=-1,
|
||||
kind="process",
|
||||
ok=True,
|
||||
detail="gap-filled fixed_size",
|
||||
attempts=1,
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def run_experiment(
|
||||
client: httpx.Client,
|
||||
*,
|
||||
doc: dict[str, Any],
|
||||
questions_file: str,
|
||||
neighbor: int,
|
||||
state: SweepState,
|
||||
) -> None:
|
||||
filename = doc["filename"]
|
||||
|
||||
def _bench() -> dict[str, Any]:
|
||||
r = client.post(
|
||||
"/benchmarks",
|
||||
json={
|
||||
"document_id": doc["id"],
|
||||
"strategies": [STRATEGY],
|
||||
"questions_file": questions_file,
|
||||
"top_k": TOP_K,
|
||||
"neighbor_prev": neighbor,
|
||||
"neighbor_next": neighbor,
|
||||
"corpus_model_id": CORPUS_MODEL_ID,
|
||||
},
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
body = with_retries(
|
||||
f"±{neighbor}/{neighbor} on {filename}",
|
||||
_bench,
|
||||
state=state,
|
||||
document=filename,
|
||||
neighbor=neighbor,
|
||||
kind="experiment",
|
||||
)
|
||||
if body is None:
|
||||
return
|
||||
|
||||
exp_id = body.get("experiment_id", "")
|
||||
state.results.append(
|
||||
UnitResult(
|
||||
document=filename,
|
||||
neighbor=neighbor,
|
||||
kind="experiment",
|
||||
ok=True,
|
||||
detail=(
|
||||
f"best={body.get('best_strategy')} "
|
||||
f"latency={body.get('total_latency_seconds')}s "
|
||||
f"cost=${body.get('estimated_cost_usd')}"
|
||||
),
|
||||
experiment_id=exp_id,
|
||||
)
|
||||
)
|
||||
log(f" ok experiment_id={exp_id}")
|
||||
|
||||
|
||||
def print_summary(state: SweepState) -> None:
|
||||
log("")
|
||||
log("=" * 88)
|
||||
log("BENCHMARK SWEEP SUMMARY")
|
||||
log("=" * 88)
|
||||
log(
|
||||
f"{'document':<40} {'±N':>4} {'kind':<11} {'status':<6} detail"
|
||||
)
|
||||
log("-" * 88)
|
||||
for r in state.results:
|
||||
n = "—" if r.neighbor < 0 else f"±{r.neighbor}"
|
||||
status = "ok" if r.ok else "FAIL"
|
||||
detail = r.experiment_id or r.detail
|
||||
if not r.ok and r.experiment_id is None:
|
||||
detail = r.detail
|
||||
elif r.ok and r.experiment_id:
|
||||
detail = f"id={r.experiment_id} {r.detail}"
|
||||
log(f"{r.document:<40} {n:>4} {r.kind:<11} {status:<6} {detail}")
|
||||
|
||||
failed = state.failed
|
||||
log("-" * 88)
|
||||
exp_ok = sum(1 for r in state.results if r.kind == "experiment" and r.ok)
|
||||
exp_fail = sum(1 for r in state.results if r.kind == "experiment" and not r.ok)
|
||||
log(f"Experiments: {exp_ok} ok, {exp_fail} failed")
|
||||
|
||||
if failed:
|
||||
log("")
|
||||
log("Manual redo hints (failed units):")
|
||||
for r in failed:
|
||||
if r.kind == "process":
|
||||
log(
|
||||
f" • Process {r.document}: "
|
||||
f'POST /documents/{{id}}/process '
|
||||
f'{{"strategies":["{STRATEGY}"],"corpus_model_id":"{CORPUS_MODEL_ID}"}}'
|
||||
)
|
||||
else:
|
||||
log(
|
||||
f" • Experiment {r.document} ±{r.neighbor}: "
|
||||
f'POST /benchmarks with strategies=["{STRATEGY}"], '
|
||||
f"neighbor_prev={r.neighbor}, neighbor_next={r.neighbor}, "
|
||||
f"corpus_model_id={CORPUS_MODEL_ID}, questions_file from SWEEP_PAIRS"
|
||||
)
|
||||
log(f" error: {r.detail}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="Run fixed_size Neighbor Expansion Benchmark Sweep")
|
||||
parser.add_argument(
|
||||
"--base-url",
|
||||
default=DEFAULT_BASE_URL,
|
||||
help=f"FastAPI base URL (default: {DEFAULT_BASE_URL})",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
base_url = args.base_url.rstrip("/")
|
||||
|
||||
state = SweepState()
|
||||
log(f"Benchmark Sweep → {base_url}")
|
||||
log(f"Strategy={STRATEGY} corpus={CORPUS_MODEL_ID} top_k={TOP_K} neighbors={list(NEIGHBOR_LEVELS)}")
|
||||
log(f"Documents: {len(SWEEP_PAIRS)}")
|
||||
|
||||
timeout = httpx.Timeout(REQUEST_TIMEOUT_S, connect=30.0)
|
||||
with httpx.Client(base_url=base_url, timeout=timeout) as client:
|
||||
try:
|
||||
health = client.get("/admin/health")
|
||||
health.raise_for_status()
|
||||
except Exception as exc:
|
||||
log(f"Cannot reach app at {base_url}: {error_detail(exc)}")
|
||||
return 2
|
||||
|
||||
docs = fetch_documents(client)
|
||||
log(f"Loaded {len(docs)} documents from app")
|
||||
|
||||
for filename, questions_file in SWEEP_PAIRS:
|
||||
log("")
|
||||
log(f"=== {filename} ({questions_file}) ===")
|
||||
doc = docs.get(filename)
|
||||
if doc is None:
|
||||
log(f" ! document not found in app — skipping")
|
||||
state.results.append(
|
||||
UnitResult(
|
||||
document=filename,
|
||||
neighbor=-1,
|
||||
kind="process",
|
||||
ok=False,
|
||||
detail="document not found in GET /documents",
|
||||
attempts=1,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
updated = ensure_corpus(client, doc, state=state)
|
||||
if updated is None:
|
||||
log(f" ! skipping experiments for {filename} (process failed)")
|
||||
continue
|
||||
docs[filename] = updated
|
||||
|
||||
for n in NEIGHBOR_LEVELS:
|
||||
run_experiment(
|
||||
client,
|
||||
doc=updated,
|
||||
questions_file=questions_file,
|
||||
neighbor=n,
|
||||
state=state,
|
||||
)
|
||||
|
||||
print_summary(state)
|
||||
return 1 if state.failed else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user