Cleanup

hotg-ai · Jan 19, 2021 · 0ac23e2 · 0ac23e2
1 parent 7c6a45a
commit 0ac23e2
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 17 deletions.
diff --git a/rune/src/run/vm.rs b/rune/src/run/vm.rs
@@ -114,14 +114,9 @@ pub fn tfm_preload_model(
     inputs: u32,
     outputs: u32
 ) -> u32 {
-    log::info!("Calling tfm_preload_model");
-
-
+    let provider: &mut Provider = unsafe { &mut *(ctx.data as *mut Provider) };
     let model_bytes = get_mem_array(ctx, model_idx, model_len);
-
-    log::info!("BYTES = {:?}", model_bytes);
-
-    return 0;
+    return provider.add_model(model_bytes, inputs, outputs);
 }
 
 pub fn tfm_model_invoke(
@@ -208,7 +203,7 @@ pub fn tfm_model_invoke(
 
 pub fn _debug(ctx: &mut Ctx, ptr:  WasmPtr<u8, Array>, len: u32) -> u32 {
 
-    log::info!("RUNE::DEBUG {}", get_mem_str(ctx, ptr, len));
+    log::info!("[Rune::Debug]{}", get_mem_str(ctx, ptr, len));
     return 0;
 }
 

diff --git a/rune/src/run/vm/provider.rs b/rune/src/run/vm/provider.rs
@@ -6,30 +6,50 @@ use runic_types::*;
 use crate::run::vm::capability::*;
 use std::cell::RefCell;
 
+use tflite::ops::builtin::BuiltinOpResolver;
+use tflite::{FlatBufferModel, InterpreterBuilder};
 
 pub struct Model {
-    weights: Vec<u8>
+    fb: FlatBufferModel,
+    inputs: u32,
+    outputs: u32
 }
 
 pub struct Provider {
     requests: RefCell<Vec<RefCell<CapabilityRequest>>>,
-    models: Vec<Model>
+    models: RefCell<Vec<RefCell<Model>>>
 
 }
 
 impl Provider {
     pub fn init() -> Provider{
         return Provider{
             requests: RefCell::new(Vec::new()),
-            models: vec![]
+            models: RefCell::new(vec![])
         };
     }
 
-    // pub fn request_capability(&mut self, requested: runic_types::CAPABILITY) -> u32 {
-    //     let idx = self.requests.len() as u32;
-    //     self.requests.push(CapabilityRequest::init(requested));
-    //     return idx;
-    // }
+    pub fn add_model(&mut self, model_weights: Vec<u8>, inputs:u32, outputs:u32) -> u32 {
+       let idx = self.models.borrow().len() as u32;
+
+       let fb = match FlatBufferModel::build_from_buffer(model_weights) {
+           Ok(fb) => fb,
+           Err(err) => {
+               log::error!("Invalid model provided {:?}", err);
+               panic!("Invalid model");
+           }
+       };
+
+       let model = Model{
+           fb,
+           inputs,
+           outputs
+       };
+
+       self.models.borrow_mut().push(RefCell::new(model));
+       log::info!("Setting Model<{},{}>({})", inputs, outputs, idx);
+       return idx;
+    }
 
 
     pub fn request_capability(&mut self, requested: u32) -> u32 {
@@ -48,7 +68,7 @@ impl Provider {
         let capability_request = match self.requests.borrow_mut().get(request_idx as usize)  {
             Some(cr) => {
                 cr.borrow_mut().set_param(key, value, value_t);
-                log::info!("Parameter set for Capability {:?}", cr.borrow());
+                log::info!("Setting capability({}) params", request_idx);
             },
             _ => {
                 log::warn!("Rune called to set param on capability_request({}) that does not exist", request_idx);