allow models without scalars to be read without errors (#808)

Co-authored-by: Corey Lowman <clowman1993@gmail.com>
coreylowman · Jul 12, 2023 · e45a308 · e45a308
1 parent b0fbaad
commit e45a308
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 25 deletions.
diff --git a/src/nn/npz.rs b/src/nn/npz.rs
@@ -1,4 +1,5 @@
 use crate::{
+    prelude::numpy::NpyError,
     shapes::{Dtype, Shape},
     tensor::{
         numpy::{read_from_npz, write_to_npz, NpzError, NumpyDtype},
@@ -162,18 +163,26 @@ impl<R: Read + Seek, E: Dtype + NumpyDtype, D: Device<E>> TensorVisitor<E, D>
 
     fn visit_scalar<N: num_traits::NumCast>(
         &mut self,
-        _opts: ScalarOptions<N>,
+        opts: ScalarOptions<N>,
         (n, full_path): (&mut N, String),
     ) -> Result<Option<N>, Self::Err> {
-        let buf: Vec<f64> = read_from_npz(self, &[], full_path)?;
-        *n = N::from(buf[0]).unwrap_or_else(|| {
-            panic!(
-                "Failed to convert f64 value {} to {} when reading from npz!",
-                buf[0],
-                std::any::type_name::<N>()
-            )
-        });
-        Ok(None)
+        match read_from_npz::<_, f64>(self, &[], full_path) {
+            Ok(buf) => {
+                *n = N::from(buf[0]).unwrap_or_else(|| {
+                    panic!(
+                        "Failed to convert f64 value {} to {} when reading from npz!",
+                        buf[0],
+                        std::any::type_name::<N>()
+                    )
+                });
+                Ok(None)
+            }
+            Err(NpyError::IoError(e)) if e.kind() == std::io::ErrorKind::NotFound => {
+                *n = opts.default;
+                Ok(None)
+            }
+            Err(x) => Err(x.into()),
+        }
     }
 }
 

diff --git a/src/nn/safetensors.rs b/src/nn/safetensors.rs
@@ -171,20 +171,29 @@ impl<'data, E: Dtype + SafeDtype, D: Device<E>> TensorVisitor<E, D> for SafeTens
 
     fn visit_scalar<N: num_traits::NumCast>(
         &mut self,
-        _: ScalarOptions<N>,
+        opts: ScalarOptions<N>,
         (n, full_path): (&mut N, String),
     ) -> Result<Option<N>, Self::Err> {
-        let data = self.tensor(&full_path)?.data();
-        let mut array = [0; 8];
-        array.copy_from_slice(data);
-        let val = f64::from_le_bytes(array);
-        *n = N::from(val).unwrap_or_else(|| {
-            panic!(
-                "Failed to convert f64 value {val} at {full_path} to {} when reading from safetensors!",
-                std::any::type_name::<N>()
-            )
-        });
-        Ok(None)
+        match self.tensor(&full_path) {
+            Ok(tensor) => {
+                let data = tensor.data();
+                let mut array = [0; 8];
+                array.copy_from_slice(data);
+                let val = f64::from_le_bytes(array);
+                *n = N::from(val).unwrap_or_else(|| {
+                    panic!(
+                        "Failed to convert f64 value {val} at {full_path} to {} when reading from safetensors!",
+                        std::any::type_name::<N>()
+                    )
+                });
+                Ok(None)
+            }
+            Err(SafeTensorError::TensorNotFound(_)) => {
+                *n = opts.default;
+                Ok(None)
+            }
+            Err(x) => Err(Error::SafeTensorError(x)),
+        }
     }
 }
 

diff --git a/src/tensor/numpy.rs b/src/tensor/numpy.rs
@@ -76,9 +76,16 @@ pub(crate) fn read_from_npz<R: Read + Seek, E: Dtype + NumpyDtype>(
         filename.push_str(".npy");
     }
 
-    let mut f = r
-        .by_name(&filename)
-        .unwrap_or_else(|_| panic!("'{filename}' not found"));
+    let mut f = match r.by_name(&filename) {
+        Ok(f) => f,
+        Err(ZipError::FileNotFound) => {
+            return Err(NpyError::IoError(io::Error::new(
+                io::ErrorKind::NotFound,
+                ZipError::FileNotFound,
+            )))
+        }
+        Err(e) => panic!("Uncaught zip error: {e}"),
+    };
 
     read_from_npy(&mut f, shape)
 }