Dev4: ONNX export pipeline for Option A deployment

api/_shared.py

@@ -164,7 +164,32 @@ def get_solver():
             _solver_error = str(exc)
             print(f"[CalculusSolver] Neural proxy init failed: {exc}", flush=True)
 
-    # 2. Second priority: Local neural solver (resolve checkpoint path first)
+    # 2. Second priority: Local ONNX solver (lightweight — no torch import,
+    #    fits Vercel's ~250MB serverless size limit). See docs/EXPORT_DECISION.md
+    #    for the measured size/correctness verification for this artifact.
+    #    Must be tried before the local torch load below, since that path
+    #    requires bundling the full PyTorch package.
+    onnx_path = os.environ.get(
+        "ONNX_MODEL_PATH", str(ROOT / "deployment" / "artifacts" / "best.onnx")
+    )
+    if os.path.exists(onnx_path):
+        try:
+            from deployment.onnx_solve import ONNXCalculusSolverInference
+            _solver = ONNXCalculusSolverInference(model_path=onnx_path)
+            _solver_mode = "neural-onnx"
+            _solver_error = None
+            print(
+                f"[CalculusSolver] ONNX neural model loaded from '{onnx_path}'",
+                flush=True,
+            )
+            return _solver, _solver_mode
+        except Exception as exc:
+            _solver_error = str(exc)
+            print(f"[CalculusSolver] ONNX load failed: {exc}", flush=True)
+
+    # 3. Third priority: Local torch-based neural solver (resolve checkpoint path first).
+    #    Heavier than the ONNX path above -- only reached if no best.onnx is present
+    #    or it failed to load.
     model_path, stage = _resolve_model_path()
     if model_path is not None:
         try:
@@ -184,7 +209,7 @@ def get_solver():
                 flush=True,
             )
 
-    # 3. Third priority: Try to load GroqSolver (Fallback intelligent model)
+    # 4. Fourth priority: Try to load GroqSolver (Fallback intelligent model)
     api_key = os.environ.get("GROQ_API_KEY")
     if api_key:
         try:
@@ -202,7 +227,7 @@ def get_solver():
             _solver_error = str(exc)
             print(f"[CalculusSolver] Groq load failed: {exc}", flush=True)
 
-    # 4. Final priority: Fallback
+    # 5. Final priority: Fallback
     from inference.fallback_solver import FallbackSolver
     _solver = FallbackSolver()
     _solver_mode = "fallback"
@@ -310,5 +335,4 @@ def normalize_solver_result(result: dict, mode: str) -> dict:
         }
     else:
         # Fallback and Groq solver results already have the correct structure
-        return {**result, "mode": mode}
-
+        return {**result, "mode": mode}

deployment/export_onnx.py

@@ -0,0 +1,92 @@
+"""
+Exports the trained SimpleCalculusModel (model/simple_transformer.py) to
+ONNX for the torch-free production deployment path (Option A -- see
+docs/EXPORT_DECISION.md for the measured size comparison vs Vercel's
+~250MB serverless limit).
+
+NOTE: model/simple_transformer.py's SimpleCalculusModel is a single
+encoder-decoder nn.Transformer with rule prediction folded into the output
+sequence (see that file's docstring). There is no separate RuleHead to
+export -- unlike the older model/transformer.py design, this model has
+exactly one set of weights and one ONNX graph.
+"""
+
+import json
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from model.simple_transformer import SimpleCalculusModel
+from inference.grammar import load_vocab
+
+
+def export_to_onnx(
+    checkpoint_path: str = os.path.join("checkpoints", "final", "best.pt"),
+    output_path: str = os.path.join("deployment", "artifacts", "best.onnx"),
+    vocab_path: str = os.path.join("tokenizer", "vocab.json"),
+    config_path: str = "config.json",
+) -> str:
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(
+            f"PyTorch checkpoint not found: {checkpoint_path}\n"
+            "This must be a checkpoint that inference/solve.py can already "
+            "load successfully -- if solve.py fails to load it, export will "
+            "fail with the identical state_dict mismatch. Confirm with "
+            "Developer 3 that this checkpoint is signed off before exporting."
+        )
+    if not os.path.exists(vocab_path):
+        raise FileNotFoundError(f"Vocab file not found: {vocab_path}")
+
+    vocab_map = load_vocab(vocab_path)
+    vocab_size = max(vocab_map["token_to_id"].values()) + 1
+    pad_id = vocab_map["token_to_id"]["[PAD]"]
+
+    hidden_dim = 128
+    max_len = 32
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            cfg = json.load(f)
+            hidden_dim = cfg.get("hidden_dim", hidden_dim)
+            max_len = cfg.get("max_len", max_len)
+
+    model = SimpleCalculusModel(
+        vocab_size=vocab_size,
+        hidden_dim=hidden_dim,
+        pad_id=pad_id,
+        max_len=max_len,
+    )
+
+    state_dict = torch.load(checkpoint_path, map_location="cpu")
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    dummy_src = torch.randint(1, vocab_size, (1, max_len), dtype=torch.long)
+    dummy_tgt_in = torch.randint(1, vocab_size, (1, max_len), dtype=torch.long)
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    torch.onnx.export(
+        model,
+        (dummy_src, dummy_tgt_in),
+        output_path,
+        input_names=["src_seq", "tgt_in_seq"],
+        output_names=["logits"],
+        dynamic_axes={
+            "src_seq": {0: "batch_size", 1: "seq_len"},
+            "tgt_in_seq": {0: "batch_size", 1: "tgt_len"},
+            "logits": {0: "batch_size", 1: "tgt_len"},
+        },
+        opset_version=14,
+    )
+
+    size_mb = os.path.getsize(output_path) / (1024 * 1024)
+    print(f"[export_onnx] Exported {checkpoint_path} -> {output_path} ({size_mb:.2f} MB)")
+    return output_path
+
+
+if __name__ == "__main__":
+    ckpt = sys.argv[1] if len(sys.argv) > 1 else os.path.join("checkpoints", "final", "best.pt")
+    out = sys.argv[2] if len(sys.argv) > 2 else os.path.join("deployment", "artifacts", "best.onnx")
+    export_to_onnx(ckpt, out)

