diff --git a/api/_shared.py b/api/_shared.py index 7adbc34..63f1276 100644 --- a/api/_shared.py +++ b/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} \ No newline at end of file diff --git a/deployment/export_onnx.py b/deployment/export_onnx.py new file mode 100644 index 0000000..aecef2d --- /dev/null +++ b/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) \ No newline at end of file diff --git a/deployment/onnx_beam_search.py b/deployment/onnx_beam_search.py new file mode 100644 index 0000000..78530fb --- /dev/null +++ b/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} \ No newline at end of file diff --git a/deployment/onnx_solve.py b/deployment/onnx_solve.py new file mode 100644 index 0000000..7a86706 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/deployment/verify_export.py b/deployment/verify_export.py new file mode 100644 index 0000000..f2bf1ae --- /dev/null +++ b/deployment/verify_export.py @@ -0,0 +1,178 @@ +""" +Verifies the ONNX export against the original PyTorch checkpoint: + 1. Runs every problem in eval/benchmarks/*.json through BOTH the PyTorch + path (inference/solve.py) and the ONNX path (deployment/onnx_solve.py). + 2. Checks token-exact match between the two -- conversion must not + silently change numerical/generation behaviour. + 3. Measures the real deployment bundle size against Vercel's ~250MB + serverless function limit (not just the raw .onnx file -- the bundle + includes onnxruntime + numpy + application code). + 4. Writes docs/EXPORT_DECISION.md from these actual measured numbers. + +Per the task's re-validation rule: this must be re-run against every new +signed-off checkpoint. A different checkpoint can convert/behave +differently even if a previous one converted cleanly -- do not assume +results from a prior run still hold. +""" + +import glob +import json +import os +import subprocess +import sys +from datetime import datetime, timezone + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +VERCEL_LIMIT_MB = 250 + + +def _load_benchmarks(pattern: str = os.path.join("eval", "benchmarks", "*.json")): + problems = [] + for path in sorted(glob.glob(pattern)): + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + items = data if isinstance(data, list) else data.get("problems", []) + for item in items: + problems.append(item) + return problems + + +def _measure_bundle_size_mb(onnx_path: str) -> float: + """Real bundle size: the .onnx weights file plus the installed + onnxruntime + numpy packages that ship alongside it in production, + per requirements-onnx.txt -- not just the raw model file.""" + total_bytes = os.path.getsize(onnx_path) + + try: + import onnxruntime, numpy + for mod in (onnxruntime, numpy): + mod_dir = os.path.dirname(mod.__file__) + for root, _, files in os.walk(mod_dir): + for fname in files: + fpath = os.path.join(root, fname) + if os.path.exists(fpath): + total_bytes += os.path.getsize(fpath) + except ImportError: + pass + + return total_bytes / (1024 * 1024) + + +def run_verification( + checkpoint_path: str = os.path.join("checkpoints", "final", "best.pt"), + onnx_path: str = os.path.join("deployment", "artifacts", "best.onnx"), +) -> dict: + from inference.solve import CalculusSolverInference + from deployment.onnx_solve import ONNXCalculusSolverInference + + problems = _load_benchmarks() + if not problems: + raise RuntimeError("No benchmark problems found under eval/benchmarks/*.json") + + pt_solver = CalculusSolverInference(model_path=checkpoint_path) + onnx_solver = ONNXCalculusSolverInference(model_path=onnx_path) + + total = 0 + exact_matches = 0 + mismatches = [] + + try: + for problem in problems: + input_env = problem.get("input", problem) + total += 1 + + pt_result = pt_solver.solve(input_env) + onnx_result = onnx_solver.solve(input_env) + + match = pt_result["output_tokens"] == onnx_result["output_tokens"] + if match: + exact_matches += 1 + else: + mismatches.append({ + "input": input_env, + "pytorch_output": pt_result["output_tokens"], + "onnx_output": onnx_result["output_tokens"], + }) + finally: + pt_solver.close() + onnx_solver.close() + + match_rate = exact_matches / total if total else 0.0 + bundle_size_mb = _measure_bundle_size_mb(onnx_path) + fits_limit = bundle_size_mb < VERCEL_LIMIT_MB + + return { + "checkpoint": checkpoint_path, + "onnx_path": onnx_path, + "total_problems": total, + "exact_matches": exact_matches, + "match_rate": match_rate, + "bundle_size_mb": bundle_size_mb, + "vercel_limit_mb": VERCEL_LIMIT_MB, + "fits_limit": fits_limit, + "mismatches": mismatches[:20], # cap for report readability + "mismatch_count": len(mismatches), + } + + +def write_decision_doc(report: dict, out_path: str = os.path.join("docs", "EXPORT_DECISION.md")) -> None: + timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + verdict = "PASS" if report["fits_limit"] and report["match_rate"] == 1.0 else "NEEDS ATTENTION" + + lines = [ + "# Export Decision — Option A (ONNX)", + "", + f"**Generated:** {timestamp}", + f"**Checkpoint verified:** `{report['checkpoint']}`", + f"**ONNX artifact:** `{report['onnx_path']}`", + "", + f"## Verdict: {verdict}", + "", + "## Numerical correctness (PyTorch vs ONNX)", + f"- Problems tested: {report['total_problems']}", + f"- Exact token match: {report['exact_matches']}/{report['total_problems']} " + f"({report['match_rate']:.1%})", + f"- Mismatches: {report['mismatch_count']}", + "", + "## Deployment size vs Vercel limit", + f"- Measured bundle size (.onnx + onnxruntime + numpy): {report['bundle_size_mb']:.1f} MB", + f"- Vercel serverless limit: {report['vercel_limit_mb']} MB", + f"- Fits limit: {'Yes' if report['fits_limit'] else 'No'}", + "", + ] + + if report["mismatch_count"] > 0: + lines.append("## Sample mismatches (first 20)") + lines.append("") + for i, m in enumerate(report["mismatches"], 1): + lines.append(f"### Mismatch {i}") + lines.append(f"- Input: `{json.dumps(m['input'])}`") + lines.append(f"- PyTorch output: `{m['pytorch_output']}`") + lines.append(f"- ONNX output: `{m['onnx_output']}`") + lines.append("") + + lines.append( + "**Re-validation rule:** this file must be regenerated against every new " + "signed-off checkpoint. Do not treat this verdict as valid for a checkpoint " + "other than the one named above." + ) + + os.makedirs(os.path.dirname(out_path), exist_ok=True) + with open(out_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + + print(f"[verify_export] Wrote {out_path} -- verdict: {verdict}") + + +if __name__ == "__main__": + ckpt = sys.argv[1] if len(sys.argv) > 1 else os.path.join("checkpoints", "final", "best.pt") + onnx = sys.argv[2] if len(sys.argv) > 2 else os.path.join("deployment", "artifacts", "best.onnx") + + report = run_verification(ckpt, onnx) + write_decision_doc(report) + + print(json.dumps( + {k: v for k, v in report.items() if k != "mismatches"}, + indent=2, + )) \ No newline at end of file diff --git a/inference/beam_search.py b/inference/beam_search.py index 2c9a418..16a0c11 100644 --- a/inference/beam_search.py +++ b/inference/beam_search.py @@ -4,192 +4,12 @@ import torch +from inference.grammar import NodeValidityPool, flatten_vocab, load_vocab, is_valid_prefix -def is_valid_prefix(tokens: List[str]) -> bool: - """Check if the given tokens form a valid prefix of a SLaNg AST. - First token, if present, may be a RULE:xxx token (SimpleCalculusModel - prepends one) -- skip it before running the AST grammar check.""" - if not tokens: - return True - - check_tokens = tokens - if tokens[0].startswith("RULE:"): - check_tokens = tokens[1:] - if not check_tokens: - return True - - tokens = check_tokens - - def parse_term(index: int) -> dict: - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] != "NODE:TERM": - return {"status": "invalid"} - index += 1 - if index >= len(tokens): - return {"status": "incomplete"} - if not tokens[index].startswith("COEF:"): - return {"status": "invalid"} - index += 1 - while index < len(tokens): - token = tokens[index] - if token.startswith("VAR:"): - index += 1 - if index >= len(tokens): - return {"status": "incomplete"} - if not tokens[index].startswith("EXP:"): - return {"status": "invalid"} - index += 1 - continue - break - return {"status": "complete", "next": index} - - def parse_term_list(index: int) -> dict: - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] == "STRUCT:CLOSE": - return {"status": "complete", "next": index} - current = index - while True: - node = parse_node(current) - if node["status"] == "invalid": - return {"status": "invalid"} - if node["status"] == "incomplete": - return {"status": "incomplete"} - current = node["next"] - if current >= len(tokens): - return {"status": "incomplete"} - if tokens[current] == "STRUCT:SEP": - current += 1 - continue - if tokens[current] == "STRUCT:CLOSE": - return {"status": "complete", "next": current} - return {"status": "invalid"} - - def parse_fraction(index: int) -> dict: - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] != "NODE:FRAC": - return {"status": "invalid"} - index += 1 - for expected in ["STRUCT:OPEN", "STRUCT:NUMI", "STRUCT:OPEN"]: - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] != expected: - return {"status": "invalid"} - index += 1 - numerator = parse_term_list(index) - if numerator["status"] != "complete": - return numerator - index = numerator["next"] - for expected in ["STRUCT:CLOSE", "STRUCT:SEP", "STRUCT:DENO", "STRUCT:OPEN"]: - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] != expected: - return {"status": "invalid"} - index += 1 - denominator = parse_term_list(index) - if denominator["status"] != "complete": - return denominator - index = denominator["next"] - for expected in ["STRUCT:CLOSE", "STRUCT:CLOSE"]: - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] != expected: - return {"status": "invalid"} - index += 1 - return {"status": "complete", "next": index} - - def parse_op_node(index: int) -> dict: - if index >= len(tokens): - return {"status": "incomplete"} - token = tokens[index] - if not isinstance(token, str) or not token.startswith("OP:"): - return {"status": "invalid"} - index += 1 - while ( - index < len(tokens) - and isinstance(tokens[index], str) - and tokens[index].startswith("OPVAR:") - ): - index += 1 - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] != "STRUCT:OPEN": - return {"status": "invalid"} - index += 1 - seen_child = False - while True: - node = parse_node(index) - if node["status"] == "invalid": - return {"status": "invalid"} - if node["status"] == "incomplete": - return {"status": "incomplete"} - seen_child = True - index = node["next"] - if index >= len(tokens): - return {"status": "incomplete"} - if tokens[index] == "STRUCT:SEP": - index += 1 - continue - if tokens[index] == "STRUCT:CLOSE": - if not seen_child: - return {"status": "invalid"} - index += 1 - return {"status": "complete", "next": index} - return {"status": "invalid"} - - def parse_node(index: int) -> dict: - if index >= len(tokens): - return {"status": "incomplete"} - token = tokens[index] - if token == "NODE:TERM": - return parse_term(index) - if token == "NODE:FRAC": - return parse_fraction(index) - if isinstance(token, str) and token.startswith("OP:"): - return parse_op_node(index) - return {"status": "invalid"} - - result = parse_node(0) - if result["status"] == "invalid": - return False - if result["status"] == "incomplete": - return True - return result["status"] == "complete" and result["next"] == len(tokens) - - -class NodeValidityPool: - def __init__(self, script_path: str = "", num_workers: int = 1): - pass - - def mask(self, tokens: List[str], candidate_tokens: List[str]) -> List[bool]: - return [is_valid_prefix(tokens + [candidate]) for candidate in candidate_tokens] - - def close(self) -> None: - pass - - -def flatten_vocab(vocab: Dict[str, Any]) -> Dict[str, int]: - token_to_id = {} - for key, value in vocab.items(): - if key.startswith("_"): - continue - if isinstance(value, dict): - token_to_id.update(value) - return token_to_id - - -def load_vocab(vocab_path: str) -> Dict[str, Any]: - with open(vocab_path, "r", encoding="utf-8") as f: - raw = json.load(f) - flat = flatten_vocab(raw) - id_to_token = {idx: token for token, idx in flat.items()} - return { - "token_to_id": flat, - "id_to_token": id_to_token, - "special": raw.get("special_tokens", {}), - } +# NOTE: NodeValidityPool, flatten_vocab, load_vocab, and is_valid_prefix now +# live in inference/grammar.py (torch-free) so the ONNX deployment path +# (deployment/onnx_beam_search.py) can reuse them without importing torch. +# This is a pure move -- no logic changed from the previous inline versions. def beam_search( diff --git a/inference/grammar.py b/inference/grammar.py new file mode 100644 index 0000000..87c8c46 --- /dev/null +++ b/inference/grammar.py @@ -0,0 +1,200 @@ +""" +Torch-free SLaNg grammar and vocab helpers, split out of inference/beam_search.py. + +Purpose: the ONNX deployment path (deployment/onnx_beam_search.py, +deployment/onnx_solve.py) must never import torch -- that's the entire +point of Option A (ONNX export) as a fix for Vercel's ~250MB serverless +size limit. Everything in this file is pure Python / stdlib only, so both +the PyTorch path (inference/beam_search.py) and the ONNX path +(deployment/onnx_beam_search.py) can import it without pulling torch in. +""" + +import json +from typing import Any, Dict, List + + +def is_valid_prefix(tokens: List[str]) -> bool: + """Check if the given tokens form a valid prefix of a SLaNg AST. + First token, if present, may be a RULE:xxx token (SimpleCalculusModel + prepends one) -- skip it before running the AST grammar check.""" + if not tokens: + return True + + check_tokens = tokens + if tokens[0].startswith("RULE:"): + check_tokens = tokens[1:] + if not check_tokens: + return True + + tokens = check_tokens + + def parse_term(index: int) -> dict: + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] != "NODE:TERM": + return {"status": "invalid"} + index += 1 + if index >= len(tokens): + return {"status": "incomplete"} + if not tokens[index].startswith("COEF:"): + return {"status": "invalid"} + index += 1 + while index < len(tokens): + token = tokens[index] + if token.startswith("VAR:"): + index += 1 + if index >= len(tokens): + return {"status": "incomplete"} + if not tokens[index].startswith("EXP:"): + return {"status": "invalid"} + index += 1 + continue + break + return {"status": "complete", "next": index} + + def parse_term_list(index: int) -> dict: + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] == "STRUCT:CLOSE": + return {"status": "complete", "next": index} + current = index + while True: + node = parse_node(current) + if node["status"] == "invalid": + return {"status": "invalid"} + if node["status"] == "incomplete": + return {"status": "incomplete"} + current = node["next"] + if current >= len(tokens): + return {"status": "incomplete"} + if tokens[current] == "STRUCT:SEP": + current += 1 + continue + if tokens[current] == "STRUCT:CLOSE": + return {"status": "complete", "next": current} + return {"status": "invalid"} + + def parse_fraction(index: int) -> dict: + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] != "NODE:FRAC": + return {"status": "invalid"} + index += 1 + for expected in ["STRUCT:OPEN", "STRUCT:NUMI", "STRUCT:OPEN"]: + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] != expected: + return {"status": "invalid"} + index += 1 + numerator = parse_term_list(index) + if numerator["status"] != "complete": + return numerator + index = numerator["next"] + for expected in ["STRUCT:CLOSE", "STRUCT:SEP", "STRUCT:DENO", "STRUCT:OPEN"]: + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] != expected: + return {"status": "invalid"} + index += 1 + denominator = parse_term_list(index) + if denominator["status"] != "complete": + return denominator + index = denominator["next"] + for expected in ["STRUCT:CLOSE", "STRUCT:CLOSE"]: + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] != expected: + return {"status": "invalid"} + index += 1 + return {"status": "complete", "next": index} + + def parse_op_node(index: int) -> dict: + if index >= len(tokens): + return {"status": "incomplete"} + token = tokens[index] + if not isinstance(token, str) or not token.startswith("OP:"): + return {"status": "invalid"} + index += 1 + while ( + index < len(tokens) + and isinstance(tokens[index], str) + and tokens[index].startswith("OPVAR:") + ): + index += 1 + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] != "STRUCT:OPEN": + return {"status": "invalid"} + index += 1 + seen_child = False + while True: + node = parse_node(index) + if node["status"] == "invalid": + return {"status": "invalid"} + if node["status"] == "incomplete": + return {"status": "incomplete"} + seen_child = True + index = node["next"] + if index >= len(tokens): + return {"status": "incomplete"} + if tokens[index] == "STRUCT:SEP": + index += 1 + continue + if tokens[index] == "STRUCT:CLOSE": + if not seen_child: + return {"status": "invalid"} + index += 1 + return {"status": "complete", "next": index} + return {"status": "invalid"} + + def parse_node(index: int) -> dict: + if index >= len(tokens): + return {"status": "incomplete"} + token = tokens[index] + if token == "NODE:TERM": + return parse_term(index) + if token == "NODE:FRAC": + return parse_fraction(index) + if isinstance(token, str) and token.startswith("OP:"): + return parse_op_node(index) + return {"status": "invalid"} + + result = parse_node(0) + if result["status"] == "invalid": + return False + if result["status"] == "incomplete": + return True + return result["status"] == "complete" and result["next"] == len(tokens) + + +class NodeValidityPool: + def __init__(self, script_path: str = "", num_workers: int = 1): + pass + + def mask(self, tokens: List[str], candidate_tokens: List[str]) -> List[bool]: + return [is_valid_prefix(tokens + [candidate]) for candidate in candidate_tokens] + + def close(self) -> None: + pass + + +def flatten_vocab(vocab: Dict[str, Any]) -> Dict[str, int]: + token_to_id = {} + for key, value in vocab.items(): + if key.startswith("_"): + continue + if isinstance(value, dict): + token_to_id.update(value) + return token_to_id + + +def load_vocab(vocab_path: str) -> Dict[str, Any]: + with open(vocab_path, "r", encoding="utf-8") as f: + raw = json.load(f) + flat = flatten_vocab(raw) + id_to_token = {idx: token for token, idx in flat.items()} + return { + "token_to_id": flat, + "id_to_token": id_to_token, + "special": raw.get("special_tokens", {}), + } \ No newline at end of file diff --git a/requirements-onnx.txt b/requirements-onnx.txt new file mode 100644 index 0000000..98ffd3d --- /dev/null +++ b/requirements-onnx.txt @@ -0,0 +1,10 @@ +# Production requirements for Option A (ONNX export) deployment. +# Deliberately does NOT include torch -- that's the entire point of this +# path. See docs/EXPORT_DECISION.md for the measured size comparison +# against Vercel's ~250MB serverless limit. +onnxruntime==1.18.1 +numpy==1.26.4 +starlette==1.3.1 +uvicorn==0.49.0 +python-dotenv==1.2.2 +groq==1.5.0 \ No newline at end of file