Temporarily compare PyTorch and ONNX on the Space's hardware

server.py gains an "onnx" engine (checkpoint exported at startup by
onnx_backend.py). With LAYA_BENCH=1, set by app.py for now, both engines
answer the same questions in a background thread after startup; results
are served on GET /bench and printed to the logs. The demo keeps running
on PyTorch meanwhile.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

app.py

@@ -28,6 +28,8 @@ os.environ.setdefault("LAYA_MODELS", "typed")
 os.environ.setdefault("LAYA_DEVICE", "cpu")
 # About 1 s per request on the Space's CPU: wait for a pause in typing before calling the model.
 os.environ.setdefault("LAYA_DEBOUNCE", "700")
+# Temporary: compare PyTorch and ONNX on the Space's hardware (see GET /bench).
+os.environ.setdefault("LAYA_BENCH", "1")
 
 import gradio as gr  # noqa: E402
 import uvicorn  # noqa: E402

onnx_backend.py

@@ -0,0 +1,102 @@
+"""ONNX Runtime engine for the demo: the original laya checkpoint exported to ONNX at startup.
+
+The export follows laya's scripts/export_onnx.py (Apache 2.0). The resulting file is not on the
+Hub, so it is rebuilt on each start (about a minute on CPU) and kept in LAYA_ONNX_DIR.
+"""
+import os
+import statistics
+import time
+
+ONNX_DIR = os.environ.get("LAYA_ONNX_DIR", os.path.join(os.path.expanduser("~"), ".cache", "laya-onnx"))
+
+
+def onnx_path_for(model_id):
+    return os.path.join(ONNX_DIR, model_id.replace("/", "--"), "model.onnx")
+
+
+def export(model_id, path):
+    """Trace the PyTorch model with variable batch and sequence sizes and save it as ONNX."""
+    import torch
+    from laya.agent import Agent
+
+    agent = Agent(model_id, compile=False, device="cpu")
+    inputs = (
+        torch.randint(0, 100, (1, 16), dtype=torch.long),  # input_ids
+        torch.ones((1, 16), dtype=torch.long),  # attention_mask
+        torch.tensor([[1, 5]], dtype=torch.long),  # marker_pos
+        torch.tensor([[True, True]], dtype=torch.bool),  # marker_mask
+        torch.tensor([0], dtype=torch.long),  # qtype
+    )
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    torch.onnx.export(
+        agent.model, inputs, path,
+        export_params=True, opset_version=18, do_constant_folding=True,
+        input_names=["input_ids", "attention_mask", "marker_pos", "marker_mask", "qtype"],
+        output_names=["logits", "act_logits"],
+        dynamic_axes={
+            "input_ids": {0: "batch_size", 1: "seq_len"},
+            "attention_mask": {0: "batch_size", 1: "seq_len"},
+            "marker_pos": {0: "batch_size", 1: "num_markers"},
+            "marker_mask": {0: "batch_size", 1: "num_markers"},
+            "qtype": {0: "batch_size"},
+            "logits": {0: "batch_size", 1: "num_markers"},
+            "act_logits": {0: "batch_size"},
+        },
+    )
+
+
+def load(model_id):
+    from laya.onnx_agent import ONNXAgent
+
+    path = onnx_path_for(model_id)
+    if not os.path.exists(path):
+        start = time.perf_counter()
+        export(model_id, path)
+        print(f"ONNX export of {model_id} took {time.perf_counter() - start:.0f} s", flush=True)
+    return ONNXAgent(model_id, onnx_path=path)
+
+
+# Texts of increasing length for the comparison; answers must not depend on the engine.
+BENCH_TEXTS = [
+    "Félicitations ! Vous avez gagné un iPhone 17. Cliquez ici pour récupérer votre cadeau.",
+    "Bonjour, je suis CFO d'une scale-up de 120 personnes. Nous devons remplacer notre outil de "
+    "reporting avant la clôture de décembre et nous avons prévu une enveloppe de 40k€.",
+    "Bonjour, je suis directrice générale d'un groupe de distribution de 1 200 salariés. Notre "
+    "plateforme e-commerce plante à chaque pic de ventes et nous perdons du chiffre d'affaires. "
+    "Le conseil a approuvé un budget de 250k€ pour la refondre, je suis seule décisionnaire et la "
+    "mise en ligne doit se faire avant le Black Friday. Pouvez-vous démarrer dans deux semaines ?",
+]
+
+
+def _probabilities(answers):
+    values = []
+    for key in sorted(answers):
+        answer = answers[key]
+        values += [answer["noul"]] if answer["type"] == "noul" else list(answer["probabilities"].values())
+    return values
+
+
+def compare(engines, cases, repeats=3):
+    """Run every case on every text with each engine; report median times and the largest gap."""
+    times = {name: [] for name in engines}
+    worst_gap = 0.0
+    for case in cases.values():
+        questions = case["questions"]()
+        for agent in engines.values():
+            agent.predict("warmup", questions)
+        for text in BENCH_TEXTS:
+            results = {}
+            for name, agent in engines.items():
+                for _ in range(repeats):
+                    start = time.perf_counter()
+                    results[name] = _probabilities(agent.predict(text, questions)["answers"])
+                    times[name].append(time.perf_counter() - start)
+            first, *others = results.values()
+            for other in others:
+                worst_gap = max(worst_gap, max(abs(a - b) for a, b in zip(first, other)))
+    return {
+        "median_ms": {name: round(statistics.median(t) * 1000) for name, t in times.items()},
+        "max_probability_gap": round(worst_gap, 4),
+        "cpu_count": os.cpu_count(),
+        "requests_per_engine": len(times[next(iter(times))]),
+    }

requirements.txt

@@ -5,3 +5,7 @@ laya==0.3.20
 fastapi
 uvicorn
 spaces
+# Temporary ONNX comparison (onnx_backend.py).
+onnx
+onnxruntime
+onnxscript

server.py

@@ -13,12 +13,15 @@ from questions import CASES
 
 WEB = Path(__file__).parent / "web"
 
-# Two engines run the same Laya weights and return the same answers:
+# Three engines run the same Laya weights and return the same answers:
 # - "mlx": laya-mlx, Apple Silicon only, fast local runs;
-# - "torch": the original laya package (PyTorch), for Linux servers.
+# - "torch": the original laya package (PyTorch), for Linux servers;
+# - "onnx": the same checkpoint exported to ONNX at startup (onnx_backend.py), CPU only.
+UPSTREAM = {"typed": "convaiinnovations/laya-typed-decisions", "english": "convaiinnovations/laya"}
 CHECKPOINTS = {
     "mlx": {"typed": "aac6fef/laya-typed-decisions-mlx", "english": "aac6fef/laya-mlx"},
-    "torch": {"typed": "convaiinnovations/laya-typed-decisions", "english": "convaiinnovations/laya"},
+    "torch": UPSTREAM,
+    "onnx": UPSTREAM,
 }
 APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
 BACKEND = os.environ.get("LAYA_BACKEND", "mlx" if APPLE_SILICON else "torch")
@@ -41,7 +44,16 @@ if origins:
 
 # Optional device override ("cpu" to mimic a GPU-less server when testing locally).
 device = {"device": os.environ["LAYA_DEVICE"]} if os.environ.get("LAYA_DEVICE") else {}
-agents = {alias: laya.load(model_id, **device) for alias, model_id in MODELS.items()}
+
+
+def load_agent(backend, model_id):
+    if backend == "onnx":
+        import onnx_backend
+        return onnx_backend.load(model_id)
+    return laya.load(model_id, **device)
+
+
+agents = {alias: load_agent(BACKEND, model_id) for alias, model_id in MODELS.items()}
 for agent in agents.values():
     agent.predict("warmup", {"q": {"type": "noul", "instructions": "Is this a test?"}})
 
@@ -93,6 +105,34 @@ def explain(name: str):
             "routing": {key: MODELS[model_for(case, key)] for key in case["questions"]()}}
 
 
+# Temporary benchmark (LAYA_BENCH=1): times PyTorch and ONNX on this machine, in the
+# background after startup since it takes about a minute. Results: GET /bench and the logs.
+bench_state = {"status": "disabled"}
+
+
+def run_bench():
+    import onnx_backend
+    bench_state["status"] = "running"
+    try:
+        model_id = UPSTREAM["typed"]
+        torch_agent = agents["typed"] if BACKEND == "torch" else load_agent("torch", model_id)
+        engines = {"torch": torch_agent, "onnx": load_agent("onnx", model_id)}
+        bench_state.update(status="done", result=onnx_backend.compare(engines, CASES))
+    except Exception as error:
+        bench_state.update(status="failed", error=f"{type(error).__name__}: {error}")
+    print(f"bench: {bench_state}", flush=True)
+
+
+if os.environ.get("LAYA_BENCH") == "1":
+    import threading
+    threading.Thread(target=run_bench, daemon=True).start()
+
+
+@app.get("/bench")
+def bench():
+    return bench_state
+
+
 @app.get("/health")
 def health():
     return {"backend": BACKEND, "models": list(MODELS.values())}