Spaces:
Running on Zero
Running on Zero
File size: 6,463 Bytes
371d90c 3e4fdce 371d90c 3e4fdce 371d90c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """Call every method through the /trace API and render each trace with the UI's own code.
Usage:
python scripts/smoke_test.py --url http://127.0.0.1:7860 --families qwen25,smollm2
python scripts/smoke_test.py --space USER/SPACE --families qwen3,llama32,gemma3,smollm2 --tokens 16
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
os.environ["DD_SKIP_INIT"] = "1" # import app.py for its renderers without loading models
from viz import plots # noqa: E402
CASES = [
("contrastive", {"method": "cd", "prompt": "Answer in one word: what is the capital of Australia?", "alpha": 1.0,
"beta": 0.1}),
("contrastive", {"method": "cad", "prompt": "Based on the passage, in which city does the Eiffel Tower stand?",
"context": "In 2031 the Eiffel Tower was rebuilt, piece by piece, in Rome.", "alpha": 1.0}),
("contrastive", {"method": "rose", "prompt": "Explain in two sentences why the sky is blue.", "alpha": 0.5}),
("contrastive", {"method": "dola", "prompt": "What is the capital of Washington State? Answer in one word.",
"rep": 1.2}),
("guided", {"method": "classifier", "attribute": "sentiment", "target": "positive", "lam": 3.0, "k": 8,
"lookahead": 1, "prompt": "Write one sentence describing a rainy Monday morning."}),
("guided", {"method": "classifier", "attribute": "topic", "topic": "astronomy", "lam": 20.0, "k": 8,
"prompt": "Write two sentences about a trip to the beach."}),
("guided", {"method": "regex", "pattern": r"\d{4}-\d{2}-\d{2}",
"prompt": "On what date did the Berlin Wall fall? Reply with the date only, in YYYY-MM-DD format."}),
("parallel", {"method": "speculative", "gamma": 4, "mode": "greedy",
"prompt": "List the numbers from 1 to 40, separated by commas."}),
("parallel", {"method": "speculative", "gamma": 4, "mode": "sampling", "temperature": 0.8,
"prompt": "Write a haiku about quantum entanglement."}),
("parallel", {"method": "pld", "num_pred": 8,
"prompt": "Fix only the typos and return the text otherwise unchanged: 'The quick brown fox jumpd "
"over the lazzy dog and ran into the forrest.'"}),
("parallel", {"method": "jacobi", "block": 6, "prompt": "Count from 1 to 30, separated by spaces."}),
]
def render_all(trace: dict) -> None:
"""Exercise the same render/step functions the UI calls for this trace."""
import app
app.time.sleep = lambda _s: None # replays animate with sleeps; skip them here
if trace["paradigm"] == "contrastive":
app.render_contrastive(trace)
n = len(trace["runs"]["method"]["steps"])
for i in {0, n // 2, max(n - 1, 0)}:
app.step_contrastive(trace, i)
list(app.replay_contrastive(trace))
elif trace["paradigm"] == "guided":
app.render_guided(trace)
n = len(trace["runs"]["method"]["steps"])
for i in {0, n // 2, max(n - 1, 0)}:
app.step_guided(trace, i)
list(app.replay_guided(trace))
else:
app.render_parallel(trace)
n = len(trace["runs"]["method"]["iters"])
for k in {0, n // 2, max(n - 1, 0)}:
app.step_parallel(trace, k)
list(app.replay_parallel(trace))
plots.speedup_curves(trace["summary"] if trace["method"] == "speculative" else None, 0.7, 0.5)
def check(trace: dict) -> list[str]:
problems = []
if trace.get("schema") != "dd.trace/v1":
problems.append("bad schema")
meth = trace["runs"]["method"]
if not meth["ids"]:
problems.append("method produced no tokens")
s = trace["summary"]
if trace["paradigm"] == "parallel" and trace["params"]["mode"] == "greedy" and s.get("identical") is False:
problems.append(f"greedy output diverged from autoregressive at token {s.get('first_divergence')} (bf16?)")
if trace["method"] == "regex" and not s["valid"]["method"] and meth["finish"] != "length":
problems.append(f"regex output does not match: {meth['text']!r}")
json.dumps(trace, allow_nan=False)
return problems
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--url")
ap.add_argument("--space")
ap.add_argument("--families", default="qwen25,smollm2")
ap.add_argument("--tokens", type=int, default=24)
ap.add_argument("--only", default="", help="comma-separated methods to run")
args = ap.parse_args()
from gradio_client import Client
from huggingface_hub import get_token
token = os.environ.get("HF_TOKEN") or get_token() # logged-in calls get the account's GPU quota
client = Client(args.space, token=token) if args.space else Client(args.url or "http://127.0.0.1:7860")
only = {m for m in args.only.split(",") if m}
failures = 0
for family in args.families.split(","):
for paradigm, params in CASES:
if only and params["method"] not in only:
continue
P = {**params, "max_new_tokens": args.tokens}
t0 = time.time()
try:
trace = client.predict(paradigm, family, json.dumps(P), api_name="/trace")
render_all(trace)
problems = check(trace)
except Exception as exc: # report and keep going
problems = [f"{type(exc).__name__}: {str(exc)[:300]}"]
trace = None
dt = time.time() - t0
tag = "OK " if not problems else "FAIL"
failures += bool(problems)
extra = ""
if trace:
s = trace["summary"]
env = trace.get("env", {})
extra = (f"gpu {env.get('gpu_seconds')}s/{env.get('requested_seconds')}s "
+ " ".join(f"{k}={s[k]}" for k in ("tau", "speedup_wall", "accept_rate", "changed", "identical")
if k in s))
extra += f" | {trace['runs']['method']['text'][:70]!r}"
print(f"{tag} {family:8} {paradigm:11} {params['method']:11} {dt: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())
|