Spaces:
Sleeping
Sleeping
| """Call the Space's /trace API and render every result with the UI's own code. | |
| Usage: | |
| python scripts/smoke_test.py --url http://127.0.0.1:7860 | |
| python scripts/smoke_test.py --space wang2226/steering-showcase [--models qwen25-05b,smollm2-135m] [--custom] | |
| Every call that is not a cache hit spends GPU time on a Space; logged-in calls use the account's quota. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) | |
| os.environ["SD_SKIP_INIT"] = "1" # import the renderers without loading models | |
| from content import examples as EX # noqa: E402 | |
| from steering.core import ALPHAS # noqa: E402 | |
| from steering.scenarios import SCENARIOS # noqa: E402 | |
| from viz import render as R # noqa: E402 | |
| def check(trace: dict, model: str, scenario: str) -> list[str]: | |
| problems = [] | |
| if trace.get("schema") != "steer.trace/v1": | |
| problems.append("bad schema") | |
| if trace.get("source") != "live": | |
| problems.append("not a live trace") | |
| if [s["alpha"] for s in trace["states"]] != ALPHAS: | |
| problems.append("states off the alpha grid") | |
| if trace.get("direction") != scenario: | |
| problems.append(f"direction {trace.get('direction')!r}") | |
| for s in trace["states"]: | |
| if not s["continuation"].startswith(" ".join(trace["candidates"][s["selected_index"]].split())): | |
| problems.append(f"alpha {s['alpha']}: response does not start with the highlighted token") | |
| if abs(sum(s["percents"]) + s["other"] - 100) > 0.01: | |
| problems.append(f"alpha {s['alpha']}: probabilities do not sum to 100") | |
| if not trace.get("depth") or not trace.get("negative_examples"): | |
| problems.append("depth or contrast continuations missing") | |
| for alpha in ALPHAS: # the same renderers the page uses | |
| R.formula_md(trace, alpha), R.cards_html(trace, alpha), R.response_html(trace, alpha) | |
| R.chart_html(trace, alpha), R.sweep_html(trace, alpha), R.schematic_html(trace, alpha) | |
| R.takeaway_html(trace), R.contrast_html(trace), R.explain_html(trace) | |
| return problems | |
| def main() -> int: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--url") | |
| ap.add_argument("--space") | |
| ap.add_argument("--models", default="qwen25-05b") | |
| ap.add_argument("--custom", action="store_true", help="also run one visitor-style prompt per model") | |
| ap.add_argument("--fresh", action="store_true", help="bypass the Space's cache") | |
| args = ap.parse_args() | |
| from gradio_client import Client | |
| from huggingface_hub import get_token | |
| token = os.environ.get("HF_TOKEN") or get_token() | |
| client = Client(args.space, token=token) if args.space else Client(args.url or "http://127.0.0.1:7860") | |
| cases = [(s.id, "", "") for s in SCENARIOS] | |
| if args.custom: | |
| e = EX.EXAMPLES[-1] | |
| cases.append((e["scenario"], e["prompt"], e["prefix"])) | |
| failures = 0 | |
| for model in args.models.split(","): | |
| for scenario, prompt, prefix in cases: | |
| t0 = time.time() | |
| try: | |
| trace = client.predict(model, scenario, prompt, prefix, args.fresh, api_name="/trace") | |
| problems = check(trace, model, scenario) | |
| except Exception as exc: | |
| trace, problems = None, [f"{type(exc).__name__}: {str(exc)[:300]}"] | |
| failures += bool(problems) | |
| extra = "" | |
| if trace: | |
| env = trace.get("env", {}) | |
| winners = " → ".join(dict.fromkeys(R.label(trace["candidates"][s["selected_index"]]) for s in trace["states"])) | |
| extra = (f"{'cached' if trace.get('cached') else 'fresh'} {env.get('device')} {env.get('gpu_seconds')}s" | |
| f" | {winners}") | |
| print(f"{'OK ' if not problems else 'FAIL'} {model:13} {scenario:18} {'custom' if prompt else ' '} " | |
| f"{time.time() - t0:5.1f}s {extra}") | |
| for p in problems: | |
| print(f" ! {p}") | |
| print(f"\n{failures} failure(s)") | |
| return 1 if failures else 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |