Merge pull request #58 from marshallpierce/mp/string-type

Add String data type support

Author: nbigaouette

onnxruntime-sys/build.rs

@@ -105,7 +105,10 @@ fn generate_bindings(include_dir: &Path) {
         .expect("Couldn't write bindings!");
 }
 
-fn download<P: AsRef<Path>>(source_url: &str, target_file: P) {
+fn download<P>(source_url: &str, target_file: P)
+where
+    P: AsRef<Path>,
+{
     let resp = ureq::get(source_url)
         .timeout_connect(1_000) // 1 second
         .timeout(std::time::Duration::from_secs(300))

onnxruntime/examples/print_structure.rs

@@ -0,0 +1,39 @@
+//! Display the input and output structure of an ONNX model.
+use onnxruntime::environment;
+use std::error::Error;
+
+fn main() -> Result<(), Box<dyn Error>> {
+    // provide path to .onnx model on disk
+    let path = std::env::args()
+        .skip(1)
+        .next()
+        .expect("Must provide an .onnx file as the first arg");
+
+    let environment = environment::Environment::builder()
+        .with_name("onnx metadata")
+        .with_log_level(onnxruntime::LoggingLevel::Verbose)
+        .build()?;
+
+    let session = environment
+        .new_session_builder()?
+        .with_optimization_level(onnxruntime::GraphOptimizationLevel::Basic)?
+        .with_model_from_file(path)?;
+
+    println!("Inputs:");
+    for (index, input) in session.inputs.iter().enumerate() {
+        println!(
+            "  {}:\n    name = {}\n    type = {:?}\n    dimensions = {:?}",
+            index, input.name, input.input_type, input.dimensions
+        )
+    }
+
+    println!("Outputs:");
+    for (index, output) in session.outputs.iter().enumerate() {
+        println!(
+            "  {}:\n    name = {}\n    type = {:?}\n    dimensions = {:?}",
+            index, output.name, output.output_type, output.dimensions
+        );
+    }
+
+    Ok(())
+}

onnxruntime/src/error.rs

@@ -56,9 +56,15 @@ pub enum OrtError {
     /// Error occurred when creating CPU memory information
     #[error("Failed to get dimensions: {0}")]
     CreateCpuMemoryInfo(OrtApiError),
+    /// Error occurred when creating ONNX tensor
+    #[error("Failed to create tensor: {0}")]
+    CreateTensor(OrtApiError),
     /// Error occurred when creating ONNX tensor with specific data
     #[error("Failed to create tensor with data: {0}")]
     CreateTensorWithData(OrtApiError),
+    /// Error occurred when filling a tensor with string data
+    #[error("Failed to fill string tensor: {0}")]
+    FillStringTensor(OrtApiError),
     /// Error occurred when checking if ONNX tensor was properly initialized
     #[error("Failed to check if tensor: {0}")]
     IsTensor(OrtApiError),
@@ -183,3 +189,11 @@ pub(crate) fn status_to_result(
     let status_wrapper: OrtStatusWrapper = status.into();
     status_wrapper.into()
 }
+
+/// A wrapper around a function on OrtApi that maps the status code into [OrtApiError]
+pub(crate) unsafe fn call_ort<F>(mut f: F) -> std::result::Result<(), OrtApiError>
+where
+    F: FnMut(sys::OrtApi) -> *const sys::OrtStatus,
+{
+    status_to_result(f(g_ort()))
+}

onnxruntime/src/lib.rs

@@ -343,8 +343,8 @@ pub enum TensorElementDataType {
     Int32 = sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 as OnnxEnumInt,
     /// Signed 64-bit int, equivalent to Rust's `i64`
     Int64 = sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 as OnnxEnumInt,
-    // /// String, equivalent to Rust's `String`
-    // String = sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING as OnnxEnumInt,
+    /// String, equivalent to Rust's `String`
+    String = sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING as OnnxEnumInt,
     // /// Boolean, equivalent to Rust's `bool`
     // Bool = sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL as OnnxEnumInt,
     // /// 16-bit floating point, equivalent to Rust's `f16`
@@ -374,9 +374,7 @@ impl Into<sys::ONNXTensorElementDataType> for TensorElementDataType {
             Int16 => sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16,
             Int32 => sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32,
             Int64 => sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,
-            // String => {
-            //     sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING
-            // }
+            String => sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING,
             // Bool => {
             //     sys::ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL
             // }
@@ -402,15 +400,22 @@ impl Into<sys::ONNXTensorElementDataType> for TensorElementDataType {
 /// Trait used to map Rust types (for example `f32`) to ONNX types (for example `Float`)
 pub trait TypeToTensorElementDataType {
     /// Return the ONNX type for a Rust type
-    fn tensor_element_data_type() -> sys::ONNXTensorElementDataType;
+    fn tensor_element_data_type() -> TensorElementDataType;
+
+    /// If the type is `String`, returns `Some` with utf8 contents, else `None`.
+    fn try_utf8_bytes(&self) -> Option<&[u8]>;
 }
 
 macro_rules! impl_type_trait {
     ($type_:ty, $variant:ident) => {
         impl TypeToTensorElementDataType for $type_ {
-            fn tensor_element_data_type() -> sys::ONNXTensorElementDataType {
+            fn tensor_element_data_type() -> TensorElementDataType {
                 // unsafe { std::mem::transmute(TensorElementDataType::$variant) }
-                TensorElementDataType::$variant.into()
+                TensorElementDataType::$variant
+            }
+
+            fn try_utf8_bytes(&self) -> Option<&[u8]> {
+                None
             }
         }
     };
@@ -423,7 +428,6 @@ impl_type_trait!(u16, Uint16);
 impl_type_trait!(i16, Int16);
 impl_type_trait!(i32, Int32);
 impl_type_trait!(i64, Int64);
-// impl_type_trait!(String, String);
 // impl_type_trait!(bool, Bool);
 // impl_type_trait!(f16, Float16);
 impl_type_trait!(f64, Double);
@@ -433,6 +437,39 @@ impl_type_trait!(u64, Uint64);
 // impl_type_trait!(, Complex128);
 // impl_type_trait!(, Bfloat16);
 
+/// Adapter for common Rust string types to Onnx strings.
+///
+/// It should be easy to use both `String` and `&str` as [TensorElementDataType::String] data, but
+/// we can't define an automatic implementation for anything that implements `AsRef<str>` as it
+/// would conflict with the implementations of [TypeToTensorElementDataType] for primitive numeric
+/// types (which might implement `AsRef<str>` at some point in the future).
+pub trait Utf8Data {
+    /// Returns the utf8 contents.
+    fn utf8_bytes(&self) -> &[u8];
+}
+
+impl Utf8Data for String {
+    fn utf8_bytes(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+
+impl<'a> Utf8Data for &'a str {
+    fn utf8_bytes(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+
+impl<T: Utf8Data> TypeToTensorElementDataType for T {
+    fn tensor_element_data_type() -> TensorElementDataType {
+        TensorElementDataType::String
+    }
+
+    fn try_utf8_bytes(&self) -> Option<&[u8]> {
+        Some(self.utf8_bytes())
+    }
+}
+
 /// Allocator type
 #[derive(Debug, Clone)]
 #[repr(i32)]

onnxruntime/src/session.rs

@@ -411,7 +411,9 @@ impl<'a> Session<'a> {
         // The C API expects pointers for the arrays (pointers to C-arrays)
         let input_ort_tensors: Vec<OrtTensor<TIn, D>> = input_arrays
             .into_iter()
-            .map(|input_array| OrtTensor::from_array(&self.memory_info, input_array))
+            .map(|input_array| {
+                OrtTensor::from_array(&self.memory_info, self.allocator_ptr, input_array)
+            })
             .collect::<Result<Vec<OrtTensor<TIn, D>>>>()?;
         let input_ort_values: Vec<*const sys::OrtValue> = input_ort_tensors
             .iter()