| """Inference backend wrapper: pyaxengine (default) or onnxruntime (host stand-in). |
| |
| Both expose an onnxruntime-like API (run / get_inputs / get_outputs). The |
| backend is selected from the model file suffix unless overridden: |
| |
| *.axmodel -> axengine.InferenceSession (pyaxengine, on-device or host sim) |
| *.onnx -> onnxruntime.InferenceSession (host numeric stand-in only) |
| |
| Encoder input dtype note: the compiled AXMODELs declare their token inputs as |
| S32 at runtime (COMPILE_NOTES §3 folds the S64 input_processors into the |
| model; SIMULATE §1 confirmed S32 binaries), while the ONNX graphs take int64. |
| The wrapper therefore adapts the feed dtype from the session's input metadata. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| class ModelSession: |
| """Minimal name-addressable inference session over axengine/onnxruntime.""" |
|
|
| def __init__(self, model_path: str | Path, backend: str = "auto") -> None: |
| self.path = Path(model_path) |
| if not self.path.is_file(): |
| raise FileNotFoundError(self.path) |
| if backend == "auto": |
| backend = "axengine" if self.path.suffix == ".axmodel" else "onnxruntime" |
| self.backend = backend |
| if backend == "axengine": |
| try: |
| from axengine import InferenceSession |
| except ImportError as exc: |
| raise ImportError( |
| "pyaxengine is required for .axmodel files " |
| "(pip install axengine on an AXera host/BSP environment)" |
| ) from exc |
| self._sess = InferenceSession(str(self.path)) |
| elif backend == "onnxruntime": |
| try: |
| import onnxruntime as ort |
| except ImportError as exc: |
| raise ImportError( |
| "onnxruntime is required for .onnx files " |
| "(host numeric stand-in only)" |
| ) from exc |
| self._sess = ort.InferenceSession( |
| str(self.path), providers=["CPUExecutionProvider"] |
| ) |
| else: |
| raise ValueError(f"unknown backend {backend!r}") |
|
|
| |
| def _inputs(self): |
| return self._sess.get_inputs() |
|
|
| def output_names(self) -> list[str]: |
| return [o.name for o in self._sess.get_outputs()] |
|
|
| def input_dtype(self, name: str) -> np.dtype: |
| """Best-effort numpy dtype of a named input. |
| |
| onnxruntime exposes `.type` strings ('tensor(int64)'); pyaxengine |
| exposes numpy-like `.dtype`. Fallback: int32 for .axmodel (see module |
| docstring), int64 for .onnx. |
| """ |
| for meta in self._inputs(): |
| if meta.name != name: |
| continue |
| dt = getattr(meta, "dtype", None) |
| if dt is not None: |
| try: |
| return np.dtype(dt) |
| except TypeError: |
| pass |
| type_str = str(getattr(meta, "type", "") or "") |
| if "int32" in type_str: |
| return np.dtype(np.int32) |
| if "int64" in type_str: |
| return np.dtype(np.int64) |
| if "float" in type_str: |
| return np.dtype(np.float32) |
| return np.dtype(np.int32 if self.path.suffix == ".axmodel" else np.int64) |
|
|
| |
| def run(self, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: |
| results = self._sess.run(None, feeds) |
| return dict(zip(self.output_names(), results)) |
|
|
