Use functions from the new crate

llama-loader/src/lib.rs

@@ -109,7 +109,7 @@ pub trait LoadHandler<T> {
     }
 
     /// # Returns
-    /// 
+    ///
     /// `None` to skip copying
     /// `Some(buf)` to provide a buffer for copying weights into
     fn get_tensor_buffer(&mut self, info: TensorInfo) -> ControlFlow<T, Option<&mut [u8]>> {
@@ -242,10 +242,9 @@ fn load_weights_ggjt<T>(
             n_dims,
             n_elements,
             ftype,
-            start_offset: offset_aligned
+            start_offset: offset_aligned,
         };
 
-
         let type_size = ggml::type_size(ftype);
         if let Some(buf) = retchk(handler.get_tensor_buffer(tensor_info))? {
             reader.seek(SeekFrom::Start(offset_aligned))?;
@@ -258,7 +257,9 @@ fn load_weights_ggjt<T>(
             reader.read_exact(buf)?;
         } else {
             // skip if no buffer is given
-            reader.seek(SeekFrom::Start(offset_aligned + (type_size * n_elements) as u64))?;
+            reader.seek(SeekFrom::Start(
+                offset_aligned + (type_size * n_elements) as u64,
+            ))?;
         }
     }
 

llama-rs/src/lib.rs

@@ -24,6 +24,7 @@ use thiserror::Error;
 #[cfg(feature = "mmap")]
 use memmap2::Mmap;
 
+use llama_loader::util::*;
 use llama_loader::{decode_element_type, ContainerType};
 
 /// dummy struct
@@ -609,48 +610,6 @@ impl Model {
         })?;
         let mut reader = BufReader::new(&file);
 
-        fn read_bytes<const N: usize>(reader: &mut impl BufRead) -> Result<[u8; N], LoadError> {
-            let mut bytes = [0u8; N];
-            reader
-                .read_exact(&mut bytes)
-                .map_err(|e| LoadError::ReadExactFailed {
-                    source: e,
-                    bytes: N,
-                })?;
-            Ok(bytes)
-        }
-
-        fn read_bytes_with_len(
-            reader: &mut impl BufRead,
-            len: usize,
-        ) -> Result<Vec<u8>, LoadError> {
-            let mut bytes = vec![0u8; len];
-            reader
-                .read_exact(&mut bytes)
-                .map_err(|e| LoadError::ReadExactFailed {
-                    source: e,
-                    bytes: len,
-                })?;
-            Ok(bytes)
-        }
-
-        fn read_i32(reader: &mut impl BufRead) -> Result<i32, LoadError> {
-            Ok(i32::from_le_bytes(read_bytes::<4>(reader)?))
-        }
-
-        fn read_u32(reader: &mut impl BufRead) -> Result<u32, LoadError> {
-            Ok(u32::from_le_bytes(read_bytes::<4>(reader)?))
-        }
-
-        fn read_f32(reader: &mut impl BufRead) -> Result<f32, LoadError> {
-            Ok(f32::from_le_bytes(read_bytes::<4>(reader)?))
-        }
-
-        /// Helper function. Reads a string from the buffer and returns it.
-        fn read_string(reader: &mut BufReader<File>, len: usize) -> Result<String, LoadError> {
-            Ok(String::from_utf8(read_bytes_with_len(reader, len)?)?)
-        }
-
         // Verify magic
         let model_type: ContainerType = match read_u32(&mut reader)? {
             ggml::FILE_MAGIC_GGMF => ContainerType::GGMF,
@@ -710,7 +669,7 @@ impl Model {
 
             for i in 0..hparams.n_vocab {
                 let len = read_i32(&mut reader)?;
-                let token = read_bytes_with_len(&mut reader, len)?;
+                let token = read_bytes_with_len(&mut reader, len.try_into()?)?;
                 max_token_length = max_token_length.max(token.len());
                 id_to_token.push(token.clone());
                 token_to_id.insert(token, TokenId::try_from(i)?);

llama-rs/src/loader.rs

@@ -6,29 +6,7 @@ use std::{
 use crate::ElementType;
 use crate::{util, LoadError, LoadProgress, Model};
 use llama_loader::decode_element_type;
-
-pub(crate) fn read_bytes<const N: usize>(reader: &mut impl BufRead) -> Result<[u8; N], LoadError> {
-    let mut bytes = [0u8; N];
-    reader
-        .read_exact(&mut bytes)
-        .map_err(|e| LoadError::ReadExactFailed {
-            source: e,
-            bytes: N,
-        })?;
-    Ok(bytes)
-}
-
-pub(crate) fn read_i32(reader: &mut impl BufRead) -> Result<i32, LoadError> {
-    Ok(i32::from_le_bytes(read_bytes::<4>(reader)?))
-}
-
-pub(crate) fn read_u32(reader: &mut impl BufRead) -> Result<u32, LoadError> {
-    Ok(u32::from_le_bytes(read_bytes::<4>(reader)?))
-}
-
-pub(crate) fn read_f32(reader: &mut impl BufRead) -> Result<f32, LoadError> {
-    Ok(f32::from_le_bytes(read_bytes::<4>(reader)?))
-}
+use llama_loader::util::*;
 
 /// Helper function. Reads a string from the buffer and returns it.
 pub(crate) fn read_string(reader: &mut impl BufRead, len: usize) -> Result<String, LoadError> {
@@ -43,11 +21,6 @@ pub(crate) fn read_string(reader: &mut impl BufRead, len: usize) -> Result<Strin
     Ok(s)
 }
 
-// NOTE: Implementation from #![feature(buf_read_has_data_left)]
-pub(crate) fn has_data_left(reader: &mut impl BufRead) -> Result<bool, std::io::Error> {
-    reader.fill_buf().map(|b| !b.is_empty())
-}
-
 pub(crate) fn load_weights_ggmf_or_unversioned(
     file_offset: u64,
     main_path: &Path,