Update to burn v0.14.0 and switch to .mpk model file

src/model/load.rs

@@ -1,5 +1,7 @@
 use npy::{self, NpyData};
 use num_traits::cast::ToPrimitive;
+use burn::tensor::cast::ToElement;
+use burn::prelude::TensorData;
 use std::error::Error;
 use std::io::Read;
 
@@ -21,7 +23,8 @@ pub fn numpy_to_tensor<B: Backend, const D: usize>(
     let shape: Vec<_> = v[0..D].into_iter().map(|&v| v as usize).collect();
     let data: Vec<B::FloatElem> = v[D..].into_iter().map(|e| e.elem()).collect();
 
-    Tensor::from_data_device(Data::new(data, shape.into()), device)
+    //Tensor::from_data_device(Data::new(data, shape.into()), device)
+    Tensor::from_data(TensorData::new(data, shape), device)
 }
 
 pub fn load_tensor<B: Backend, const D: usize>(
@@ -48,7 +51,7 @@ pub fn load_f32<B: Backend>(
     path: &str,
     device: &B::Device,
 ) -> Result<f32, Box<dyn Error>> {
-    load_tensor::<B, 1>(name, path, device).map(|t| t.into_scalar().to_f32().unwrap())
+    load_tensor::<B, 1>(name, path, device).map(|t| t.into_scalar().to_f32())
 }
 
 pub fn load_usize<B: Backend>(
@@ -56,7 +59,7 @@ pub fn load_usize<B: Backend>(
     path: &str,
     device: &B::Device,
 ) -> Result<usize, Box<dyn Error>> {
-    load_tensor::<B, 1>(name, path, device).map(|t| t.into_scalar().to_usize().unwrap())
+    load_tensor::<B, 1>(name, path, device).map(|t| t.into_scalar().to_usize())
 }
 
 pub fn load_linear<B: Backend>(
@@ -66,13 +69,10 @@ pub fn load_linear<B: Backend>(
     let weight = load_tensor::<B, 2>("weight", path, device)?;
     let bias = load_tensor::<B, 1>("bias", path, device).ok();
 
-    let record = nn::LinearRecord {
-        weight: weight.into(),
-        bias: bias.map(|t| t.into()),
-    };
-
-    let linear: nn::Linear<B> = nn::LinearConfig::new(3, 3).init_with(record);
-    Ok(linear)
+    Ok(nn::Linear {
+        weight: Param::from_tensor(weight), 
+        bias: bias.map(|t| Param::from_tensor(t)), 
+    })
 }
 
 pub fn load_embedding<B: Backend>(
@@ -80,14 +80,10 @@ pub fn load_embedding<B: Backend>(
     device: &B::Device,
 ) -> Result<nn::Embedding<B>, Box<dyn Error>> {
     let weight = load_tensor::<B, 2>("weight", path, device)?;
-    let [n_vocab, n_state] = weight.dims();
 
-    let record = nn::EmbeddingRecord {
-        weight: weight.into(),
-    };
-
-    let embedding = nn::EmbeddingConfig::new(n_vocab, n_state).init_with(record);
-    Ok(embedding)
+    Ok(nn::Embedding {
+        weight: Param::from_tensor(weight), 
+    })
 }
 
 pub fn load_layer_norm<B: Backend>(
@@ -100,13 +96,9 @@ pub fn load_layer_norm<B: Backend>(
 
     let [n_state] = weight.dims();
 
-    let record = nn::LayerNormRecord {
-        gamma: weight.into(),
-        beta: bias.into(),
-        epsilon: <f64 as Module<B>>::into_record(eps),
-    };
-
-    let layer_norm: nn::LayerNorm<B> = nn::LayerNormConfig::new(n_state).init_with(record);
+    let mut layer_norm = nn::LayerNormConfig::new(n_state).with_epsilon(eps).init(device);
+    layer_norm.gamma = Param::from_tensor(weight);
+    layer_norm.beta = Param::from_tensor(bias);
 
     Ok(layer_norm)
 }
@@ -116,7 +108,7 @@ pub fn load_layer_norm<B: Backend>(
     let eps = load_f32::<B>("eps", path, device)?.into();
 
     let rmsnorm =  RMSNorm {
-        weight: weight.into(),
+        weight: Param::from_tensor(weight),
         eps: eps
     };
 
@@ -148,40 +140,38 @@ pub fn load_conv2d<B: Backend>(
     let padding = tensor_to_array_2(padding);
     let padding = nn::PaddingConfig2d::Explicit(padding[0], padding[1]);
 
-    let record = conv::Conv2dRecord {
-        weight: weight.into(),
-        bias: bias.map(|t| t.into()),
-        stride: <[usize; 2] as Module<B>>::into_record(stride),
-        kernel_size: <[usize; 2] as Module<B>>::into_record(kernel_size),
-        dilation: <[usize; 2] as Module<B>>::into_record(dilation),
-        groups: <usize as Module<B>>::into_record(n_group),
-        padding: <nn::PaddingConfig2d as Module<B>>::into_record(padding.clone()),
-    };
-
-    let conv2d: conv::Conv2d<B> =
-        conv::Conv2dConfig::new([n_channels_in, n_channels_out], kernel_size)
+    let mut conv2d = conv::Conv2dConfig::new([n_channels_in, n_channels_out], kernel_size)
             .with_stride(stride)
             .with_dilation(dilation)
             .with_groups(n_group)
-            .with_padding(padding)
+            .with_padding(padding.clone())
             .with_bias(has_bias)
-            .init_with(record);
+            .init(device);
+
+    conv2d.weight = Param::from_tensor(weight);
+    conv2d.bias = bias.map(|t| Param::from_tensor(t));
+    conv2d.stride = stride;
+    conv2d.kernel_size = kernel_size;
+    conv2d.dilation = dilation;
+    conv2d.groups = n_group;
+    conv2d.padding = burn::module::Ignored(padding);
+        
     Ok(conv2d)
 }
 
 pub fn tensor_to_array_2<B: Backend>(x: Tensor<B, 1>) -> [usize; 2] {
-    let vec = x.into_data().value;
+    let vec: Vec<<B as Backend>::FloatElem> = x.into_data().to_vec().unwrap();
     assert!(vec.len() == 2, "Tensor length must be 2.");
-    [vec[0].to_usize().unwrap(), vec[1].to_usize().unwrap()]
+    [vec[0].to_usize(), vec[1].to_usize()]
 }
 
 pub fn tensor_to_array<const N: usize, B: Backend>(x: Tensor<B, 1>) -> [usize; N] {
-    let vec = x.into_data().value;
+    let vec: Vec<<B as Backend>::FloatElem> = x.into_data().to_vec().unwrap();
     assert!(vec.len() == N, "Tensor length must be {}.", N);
 
     let mut arr = [0; N];
     for (a, t) in arr.iter_mut().zip(vec) {
-        *a = t.to_usize().unwrap();
+        *a = t.to_usize();
     }
 
     arr