ronn/
session.rs

1//! Session class for Python bindings
2
3use crate::error::{PyResult, RonnError};
4use numpy::PyArray1;
5use pyo3::prelude::*;
6use pyo3::types::PyDict;
7use ronn_api::InferenceSession;
8use std::collections::HashMap;
9
10/// Inference session
11///
12/// # Example
13///
14/// ```python
15/// import numpy as np
16///
17/// inputs = {
18///     "input": np.array([[1.0, 2.0, 3.0]], dtype=np.float32)
19/// }
20/// outputs = session.run(inputs)
21/// print(outputs["output"])
22/// ```
23#[pyclass(name = "Session")]
24pub struct PySession {
25    inner: InferenceSession,
26}
27
28impl PySession {
29    pub fn new(session: InferenceSession) -> Self {
30        Self { inner: session }
31    }
32}
33
34#[pymethods]
35impl PySession {
36    /// Run inference
37    ///
38    /// # Arguments
39    ///
40    /// * `inputs` - Dictionary mapping input names to numpy arrays
41    ///
42    /// # Returns
43    ///
44    /// Dictionary mapping output names to numpy arrays
45    ///
46    /// # Example
47    ///
48    /// ```python
49    /// outputs = session.run({"input": input_array})
50    /// result = outputs["output"]
51    /// ```
52    fn run(&self, py: Python, inputs: &PyDict) -> PyResult<PyObject> {
53        // Convert Python dict to HashMap<String, Tensor>
54        let mut input_tensors: HashMap<String, ronn_core::Tensor> = HashMap::new();
55
56        for (key, value) in inputs.iter() {
57            let name: String = key.extract()?;
58
59            // Handle different numpy array types
60            let tensor = if let Ok(array) = value.downcast::<PyArray1<f32>>() {
61                let data: Vec<f32> = array.to_vec()?;
62                let shape = vec![array.len()];
63                ronn_core::Tensor::from_data(
64                    data,
65                    shape,
66                    ronn_core::DataType::F32,
67                    ronn_core::TensorLayout::RowMajor,
68                )
69                .map_err(RonnError::from)?
70            } else {
71                return Err(RonnError(format!("Unsupported input type for '{}'", name)));
72            };
73
74            input_tensors.insert(name, tensor);
75        }
76
77        // Convert HashMap<String, Tensor> to HashMap<&str, Tensor>
78        let inputs_ref: HashMap<&str, ronn_core::Tensor> = input_tensors
79            .iter()
80            .map(|(k, v)| (k.as_str(), v.clone()))
81            .collect();
82
83        // Run inference
84        let output_tensors = self.inner.run(inputs_ref).map_err(RonnError::from)?;
85
86        // Convert back to Python dict
87        let result = PyDict::new(py);
88        for (name, tensor) in output_tensors {
89            // Convert tensor to numpy array
90            let data: Vec<f32> = tensor.to_vec().unwrap_or_else(|_| vec![]);
91            let array = PyArray1::from_vec(py, data);
92            result.set_item(name, array)?;
93        }
94
95        Ok(result.into())
96    }
97
98    fn __repr__(&self) -> String {
99        "Session(...)".to_string()
100    }
101}
ronn/session.rs

ronn/
session.rs