deployment/onnx_beam_search.py

@@ -0,0 +1,108 @@
+"""
+numpy/onnxruntime-only mirror of inference/beam_search.py -- mirrors that
+file's beam_search() line-for-line in logic, but never imports torch.
+This is the entire point of the Option A (ONNX) deployment path: the
+production Vercel bundle only needs onnxruntime + numpy, not the full
+PyTorch package, which is what pushed the old bundle over the ~250MB
+serverless size limit.
+"""
+
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import onnxruntime as ort
+
+from inference.grammar import NodeValidityPool
+
+
+def _softmax(x: np.ndarray) -> np.ndarray:
+    x = x - np.max(x)
+    e = np.exp(x)
+    return e / e.sum()
+
+
+def onnx_beam_search(
+    session: ort.InferenceSession,
+    src_tokens: List[int],
+    vocab_map: Dict[str, Any],
+    beam_size: int = 5,
+    max_len: int = 32,
+    node_pool: Optional[NodeValidityPool] = None,
+) -> Dict[str, Any]:
+    """Mirrors inference/beam_search.py::beam_search(), but calls the
+    exported ONNX graph via onnxruntime instead of a torch.nn.Module."""
+    vocab = vocab_map["token_to_id"]
+    id_to_token = vocab_map["id_to_token"]
+    bos_id = vocab["[BOS]"]
+    eos_id = vocab["[EOS]"]
+
+    if node_pool is None:
+        node_pool = NodeValidityPool()
+
+    vocab_size = max(id_to_token.keys()) + 1
+    all_candidate_tokens = [id_to_token.get(idx, "[PAD]") for idx in range(vocab_size)]
+
+    src_arr = np.array([src_tokens], dtype=np.int64)
+
+    beams = [{"tokens": [bos_id], "score": 0.0, "finished": False}]
+    completed = []
+
+    for _ in range(max_len):
+        candidates = []
+        for beam in beams:
+            if beam["finished"]:
+                candidates.append(beam)
+                continue
+
+            current_tokens = beam["tokens"]
+            token_strings = [id_to_token[t] for t in current_tokens]
+            validity_tokens = (
+                token_strings[1:]
+                if token_strings and token_strings[0] == "[BOS]"
+                else token_strings
+            )
+
+            tgt_arr = np.array([current_tokens], dtype=np.int64)
+            logits = session.run(
+                ["logits"],
+                {"src_seq": src_arr, "tgt_in_seq": tgt_arr},
+            )[0]
+            next_logits = logits[0, -1, :]
+
+            mask = node_pool.mask(validity_tokens, all_candidate_tokens)
+            safe_logits = next_logits.copy()
+            safe_logits[[not v for v in mask]] = -np.inf
+
+            if np.all(np.isinf(safe_logits)):
+                continue
+
+            log_probs = np.log(_softmax(safe_logits) + 1e-12)
+            k = min(beam_size, safe_logits.shape[0])
+            top_idx = np.argpartition(-log_probs, k - 1)[:k]
+            top_idx = top_idx[np.argsort(-log_probs[top_idx])]
+
+            for token_id in top_idx:
+                token_id = int(token_id)
+                score = float(log_probs[token_id])
+                new_tokens = current_tokens + [token_id]
+                finished = token_id == eos_id
+                candidates.append({
+                    "tokens": new_tokens,
+                    "score": beam["score"] + score,
+                    "finished": finished,
+                })
+
+        if not candidates:
+            break
+
+        beams = sorted(candidates, key=lambda x: x["score"], reverse=True)[:beam_size]
+        if all(b["finished"] for b in beams):
+            completed.extend(beams)
+            break
+
+    best = sorted(completed, key=lambda x: x["score"], reverse=True)[0] if completed else (
+        beams[0] if beams else {"tokens": [bos_id], "score": 0.0, "finished": False}
+    )
+
+    status = "solved" if best["finished"] else "partial"
+    return {"tokens": best["tokens"], "score": best["score"], "status": status}

