From b3241414e86ea069fa85667a756b6755da7ded4a Mon Sep 17 00:00:00 2001 From: rsasaki0109 Date: Fri, 7 Aug 2026 20:09:38 +0900 Subject: [PATCH 1/2] Add Python ONNX entity embedder (Epic 150C) Expose OnnxEntityEmbedder in the Python extension: it embeds point-entity feature vectors through an existing OnnxRuntimeSession and returns a NumPy embedding plus its dimension, gated by the onnxruntime wheel feature. The double_dynamic.onnx fixture proves the real ONNX path doubles [1,2,3] to [2,4,6]. Add typed stub and wheel-gate test. --- CHANGELOG.md | 5 ++ crates/spatialrust-py/Cargo.toml | 2 +- crates/spatialrust-py/spatialrust.pyi | 18 ++++ crates/spatialrust-py/src/lib.rs | 95 ++++++++++++++++++++ crates/spatialrust-py/tests/test_bindings.py | 31 +++++++ docs/FEATURE_MATRIX.md | 2 +- docs/ROADMAP.md | 2 +- notes/2026-08-07_epic150_ai_semantic.md | 13 ++- 8 files changed, 162 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 890eb1b..f90e7ee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,11 @@ removed no sooner than the next major (see `docs/API_STABILITY.md`). ### Added +- **Epic 150C Python entity embedder**: `OnnxEntityEmbedder` in the Python + extension embeds point-entity feature vectors through an existing + `OnnxRuntimeSession` and returns a NumPy embedding vector with its dimension, + gated by the `onnxruntime` wheel feature. The `double_dynamic.onnx` fixture + proves the real ONNX path doubles `[1,2,3]` → `[2,4,6]`. - **Epic 150 real AI semantic meaning** (`semantic-model`): `OnnxEntityEmbedder` runs point-entity feature vectors through an already-open ONNX model session with explicit copy policy and produces an `Embedding` ready for the existing diff --git a/crates/spatialrust-py/Cargo.toml b/crates/spatialrust-py/Cargo.toml index b63c287..3598fa8 100644 --- a/crates/spatialrust-py/Cargo.toml +++ b/crates/spatialrust-py/Cargo.toml @@ -11,7 +11,7 @@ description = "Python bindings for SpatialRust point cloud processing" [features] default = [] -onnxruntime = ["spatialrust/ai-onnxruntime"] +onnxruntime = ["spatialrust/ai-onnxruntime", "spatialrust/semantic-model"] [lib] # Importable module name: `import spatialrust` diff --git a/crates/spatialrust-py/spatialrust.pyi b/crates/spatialrust-py/spatialrust.pyi index 2d5b43b..6084836 100644 --- a/crates/spatialrust-py/spatialrust.pyi +++ b/crates/spatialrust-py/spatialrust.pyi @@ -986,6 +986,24 @@ class OnnxRuntimeSession: def outputs(self) -> list[tuple[str, str, list[str]]]: ... def run(self, inputs: dict[str, Tensor], *, copy: bool = ...) -> dict[str, Tensor]: ... +class OnnxEntityEmbedder: + """Embeds point-entity features through an ONNX session (Epic 150).""" + + def __new__( + cls, + session: OnnxRuntimeSession, + input_name: str, + output_name: str, + input_shape: list[int], + output_shape: list[int], + copy: bool = ..., + ) -> OnnxEntityEmbedder: ... + def embed( + self, session: OnnxRuntimeSession, features: NDArray[np.float32] + ) -> dict[str, object]: + """Returns {\"embedding\": ndarray, \"dim\": int}.""" + ... + def tensor_copy_from_numpy(array: NDArray[np.generic]) -> Tensor: ... @final class DLPackTensorView: diff --git a/crates/spatialrust-py/src/lib.rs b/crates/spatialrust-py/src/lib.rs index d86e8dc..fdcd7e7 100644 --- a/crates/spatialrust-py/src/lib.rs +++ b/crates/spatialrust-py/src/lib.rs @@ -355,6 +355,100 @@ fn python_tensor_spec(spec: &spatialrust::ai::TensorSpec) -> (String, String, Ve (spec.name.clone(), tensor_dtype_name(spec.dtype), dimensions) } +/// Embeds point-entity features through an ONNX session into a searchable +/// embedding (Epic 150). +/// +/// Args: +/// session: spatialrust.OnnxRuntimeSession +/// input_name: str +/// output_name: str +/// input_shape: list[int] +/// output_shape: list[int] +/// copy: bool (default True) — permit documented host copies +#[pyclass(name = "OnnxEntityEmbedder", unsendable)] +struct PyOnnxEntityEmbedder { + #[cfg(feature = "onnxruntime")] + inner: spatialrust::semantic::OnnxEntityEmbedder, +} + +#[pymethods] +impl PyOnnxEntityEmbedder { + #[new] + #[pyo3(signature = (session, input_name, output_name, input_shape, output_shape, copy=true))] + fn new( + session: &Bound<'_, PyAny>, + input_name: String, + output_name: String, + input_shape: Vec, + output_shape: Vec, + copy: bool, + ) -> PyResult { + #[cfg(feature = "onnxruntime")] + { + let _ = session; // session ownership is retained by the caller; embed() receives it. + let copy_policy = if copy { AiCopyPolicy::Allow } else { AiCopyPolicy::Forbid }; + let descriptor = |shape: Vec| { + spatialrust::tensor::TensorDescriptor::contiguous( + spatialrust::tensor::DataType::F32, + shape, + spatialrust::tensor::Device::CPU, + ) + }; + let inner = spatialrust::semantic::OnnxEntityEmbedder::try_new( + input_name, + output_name, + descriptor(input_shape), + descriptor(output_shape), + copy_policy, + ) + .map_err(to_py_err)?; + Ok(Self { inner }) + } + #[cfg(not(feature = "onnxruntime"))] + { + let _ = (session, input_name, output_name, input_shape, output_shape, copy); + Err(PyRuntimeError::new_err( + "this SpatialRust Python module was built without the `onnxruntime` feature", + )) + } + } + + /// Embeds one feature vector; `features` length must match the input shape. + fn embed<'py>( + &self, + py: Python<'py>, + session: PyRefMut<'_, PyOnnxRuntimeSession>, + features: Vec, + ) -> PyResult> { + #[cfg(feature = "onnxruntime")] + { + let mut session = session; + let inner: &mut dyn ModelSession = session.inner.as_mut(); + let embedding = self + .inner + .embed_one(inner, &features) + .map_err(to_py_err)?; + let dims = embedding.dim(); + let values = embedding.as_slice().to_vec(); + let array = numpy::IntoPyArray::into_pyarray_bound( + numpy::ndarray::Array1::from(values), + py, + ); + let dict = PyDict::new_bound(py); + dict.set_item("embedding", array)?; + dict.set_item("dim", dims)?; + Ok(dict) + } + #[cfg(not(feature = "onnxruntime"))] + { + let _ = (py, session, features); + Err(PyRuntimeError::new_err( + "this SpatialRust Python module was built without the `onnxruntime` feature", + )) + } + } +} + #[pyclass(name = "DLPackTensorView", unsendable)] struct PyDlpackTensorView { inner: DlpackImport, @@ -4610,6 +4704,7 @@ fn spatialrust_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/spatialrust-py/tests/test_bindings.py b/crates/spatialrust-py/tests/test_bindings.py index 537a817..ebd6b81 100644 --- a/crates/spatialrust-py/tests/test_bindings.py +++ b/crates/spatialrust-py/tests/test_bindings.py @@ -1293,3 +1293,34 @@ def test_point_cloud_stream_arrow_c_stream_batches(tmp_path): # Concatenated x column matches the source. concat = pa.concat_arrays([batch.column("x") for batch in batches]) np.testing.assert_allclose(np.array(concat), xs, rtol=1e-5) + + +def test_onnx_entity_embedder_runs_real_model(tmp_path): + model = bytes( + [ + 8, 8, 18, 16, 115, 112, 97, 116, 105, 97, 108, 114, 117, 115, 116, 45, 116, + 101, 115, 116, 58, 106, 10, 27, 10, 5, 105, 110, 112, 117, 116, 10, 5, 105, + 110, 112, 117, 116, 18, 6, 111, 117, 116, 112, 117, 116, 34, 3, 65, 100, 100, + 18, 14, 100, 111, 117, 98, 108, 101, 95, 100, 121, 110, 97, 109, 105, 99, 90, + 28, 10, 5, 105, 110, 112, 117, 116, 18, 19, 10, 17, 8, 1, 18, 13, 10, 7, + 18, 5, 98, 97, 116, 99, 104, 10, 2, 8, 3, 98, 29, 10, 6, 111, 117, 116, + 112, 117, 116, 18, 19, 10, 17, 8, 1, 18, 13, 10, 7, 18, 5, 98, 97, 116, + 99, 104, 10, 2, 8, 3, 66, 4, 10, 0, 16, 13, + ] + ) + path = tmp_path / "double_dynamic.onnx" + path.write_bytes(model) + try: + session = sr.OnnxRuntimeSession(str(path), deterministic=True) + except RuntimeError as error: + if "without the `onnxruntime` feature" in str(error): + pytest.skip("extension was intentionally built without ONNX Runtime") + raise + + embedder = sr.OnnxEntityEmbedder( + session, "input", "output", [1, 3], [1, 3], copy=True + ) + result = embedder.embed(session, np.array([1.0, 2.0, 3.0], dtype=np.float32)) + # The fixture doubles its input. + np.testing.assert_allclose(result["embedding"], np.array([2.0, 4.0, 6.0])) + assert result["dim"] == 3 diff --git a/docs/FEATURE_MATRIX.md b/docs/FEATURE_MATRIX.md index 86b6444..c8cdbd3 100644 --- a/docs/FEATURE_MATRIX.md +++ b/docs/FEATURE_MATRIX.md @@ -39,7 +39,7 @@ workspace because its build requires a Python toolchain. | `spatialrust-pipeline` | MVP pipeline | GPU MVP stages | algorithm crates only | | `spatialrust-interchange` | `interchange-gltf`, `interchange-openusd` | `tiles3d`: deterministic OGC 3D Tiles 1.1 `tileset.json` + `pnts` octree export; `tiles3d-copc`: bounded COPC hierarchy → tileset | `tiles3d-copc` pulls `spatialrust-io` + `spatialrust-core` for COPC node reads | | `spatialrust-arrow` | no capability by default | C Data `__arrow_c_array__`, C Stream `__arrow_c_stream__`, C Device `__arrow_c_device_array__` exports/imports for point records | `arrow-c-data`, `arrow-c-stream`, `arrow-c-device`; no Arrow runtime dependency | -| `spatialrust-py` | Python binding surface | selected meta-crate features, including `export_tiles3d` / `export_copc_tiles3d`, Arrow C Data `__arrow_c_array__` / `__arrow_c_stream__` | PyO3/NumPy | +| `spatialrust-py` | Python binding surface | selected meta-crate features, including `export_tiles3d` / `export_copc_tiles3d`, Arrow C Data `__arrow_c_array__` / `__arrow_c_stream__`, `OnnxEntityEmbedder` (onnxruntime) | PyO3/NumPy | ## Execution contract diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 0e830c1..d75f3ee 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -139,7 +139,7 @@ ONNX model committed as fixture bytes so no external download is required. | --- | --- | --- | --- | | 150A | Complete | `OnnxEntityEmbedder`: entity feature tensors → ONNX embedding via an existing `ModelSession`, with explicit copy policy | `semantic`, `ai-onnxruntime` | | 150B | Complete | Fixture ONNX model bytes + embedder correctness tests against the mock/real boundary | fixture, tests | -| 150C | Active | Facade wiring, Python binding, FEATURE_MATRIX/CHANGELOG/notes | facade | +| 150C | Complete | Facade wiring, Python binding, FEATURE_MATRIX/CHANGELOG/notes | facade | The embedder never loads a model itself; it consumes an already-open session so backend identity, device placement, and copy policy stay explicit. The fixture diff --git a/notes/2026-08-07_epic150_ai_semantic.md b/notes/2026-08-07_epic150_ai_semantic.md index 76d2d4c..53bcf88 100644 --- a/notes/2026-08-07_epic150_ai_semantic.md +++ b/notes/2026-08-07_epic150_ai_semantic.md @@ -1,6 +1,6 @@ # Epic 150: real AI semantic meaning via ONNX entity embeddings -Date: 2026-08-07. Slices 150A/150B complete. +Date: 2026-08-07. Slices 150A–150C complete. ## Why @@ -26,6 +26,11 @@ data model. - real ONNX Runtime CPU session embeds `[1,2,3]` → `[2,4,6]` exactly; - the embedding is inserted into `SemanticSearchIndex` and found by a query embedding, proving the full feature → model → search path. +- Python binding (150C): `PyOnnxEntityEmbedder` takes an existing + `OnnxRuntimeSession`, input/output names and shapes, and `copy` policy; + `embed()` returns a NumPy embedding plus its dimension. Gated by the + `onnxruntime` wheel feature; a `double_dynamic.onnx`-style Python test + doubles `[1,2,3]` → `[2,4,6]`. ## Contract notes @@ -44,5 +49,7 @@ data model. ## Next slices -150C: Python binding for the embedder and open-vocabulary search over indexed -entities, plus a note with an example using a real public embedding model. +Epic 150 is complete. A follow-up note can document using a real public +open-vocabulary embedding model (e.g. CLIP image/text) with this surface; the +wiring and tests already cover the full feature → model → embedding → search +path in both Rust and Python. From 234c6ec2ab040e9450684d5065b9cba4c646425c Mon Sep 17 00:00:00 2001 From: rsasaki0109 Date: Fri, 7 Aug 2026 20:15:42 +0900 Subject: [PATCH 2/2] Mark OnnxEntityEmbedder final in the type stub --- crates/spatialrust-py/spatialrust.pyi | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/spatialrust-py/spatialrust.pyi b/crates/spatialrust-py/spatialrust.pyi index 6084836..939331f 100644 --- a/crates/spatialrust-py/spatialrust.pyi +++ b/crates/spatialrust-py/spatialrust.pyi @@ -986,6 +986,7 @@ class OnnxRuntimeSession: def outputs(self) -> list[tuple[str, str, list[str]]]: ... def run(self, inputs: dict[str, Tensor], *, copy: bool = ...) -> dict[str, Tensor]: ... +@final class OnnxEntityEmbedder: """Embeds point-entity features through an ONNX session (Epic 150)."""