Download vons/onnx_runtime.py from INLEVEL9/Vons: direct link, hf CLI and curl.
- Browser
- Download file 23.5 kB
-
https://huggingface.co/INLEVEL9/Vons/resolve/main/vons/onnx_runtime.py
- Command line
-
hf download hf://INLEVEL9/Vons/vons/onnx_runtime.py
-
curl -L -o onnx_runtime.py https://huggingface.co/INLEVEL9/Vons/resolve/main/vons/onnx_runtime.py
23.5 kB
| """ONNX-backed decision runtime with hash-verified bundles. | |
| Host-owned runtime layer. The caller supplies a bundle directory with a valid | |
| ``vons.bundle.manifest/v1`` manifest; this module verifies every recorded byte | |
| before executing the graph and strictly validates request/response against the | |
| public ``DecisionRequest`` / ``DecisionResponse`` contract. | |
| Synthetic backends (``SyntheticAbstainBackend``, ``SyntheticChoiceBackend``) | |
| exist for protocol testing and must never be used as substitutes for a verified | |
| ONNX decision run. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| from abc import ABC, abstractmethod | |
| from collections.abc import Mapping | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any | |
| from .bundle import verify_bundle_manifest | |
| from .contract import ( | |
| Backend, | |
| DecisionRequest, | |
| DecisionResponse, | |
| OptionProbability, | |
| QuestionAnswer, | |
| ResponseStatus, | |
| validate_request, | |
| ) | |
| class OnnxRuntimeError(RuntimeError): | |
| """Raised when the ONNX runtime or bundle verification fails. | |
| The message is intentionally generic. Callers must not echo exception data | |
| or request payloads into shared logs. | |
| """ | |
| class OnnxRuntimeUnavailableError(OnnxRuntimeError): | |
| """Raised when an optional local inference dependency is unavailable.""" | |
| class VerifiedBundle: | |
| """A manifest-verified bundle directory, ready for inference.""" | |
| root: Path | |
| manifest_path: Path | |
| manifest: Mapping[str, Any] | |
| model_id: str | |
| def load(cls, bundle_dir: str | Path) -> VerifiedBundle: | |
| try: | |
| root = Path(bundle_dir).resolve(strict=True) | |
| except (OSError, FileNotFoundError): | |
| raise OnnxRuntimeError("bundle path does not exist") from None | |
| if not root.is_dir(): | |
| raise OnnxRuntimeError("bundle path is not a directory") | |
| manifest_path = None | |
| for candidate in ( | |
| root / "bundle-manifest.json", | |
| root / "bundle-manifest-v1.json", | |
| ): | |
| if not candidate.is_file(): | |
| continue | |
| try: | |
| candidate_payload = json.loads(candidate.read_text(encoding="utf-8")) | |
| except (OSError, json.JSONDecodeError): | |
| continue | |
| if ( | |
| isinstance(candidate_payload, Mapping) | |
| and candidate_payload.get("schema_version") == "vons.bundle.manifest/v1" | |
| ): | |
| manifest_path = candidate | |
| break | |
| if manifest_path is None: | |
| raise OnnxRuntimeError("bundle manifest is missing") | |
| try: | |
| verification = verify_bundle_manifest(manifest_path) | |
| except Exception: # noqa: BLE001 - normalize to sanitized form | |
| raise OnnxRuntimeError("bundle manifest verification failed") from None | |
| _ = verification # passed; record consumed, payload never echoed | |
| try: | |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| if not isinstance(manifest, Mapping): | |
| raise TypeError("bundle manifest must be an object") | |
| model_id = hashlib.sha256(manifest_path.read_bytes()).hexdigest() | |
| except (OSError, TypeError, ValueError): | |
| raise OnnxRuntimeError("bundle manifest could not be read") from None | |
| return cls(root=root, manifest_path=manifest_path, manifest=manifest, model_id=model_id) | |
| class DecisionBackend(ABC): | |
| """Strictly-typed decision backend. | |
| Implementations validate request input against the public contract before | |
| execution and validate every response field before returning to the host. | |
| No backend may emit a fallback or coerced choice when validation fails. | |
| """ | |
| backend_type: Backend | |
| def decide(self, request: DecisionRequest) -> DecisionResponse: | |
| """Execute a validated request and return a validated response.""" | |
| def _validate_request(self, request: DecisionRequest) -> None: | |
| if not isinstance(request, DecisionRequest): | |
| raise OnnxRuntimeError("invalid decision request type") | |
| try: | |
| validate_request(request) | |
| except (TypeError, ValueError) as exc: | |
| raise OnnxRuntimeError("decision request failed validation") from exc | |
| class SyntheticAbstainBackend(DecisionBackend): | |
| """Test-only backend that always abstains with a deterministic reason. | |
| The reason byte is preserved exactly so protocol round-trips can be | |
| asserted without running model weights. | |
| """ | |
| backend_type = Backend.DIRECT | |
| def __init__(self, abstain_reason: str = "synthetic backend requested abstention") -> None: | |
| self.abstain_reason = abstain_reason | |
| def decide(self, request: DecisionRequest) -> DecisionResponse: | |
| self._validate_request(request) | |
| answers = tuple( | |
| QuestionAnswer( | |
| question_id=question.id, | |
| choice=None, | |
| probabilities=(), | |
| confidence=0.0, | |
| status=ResponseStatus.ABSTAIN, | |
| abstain_reason=self.abstain_reason, | |
| ) | |
| for question in request.questions | |
| ) | |
| return DecisionResponse( | |
| answers=answers, | |
| backend=self.backend_type, | |
| model_id="synthetic-abstain-v1", | |
| metadata={"synthetic": True}, | |
| ) | |
| class SyntheticChoiceBackend(DecisionBackend): | |
| """Test-only backend that picks a fixed option by position. | |
| Used to validate exact question-id reconciliation and response shape | |
| without loading model weights. Callers set ``choice_index`` to a valid | |
| position for every question in the request; ``ValueError`` is raised | |
| when the index is out of range. | |
| """ | |
| backend_type = Backend.DIRECT | |
| def __init__(self, choice_index: int = 0, confidence: float = 0.9) -> None: | |
| if not isinstance(choice_index, int) or choice_index < 0: | |
| raise ValueError("choice_index must be a non-negative integer") | |
| if not 0 <= confidence <= 1: | |
| raise ValueError("confidence must be in [0, 1]") | |
| self.choice_index = choice_index | |
| self.confidence = float(confidence) | |
| def decide(self, request: DecisionRequest) -> DecisionResponse: | |
| self._validate_request(request) | |
| answers: list[QuestionAnswer] = [] | |
| for question in request.questions: | |
| options = question.options or ("true", "false") | |
| if self.choice_index >= len(options): | |
| raise OnnxRuntimeError("synthetic backend choice index out of range") | |
| choice = options[self.choice_index] | |
| probs: list[OptionProbability] = [] | |
| remaining = 1.0 - self.confidence | |
| for i, option in enumerate(options): | |
| if i == self.choice_index: | |
| p = self.confidence | |
| else: | |
| p = remaining / max(1, len(options) - 1) | |
| probs.append(OptionProbability(option=option, probability=float(p))) | |
| answers.append( | |
| QuestionAnswer( | |
| question_id=question.id, | |
| choice=choice, | |
| probabilities=tuple(probs), | |
| confidence=self.confidence, | |
| status=ResponseStatus.OK, | |
| ) | |
| ) | |
| return DecisionResponse( | |
| answers=tuple(answers), | |
| backend=self.backend_type, | |
| model_id="synthetic-choice-v1", | |
| metadata={"synthetic": True, "choice_index": self.choice_index}, | |
| ) | |
| class OnnxRuntimeBackend(DecisionBackend): | |
| """Verified ONNX Runtime backend. | |
| Requires the optional local-inference dependencies. The bundle is verified | |
| byte-for-byte before the session is created. Tokenizer files are resolved | |
| from the verified bundle only; remote tokenizer discovery is never used. | |
| """ | |
| def __init__( | |
| self, | |
| bundle: VerifiedBundle, | |
| *, | |
| tokenizer: Any | None = None, | |
| session: Any | None = None, | |
| ) -> None: | |
| if not isinstance(bundle, VerifiedBundle): | |
| raise OnnxRuntimeError("ONNX backend requires a verified bundle") | |
| self.bundle = bundle | |
| metadata = bundle.manifest.get("metadata", {}) | |
| if not isinstance(metadata, Mapping): | |
| raise OnnxRuntimeError("bundle runtime metadata is invalid") | |
| try: | |
| self.backend_type = Backend(metadata.get("backend", Backend.DIRECT.value)) | |
| except (TypeError, ValueError): | |
| raise OnnxRuntimeError("bundle backend is unsupported") from None | |
| if self.backend_type not in {Backend.DIRECT, Backend.DIFFUSION}: | |
| raise OnnxRuntimeError("bundle backend is unsupported") | |
| option_count = metadata.get("option_count", 32) | |
| sequence_length = metadata.get("sequence_length", 512) | |
| if ( | |
| isinstance(option_count, bool) | |
| or not isinstance(option_count, int) | |
| or not 2 <= option_count <= 32 | |
| or isinstance(sequence_length, bool) | |
| or not isinstance(sequence_length, int) | |
| or not 1 <= sequence_length <= 512 | |
| ): | |
| raise OnnxRuntimeError("bundle inference dimensions are invalid") | |
| self.option_count = option_count | |
| self.sequence_length = sequence_length | |
| self._metadata = metadata | |
| self._tokenizer = tokenizer | |
| self._session = session | |
| self._input_names = self._string_tuple(metadata.get("input_names"), "input_names") | |
| self._output_names = self._string_tuple(metadata.get("output_names"), "output_names") | |
| self._abstain_threshold = self._threshold(metadata.get("abstain_threshold", 0.55)) | |
| self._answerability_threshold = self._threshold( | |
| metadata.get("answerability_threshold", 0.5) | |
| ) | |
| conditioning = metadata.get("diffusion_conditioning", "pooled") | |
| if self.backend_type is Backend.DIFFUSION and conditioning not in { | |
| "pooled", | |
| "token_cross_attention", | |
| }: | |
| raise OnnxRuntimeError("bundle diffusion conditioning is unsupported") | |
| self._diffusion_conditioning = conditioning | |
| def _string_tuple(value: Any, name: str) -> tuple[str, ...]: | |
| if not isinstance(value, (list, tuple)) or any( | |
| not isinstance(item, str) or not item for item in value | |
| ): | |
| raise OnnxRuntimeError(f"bundle {name} is invalid") | |
| if len(value) != len(set(value)): | |
| raise OnnxRuntimeError(f"bundle {name} contains duplicates") | |
| return tuple(value) | |
| def _threshold(value: Any) -> float: | |
| if isinstance(value, bool) or not isinstance(value, (int, float)): | |
| raise OnnxRuntimeError("bundle decision threshold is invalid") | |
| result = float(value) | |
| if not 0.0 <= result <= 1.0: | |
| raise OnnxRuntimeError("bundle decision threshold is invalid") | |
| return result | |
| def _numpy() -> Any: | |
| try: | |
| import numpy as np | |
| except ImportError as exc: # pragma: no cover - depends on optional install | |
| raise OnnxRuntimeUnavailableError("numpy is not installed") from exc | |
| return np | |
| def _tokenizer_path(self, *, suffix: str) -> Path: | |
| files = self.bundle.manifest.get("files", ()) | |
| if not isinstance(files, (list, tuple)): | |
| raise OnnxRuntimeError("bundle tokenizer inventory is invalid") | |
| candidates = [ | |
| item | |
| for item in files | |
| if isinstance(item, Mapping) | |
| and item.get("role") == "tokenizer" | |
| and isinstance(item.get("path"), str) | |
| and str(item["path"]).endswith(suffix) | |
| ] | |
| if len(candidates) != 1: | |
| raise OnnxRuntimeError("bundle tokenizer file is missing or ambiguous") | |
| try: | |
| path = (self.bundle.root / str(candidates[0]["path"])).resolve(strict=True) | |
| path.relative_to(self.bundle.root.resolve(strict=True)) | |
| except (OSError, ValueError): | |
| raise OnnxRuntimeError("bundle tokenizer path escapes the bundle") from None | |
| if not path.is_file(): | |
| raise OnnxRuntimeError("bundle tokenizer file is missing") | |
| return path | |
| def _require_tokenizer(self) -> Any: | |
| if self._tokenizer is not None: | |
| return self._tokenizer | |
| try: | |
| from tokenizers import Tokenizer | |
| except ImportError as exc: # pragma: no cover - depends on optional install | |
| raise OnnxRuntimeUnavailableError("tokenizers is not installed") from exc | |
| tokenizer_path = self._tokenizer_path(suffix="tokenizer.json") | |
| config_path = self._tokenizer_path(suffix="tokenizer_config.json") | |
| try: | |
| config = json.loads(config_path.read_text(encoding="utf-8")) | |
| raw_tokenizer = Tokenizer.from_file(str(tokenizer_path)) | |
| pad_token = self._special_token_content(config.get("pad_token")) | |
| cls_token = self._special_token_content(config.get("cls_token")) | |
| pad_token_id = raw_tokenizer.token_to_id(pad_token) if pad_token else None | |
| cls_token_id = raw_tokenizer.token_to_id(cls_token) if cls_token else None | |
| except Exception: # noqa: BLE001 - sanitize third-party errors | |
| raise OnnxRuntimeError("bundle tokenizer could not be loaded") from None | |
| if pad_token_id is None or cls_token_id is None: | |
| raise OnnxRuntimeError("bundle tokenizer is missing required special tokens") | |
| self._tokenizer = _LocalTokenizer(raw_tokenizer, pad_token_id, cls_token_id) | |
| return self._tokenizer | |
| def _special_token_content(value: Any) -> str | None: | |
| if isinstance(value, str): | |
| return value | |
| if isinstance(value, Mapping) and isinstance(value.get("content"), str): | |
| return value["content"] | |
| return None | |
| def _validate_session_contract(self, session: Any) -> None: | |
| try: | |
| input_names = {item.name for item in session.get_inputs()} | |
| output_names = {item.name for item in session.get_outputs()} | |
| except Exception: # noqa: BLE001 - sanitize third-party errors | |
| raise OnnxRuntimeError("ONNX session metadata is unavailable") from None | |
| if not self._input_names or input_names != set(self._input_names): | |
| raise OnnxRuntimeError("ONNX graph inputs disagree with the verified manifest") | |
| if not self._output_names or output_names != set(self._output_names): | |
| raise OnnxRuntimeError("ONNX graph outputs disagree with the verified manifest") | |
| expected_inputs = { | |
| "input_ids", | |
| "attention_mask", | |
| "token_type_ids", | |
| "option_mask", | |
| } | |
| if self.backend_type is Backend.DIFFUSION: | |
| expected_inputs.add("initial_noise") | |
| if input_names != expected_inputs: | |
| raise OnnxRuntimeError("ONNX graph input contract is unsupported") | |
| score_name = "logits" if self.backend_type is Backend.DIRECT else "scores" | |
| if output_names != {score_name, "answerability"}: | |
| raise OnnxRuntimeError("ONNX graph output contract is unsupported") | |
| def prepare(self) -> None: | |
| """Load all local runtime dependencies before a client request is sent.""" | |
| self._numpy() | |
| self._require_tokenizer() | |
| self._validate_session_contract(self._require_session()) | |
| def _require_session(self) -> Any: | |
| if self._session is not None: | |
| return self._session | |
| try: | |
| import onnxruntime as ort # type: ignore[import-not-found] | |
| except ImportError as exc: # pragma: no cover - depends on optional install | |
| raise OnnxRuntimeUnavailableError("onnxruntime is not installed") from exc | |
| graph_item = None | |
| for item in self.bundle.manifest.get("files", ()): | |
| if isinstance(item, Mapping) and item.get("role") == "model_graph": | |
| graph_item = item | |
| break | |
| if graph_item is None: | |
| raise OnnxRuntimeError("bundle has no model_graph entry") | |
| try: | |
| graph_path = (self.bundle.root / str(graph_item["path"])).resolve(strict=True) | |
| graph_path.relative_to(self.bundle.root.resolve(strict=True)) | |
| except (OSError, ValueError): | |
| raise OnnxRuntimeError("bundle graph path is invalid") from None | |
| if not graph_path.is_file(): | |
| raise OnnxRuntimeError("bundle graph file is missing") | |
| try: | |
| self._session = ort.InferenceSession( | |
| str(graph_path), providers=["CPUExecutionProvider"] | |
| ) | |
| self._validate_session_contract(self._session) | |
| except OnnxRuntimeError: | |
| raise | |
| except Exception: # noqa: BLE001 - sanitize third-party errors | |
| raise OnnxRuntimeError("ONNX session could not be created") from None | |
| return self._session | |
| def decide(self, request: DecisionRequest) -> DecisionResponse: | |
| self._validate_request(request) | |
| if request.backend is not self.backend_type: | |
| raise OnnxRuntimeError("request backend does not match the verified bundle") | |
| try: | |
| from .inference import prepare_question_tensors | |
| np = self._numpy() | |
| tokenizer = self._require_tokenizer() | |
| session = self._require_session() | |
| self._validate_session_contract(session) | |
| score_name = "logits" if self.backend_type is Backend.DIRECT else "scores" | |
| answers: list[QuestionAnswer] = [] | |
| for question_index, question in enumerate(request.questions): | |
| padded, _ = prepare_question_tensors( | |
| tokenizer, | |
| request, | |
| question, | |
| option_count=self.option_count, | |
| max_sequence_budget=self.sequence_length, | |
| ) | |
| feed: dict[str, Any] = { | |
| "input_ids": np.asarray(padded.input_ids, dtype=np.int64)[None, :, :], | |
| "attention_mask": np.asarray(padded.attention_mask, dtype=np.int64)[None, :, :], | |
| "token_type_ids": np.asarray(padded.token_type_ids, dtype=np.int64)[None, :, :], | |
| "option_mask": np.asarray(padded.option_mask, dtype=np.bool_)[None, :], | |
| } | |
| if self.backend_type is Backend.DIFFUSION: | |
| seed = None if request.seed is None else request.seed + question_index | |
| rng = np.random.default_rng(seed) | |
| feed["initial_noise"] = rng.standard_normal( | |
| (1, padded.allocated_candidates), dtype=np.float32 | |
| ) | |
| outputs = session.run(list(self._output_names), feed) | |
| if not isinstance(outputs, (list, tuple)) or len(outputs) != len( | |
| self._output_names | |
| ): | |
| raise OnnxRuntimeError("ONNX graph returned an invalid output list") | |
| output_map = dict(zip(self._output_names, outputs, strict=True)) | |
| scores = np.asarray(output_map[score_name], dtype=np.float64) | |
| raw_answerability = np.asarray(output_map["answerability"], dtype=np.float64) | |
| live_options = question.options or ("true", "false") | |
| if scores.shape != (1, padded.allocated_candidates): | |
| raise OnnxRuntimeError("ONNX graph returned invalid score dimensions") | |
| if raw_answerability.size != 1: | |
| raise OnnxRuntimeError("ONNX graph returned invalid answerability") | |
| live_scores = scores[0, : len(live_options)] | |
| if not np.isfinite(live_scores).all() or not np.isfinite(raw_answerability).all(): | |
| raise OnnxRuntimeError("ONNX graph returned non-finite scores") | |
| shifted = live_scores - np.max(live_scores) | |
| exp_scores = np.exp(shifted) | |
| probabilities = exp_scores / exp_scores.sum() | |
| answerability_logit = float(raw_answerability.reshape(-1)[0]) | |
| if answerability_logit >= 0: | |
| answerability = 1.0 / (1.0 + np.exp(-answerability_logit)) | |
| else: | |
| exp_logit = np.exp(answerability_logit) | |
| answerability = exp_logit / (1.0 + exp_logit) | |
| choice_index = int(np.argmax(probabilities)) | |
| max_probability = float(probabilities[choice_index]) | |
| confidence = max_probability * float(answerability) | |
| abstain = ( | |
| confidence < self._abstain_threshold | |
| or answerability < self._answerability_threshold | |
| ) | |
| answers.append( | |
| QuestionAnswer( | |
| question_id=question.id, | |
| choice=None if abstain else live_options[choice_index], | |
| probabilities=tuple( | |
| OptionProbability( | |
| option=option, probability=float(probabilities[index]) | |
| ) | |
| for index, option in enumerate(live_options) | |
| ), | |
| confidence=confidence, | |
| status=ResponseStatus.ABSTAIN if abstain else ResponseStatus.OK, | |
| abstain_reason=( | |
| "below configured local decision threshold" if abstain else None | |
| ), | |
| ) | |
| ) | |
| return DecisionResponse( | |
| answers=tuple(answers), | |
| backend=self.backend_type, | |
| model_id=self.bundle.model_id, | |
| metadata={"verified_bundle_sha256": self.bundle.model_id}, | |
| ) | |
| except OnnxRuntimeError: | |
| raise | |
| except Exception: # noqa: BLE001 - sanitize tokenizer/ORT failures | |
| raise OnnxRuntimeError("local ONNX inference failed") from None | |
| class _LocalEncoding: | |
| ids: list[int] | |
| attention_mask: list[int] | |
| token_type_ids: list[int] | |
| class _LocalTokenizer: | |
| """Small adapter from the bundled Rust tokenizer to Vons tensor preparation.""" | |
| def __init__(self, tokenizer: Any, pad_token_id: int, cls_token_id: int) -> None: | |
| self._tokenizer = tokenizer | |
| self.pad_token_id = pad_token_id | |
| self.cls_token_id = cls_token_id | |
| def encode( | |
| self, | |
| text: str, | |
| *, | |
| add_special_tokens: bool = True, | |
| truncation: bool = False, | |
| max_length: int = 512, | |
| **_: Any, | |
| ) -> _LocalEncoding: | |
| if not isinstance(text, str): | |
| raise TypeError("tokenizer input must be text") | |
| if truncation: | |
| self._tokenizer.enable_truncation(max_length=max_length) | |
| try: | |
| encoded = self._tokenizer.encode(text, add_special_tokens=add_special_tokens) | |
| finally: | |
| if truncation: | |
| self._tokenizer.no_truncation() | |
| return _LocalEncoding( | |
| ids=list(encoded.ids), | |
| attention_mask=list(encoded.attention_mask), | |
| token_type_ids=list(encoded.type_ids), | |
| ) | |