deployment/onnx_solve.py

@@ -0,0 +1,121 @@
+"""
+Torch-free mirror of inference/solve.py::CalculusSolverInference, using
+onnxruntime instead of a loaded PyTorch model. Mirrors that file's solve()
+line-for-line in logic. This is what the production Vercel API imports
+under Option A -- inference/solve.py (and torch) never gets imported in
+that deployment.
+"""
+
+import json
+import os
+from typing import Any, Dict, List
+
+import onnxruntime as ort
+
+from inference.grammar import NodeValidityPool, load_vocab
+from deployment.onnx_beam_search import onnx_beam_search
+
+
+class ONNXCalculusSolverInference:
+    def __init__(
+        self,
+        model_path: str = os.path.join("deployment", "artifacts", "best.onnx"),
+        vocab_path: str = os.path.join("tokenizer", "vocab.json"),
+        beam_size: int = 5,
+        max_len: int = 32,
+    ):
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"ONNX model not found: {model_path}")
+        if not os.path.exists(vocab_path):
+            raise FileNotFoundError(f"Vocab file not found: {vocab_path}")
+
+        self.vocab_map = load_vocab(vocab_path)
+        self.session = ort.InferenceSession(
+            model_path, providers=["CPUExecutionProvider"]
+        )
+
+        config_path = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.json"
+        )
+        if os.path.exists(config_path):
+            with open(config_path, "r") as f:
+                cfg = json.load(f)
+                max_len = cfg.get("max_len", max_len)
+
+        self.beam_size = beam_size
+        self.max_len = max_len
+        self.node_pool = NodeValidityPool()
+        self.bos_id = self.vocab_map["token_to_id"]["[BOS]"]
+        self.eos_id = self.vocab_map["token_to_id"]["[EOS]"]
+        self.pad_id = self.vocab_map["token_to_id"]["[PAD]"]
+
+    def close(self) -> None:
+        self.node_pool.close()
+
+    def _serialize_input(self, input_env: Dict[str, Any]) -> List[str]:
+        from tokenizer.slang_serializer import serialize_slang_math
+        return serialize_slang_math(input_env)
+
+    def _verify_output(self, input_env: Dict[str, Any], output_tokens: List[str]) -> Dict[str, Any]:
+        from inference.verifier import verify
+        return verify(input_env, output_tokens)
+
+    def solve(self, input_env: Dict[str, Any]) -> Dict[str, Any]:
+        token_strings = self._serialize_input(input_env)
+        token_ids = [
+            self.vocab_map["token_to_id"].get(token, self.pad_id)
+            for token in token_strings
+        ]
+        token_ids = token_ids[: self.max_len]
+        padded_tokens = token_ids + [self.pad_id] * (self.max_len - len(token_ids))
+
+        result = onnx_beam_search(
+            session=self.session,
+            src_tokens=padded_tokens,
+            vocab_map=self.vocab_map,
+            beam_size=self.beam_size,
+            max_len=self.max_len,
+            node_pool=self.node_pool,
+        )
+
+        output_token_strings = [
+            self.vocab_map["id_to_token"][t]
+            for t in result["tokens"]
+            if t in self.vocab_map["id_to_token"]
+        ]
+
+        if output_token_strings and output_token_strings[0] == "[BOS]":
+            output_token_strings = output_token_strings[1:]
+
+        predicted_rule = None
+        if output_token_strings and output_token_strings[0].startswith("RULE:"):
+            predicted_rule = output_token_strings[0]
+            output_token_strings = output_token_strings[1:]
+
+        verifier_result = self._verify_output(input_env, output_token_strings)
+        status = verifier_result.get("status", result.get("status"))
+        warning = verifier_result.get("error")
+
+        return {
+            "input": input_env,
+            "output_tokens": output_token_strings,
+            "status": status,
+            "verified": verifier_result.get("verified", False),
+            "confidence": verifier_result.get("confidence", 0),
+            "rule": predicted_rule,
+            "output": verifier_result.get("output"),
+            "warning": warning,
+        }
+
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) < 2:
+        raise SystemExit("Usage: python deployment/onnx_solve.py input.json")
+    with open(sys.argv[1], "r", encoding="utf-8") as f:
+        payload = json.load(f)
+    solver = ONNXCalculusSolverInference()
+    try:
+        print(json.dumps(solver.solve(payload), indent=2))
+    finally:
+        solver.close()