spatialrust_semantic/
model.rs

1//! Real-model entity embedding through an explicit model session.
2
3use spatialrust_ai::{CopyPolicy, ModelSession, NamedTensors, RunOptions};
4use spatialrust_tensor::{Device, TensorBuffer, TensorDescriptor};
5
6use crate::{Embedding, SemanticError, SemanticResult};
7
8/// Runs entity features through an already-open model session to produce an
9/// [`Embedding`].
10///
11/// The embedder never loads a model, chooses a backend, or moves data across a
12/// device boundary: the caller supplies an open session, names the input and
13/// output tensors, and selects the copy policy. This keeps backend identity and
14/// transfer semantics explicit and auditable.
15#[derive(Clone, Debug)]
16pub struct OnnxEntityEmbedder {
17    input_name: String,
18    output_name: String,
19    /// Input tensor descriptor (feature shape).
20    input_descriptor: TensorDescriptor,
21    /// Output tensor descriptor (embedding shape).
22    output_descriptor: TensorDescriptor,
23    copy_policy: CopyPolicy,
24}
25
26impl OnnxEntityEmbedder {
27    /// Creates an embedder for one model input/output pair.
28    pub fn try_new(
29        input_name: impl Into<String>,
30        output_name: impl Into<String>,
31        input_descriptor: TensorDescriptor,
32        output_descriptor: TensorDescriptor,
33        copy_policy: CopyPolicy,
34    ) -> SemanticResult<Self> {
35        if input_descriptor.device() != Device::CPU || output_descriptor.device() != Device::CPU {
36            return Err(SemanticError::InvalidConfiguration(
37                "entity embedder requires CPU-hosted input and output tensors".into(),
38            ));
39        }
40        if output_descriptor.shape().is_empty() || output_descriptor.shape().last() == Some(&0) {
41            return Err(SemanticError::InvalidConfiguration(
42                "embedding output must have a non-empty trailing dimension".into(),
43            ));
44        }
45        Ok(Self {
46            input_name: input_name.into(),
47            output_name: output_name.into(),
48            input_descriptor,
49            output_descriptor,
50            copy_policy,
51        })
52    }
53
54    /// Embeds one feature vector (flattened `f32` values) into an embedding.
55    ///
56    /// `features` must contain exactly the element count implied by
57    /// `input_descriptor`. The session's `output_name` tensor must match
58    /// `output_descriptor`.
59    pub fn embed_one(
60        &self,
61        session: &mut dyn ModelSession,
62        features: &[f32],
63    ) -> SemanticResult<Embedding> {
64        let expected = self.input_descriptor.element_count().map_err(|error| {
65            SemanticError::InvalidConfiguration(format!("input descriptor: {error}"))
66        })?;
67        if features.len() != expected {
68            return Err(SemanticError::InvalidConfiguration(format!(
69                "entity features have {} elements; expected {}",
70                features.len(),
71                expected
72            )));
73        }
74        let bytes = features_to_bytes(features)?;
75        let tensor = TensorBuffer::try_new(bytes, self.input_descriptor.clone())
76            .map_err(|error| SemanticError::InvalidConfiguration(error.to_string()))?;
77        let mut inputs = NamedTensors::new();
78        inputs
79            .insert(self.input_name.clone(), tensor)
80            .map_err(|error| SemanticError::InvalidConfiguration(error.to_string()))?;
81
82        let options = RunOptions { input_copy: self.copy_policy, output_copy: self.copy_policy };
83        let outputs = session
84            .run_with_options(inputs, options)
85            .map_err(|error| SemanticError::InvalidConfiguration(error.to_string()))?;
86        let output = outputs.get(&self.output_name).ok_or_else(|| {
87            SemanticError::InvalidConfiguration(format!(
88                "model output `{}` not found",
89                self.output_name
90            ))
91        })?;
92
93        let expected_output = self.output_descriptor.element_count().map_err(|error| {
94            SemanticError::InvalidConfiguration(format!("output descriptor: {error}"))
95        })?;
96        let output_shape = output.descriptor().shape();
97        if output_shape.iter().product::<usize>() != expected_output {
98            return Err(SemanticError::InvalidConfiguration(format!(
99                "model output `{}` has shape {output_shape:?}; expected {expected_output} elements",
100                self.output_name
101            )));
102        }
103        let bytes = output.allocation_bytes();
104        if bytes.len() != expected_output * 4 {
105            return Err(SemanticError::InvalidConfiguration(format!(
106                "model output `{}` has {} bytes; expected {}",
107                self.output_name,
108                bytes.len(),
109                expected_output * 4
110            )));
111        }
112        let mut values = Vec::with_capacity(expected_output);
113        for chunk in bytes.chunks_exact(4) {
114            values.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
115        }
116        Embedding::try_new(values)
117    }
118
119    /// Returns the configured input descriptor.
120    #[must_use]
121    pub fn input_descriptor(&self) -> &TensorDescriptor {
122        &self.input_descriptor
123    }
124
125    /// Returns the configured output descriptor.
126    #[must_use]
127    pub fn output_descriptor(&self) -> &TensorDescriptor {
128        &self.output_descriptor
129    }
130}
131
132fn features_to_bytes(features: &[f32]) -> SemanticResult<Vec<u8>> {
133    let mut bytes = Vec::with_capacity(features.len() * 4);
134    for value in features {
135        if !value.is_finite() {
136            return Err(SemanticError::InvalidConfiguration(
137                "entity features must contain finite values".into(),
138            ));
139        }
140        bytes.extend_from_slice(&value.to_le_bytes());
141    }
142    Ok(bytes)
143}
144
145#[cfg(test)]
146mod tests {
147    use super::OnnxEntityEmbedder;
148    use crate::Embedding;
149    use spatialrust_ai::{CopyPolicy, ModelInfo, ModelSession, NamedTensors, RunOptions};
150    use spatialrust_tensor::{DataType, Device, TensorDescriptor};
151
152    #[derive(Clone, Debug)]
153    struct IdentitySession {
154        info: ModelInfo,
155    }
156
157    impl Default for IdentitySession {
158        fn default() -> Self {
159            Self { info: ModelInfo { name: None, inputs: Vec::new(), outputs: Vec::new() } }
160        }
161    }
162
163    impl ModelSession for IdentitySession {
164        fn backend_name(&self) -> &str {
165            "test-identity"
166        }
167
168        fn model_info(&self) -> &ModelInfo {
169            &self.info
170        }
171
172        fn run_with_options(
173            &mut self,
174            inputs: NamedTensors,
175            _options: RunOptions,
176        ) -> spatialrust_ai::AiResult<NamedTensors> {
177            let mut outputs = NamedTensors::new();
178            for (name, tensor) in inputs.into_values() {
179                let output_name = if name == "input" { "output".to_owned() } else { name };
180                outputs.insert(output_name, tensor)?;
181            }
182            Ok(outputs)
183        }
184    }
185
186    fn descriptor(shape: &[usize]) -> TensorDescriptor {
187        TensorDescriptor::contiguous(DataType::F32, shape.to_vec(), Device::CPU)
188    }
189
190    #[test]
191    fn embeds_identity_features() {
192        let embedder = OnnxEntityEmbedder::try_new(
193            "input",
194            "output",
195            descriptor(&[1, 4]),
196            descriptor(&[1, 4]),
197            CopyPolicy::Allow,
198        )
199        .unwrap();
200        let mut session = IdentitySession::default();
201        let embedding = embedder.embed_one(&mut session, &[0.1, 0.2, 0.3, 0.4]).unwrap();
202        assert_eq!(embedding, Embedding::try_new(vec![0.1, 0.2, 0.3, 0.4]).unwrap());
203    }
204
205    #[test]
206    fn rejects_feature_count_mismatch() {
207        let embedder = OnnxEntityEmbedder::try_new(
208            "input",
209            "output",
210            descriptor(&[1, 4]),
211            descriptor(&[1, 4]),
212            CopyPolicy::Allow,
213        )
214        .unwrap();
215        let mut session = IdentitySession::default();
216        assert!(embedder.embed_one(&mut session, &[0.1, 0.2]).is_err());
217    }
218
219    #[test]
220    fn rejects_non_finite_features() {
221        let embedder = OnnxEntityEmbedder::try_new(
222            "input",
223            "output",
224            descriptor(&[1, 3]),
225            descriptor(&[1, 3]),
226            CopyPolicy::Allow,
227        )
228        .unwrap();
229        let mut session = IdentitySession::default();
230        assert!(embedder.embed_one(&mut session, &[f32::NAN, 0.0, 0.0]).is_err());
231    }
232
233    #[test]
234    fn rejects_device_mismatch() {
235        let gpu = TensorDescriptor::contiguous(
236            DataType::F32,
237            vec![1, 3],
238            Device { kind: spatialrust_tensor::DeviceKind::Cuda, id: 0 },
239        );
240        assert!(OnnxEntityEmbedder::try_new(
241            "input",
242            "output",
243            gpu,
244            descriptor(&[1, 3]),
245            CopyPolicy::Allow,
246        )
247        .is_err());
248    }
249
250    #[test]
251    fn rejects_zero_dim_output() {
252        assert!(OnnxEntityEmbedder::try_new(
253            "input",
254            "output",
255            descriptor(&[1, 3]),
256            descriptor(&[0]),
257            CopyPolicy::Allow,
258        )
259        .is_err());
260    }
261}
spatialrust_semantic/model.rs

spatialrust_semantic/
model.rs
