update snapshot to match wasmtime-wasi-nn v36 (#154)

Author: ancaantochi

samples/wasm/rust/examples/snapshot/Cargo.lock

@@ -7,7 +7,7 @@ edition = "2021"
 
 [dependencies]
 ndarray = "0.16.1"
-wit-bindgen = "0.22"
+wit-bindgen = "0.32"
 wasm_graph_sdk = { version = "=1.1.1", registry = "aio-sdks" }
 serde = { version = "1", features = [
 	"derive",
@@ -16,7 +16,6 @@ serde = { version = "1", features = [
 serde_json = { version = "1.0", default-features = false, features = [
 	"alloc",
 ] }
-wasi-nn = {git = "https://github.com/bytecodealliance/wasi-nn.git"}
 
 [lib]
 crate-type = ["cdylib"]

samples/wasm/rust/examples/snapshot/Cargo.toml

@@ -4,6 +4,11 @@
 // Generated by `wit_bindgen::generate` expansion.
 #![allow(clippy::missing_safety_doc)]
 
+wit_bindgen::generate!({
+    path: "wit",
+    world: "ml",
+});
+
 mod map_snapshot {
     use std::sync::LazyLock;
 
@@ -23,7 +28,7 @@ mod map_snapshot {
     use ndarray::{Array, Dim};
     use std::io::BufRead;
 
-    use wasi_nn::{
+    use crate::wasi::nn::{
         graph::{load, ExecutionTarget, Graph, GraphEncoding, GraphExecutionContext},
         tensor::{Tensor, TensorData, TensorDimensions, TensorType},
     };
@@ -76,21 +81,14 @@ mod map_snapshot {
         let tensor = Tensor::new(&dimensions, TensorType::Fp32, &data);
 
         unsafe {
-            match (*exec_context).set_input("data", tensor) {
-                Ok(()) => {}
-                Err(e) => {
-                    logger::log(
-                        Level::Error,
-                        "module-snapshot/map",
-                        &format!("Error setting input tensor: {e:?}"),
-                    );
-                    panic!("Error setting input tensor: {:?}", e);
-                }
-            }
-
             // Execute the inferencing
-            match (*exec_context).compute() {
-                Ok(()) => {}
+            let output_data = match (*exec_context).compute(vec![("data".to_owned(), tensor)]) {
+                Ok(result) => result
+                    .iter()
+                    .find(|p| p.0 == MODEL_OUTPUT_NAME)
+                    .unwrap()
+                    .1
+                    .data(),
                 Err(e) => {
                     logger::log(
                         Level::Error,
@@ -99,13 +97,8 @@ mod map_snapshot {
                     );
                     panic!("Error executing graph inference: {e:?}");
                 }
-            }
+            };
 
-            // Get the inferencing result (bytes) and convert it to f32
-            let output_data = (*exec_context)
-                .get_output(MODEL_OUTPUT_NAME)
-                .unwrap()
-                .data();
             let output_f32 = bytes_to_f32_vec(&output_data);
 
             let output_shape = [1, 1000, 1, 1];