Move loader to standalone crate llama-loader

Cargo.toml

@@ -2,6 +2,7 @@
 members = [
     "ggml-sys",
     "ggml",
+    "llama-loader",
     "llama-rs",
     "llama-cli",
     "generate-ggml-bindings"

ggml/src/lib.rs

@@ -440,7 +440,7 @@ impl Tensor {
     // /// Set the tensor's data pointer (useful for mmap-ed data)
     // ///
     // /// # Safety
-    // /// 
+    // ///
     // /// The memory region from `data_ptr` to `data_ptr.offset(tensor.nbytes())` will be read from.
     // pub unsafe fn set_data(&self, data_ptr: *mut c_void) {
     //     self.with_alive_ctx(|| {

llama-loader/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "llama-loader"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+ggml = { path = "../ggml" }
+thiserror = "*"

llama-rs/src/loader2.rs → llama-loader/src/lib.rs

@@ -1,15 +1,31 @@
-#![allow(missing_docs)]
+//! standalone model loader
+//!
+//! Only the hyperparameter is llama-specific. Everything else can be reused for other LLM.
+#![allow(clippy::nonminimal_bool)]
 
-//! standalone model loader 
+pub mod util;
 
 use std::{
     io::{BufRead, Seek, SeekFrom},
     ops::ControlFlow,
 };
+use util::*;
+
+pub type ElementType = ggml::Type;
+
+/// file type containing the model
+#[derive(Debug, PartialEq, Clone, Copy)]
+#[allow(clippy::upper_case_acronyms)]
+pub enum ContainerType {
+    /// legacy format, oldest ggml tensor file format
+    GGML,
+    /// also legacy format, newer than GGML, older than GGJT
+    GGMF,
+    /// mmap-able format
+    GGJT,
+}
 
-use crate::{loader::has_data_left, ElementType, ModelContainerType};
-
-pub(crate) fn decode_element_type(ftype: i32) -> Option<ElementType> {
+pub fn decode_element_type(ftype: i32) -> Option<ElementType> {
     match ftype {
         0 => Some(ggml::Type::F32),
         1 => Some(ggml::Type::F16),
@@ -19,7 +35,7 @@ pub(crate) fn decode_element_type(ftype: i32) -> Option<ElementType> {
     }
 }
 
-pub(crate) fn encode_element_type(element_type: ElementType) -> Option<i32> {
+pub fn encode_element_type(element_type: ElementType) -> Option<i32> {
     match element_type {
         ggml::Type::F32 => Some(0),
         ggml::Type::F16 => Some(1),
@@ -29,38 +45,9 @@ pub(crate) fn encode_element_type(element_type: ElementType) -> Option<i32> {
     }
 }
 
-pub(crate) fn read_bytes<const N: usize>(
-    reader: &mut impl BufRead,
-) -> Result<[u8; N], std::io::Error> {
-    let mut bytes = [0u8; N];
-    reader.read_exact(&mut bytes)?;
-    Ok(bytes)
-}
-
-pub(crate) fn read_i32(reader: &mut impl BufRead) -> Result<i32, std::io::Error> {
-    Ok(i32::from_le_bytes(read_bytes::<4>(reader)?))
-}
-
-pub(crate) fn read_u32(reader: &mut impl BufRead) -> Result<u32, std::io::Error> {
-    Ok(u32::from_le_bytes(read_bytes::<4>(reader)?))
-}
-
-pub(crate) fn read_f32(reader: &mut impl BufRead) -> Result<f32, std::io::Error> {
-    Ok(f32::from_le_bytes(read_bytes::<4>(reader)?))
-}
-
-pub(crate) fn read_bytes_with_len(
-    reader: &mut impl BufRead,
-    len: usize,
-) -> Result<Vec<u8>, std::io::Error> {
-    let mut bytes = vec![0u8; len];
-    reader.read_exact(&mut bytes)?;
-    Ok(bytes)
-}
-
 /// The hyperparameters of the model.
 #[derive(Debug, Clone)]
-pub struct FixedHyperparameters {
+pub struct LlamaHyperparameters {
     pub n_vocab: usize,
     pub n_embd: usize,
     pub n_mult: usize,
@@ -90,7 +77,7 @@ pub enum LoadError<T> {
 
     #[error("unsupported tensor dtype/f16_: {0}")]
     UnsupportedElementtype(i32),
-    
+
     /// sanity check failed
     #[error("invariant broken: {0}")]
     InvariantBroken(String),
@@ -107,11 +94,11 @@ pub struct TensorInfo {
 
 #[allow(unused_variables)]
 pub trait LoadHandler<T> {
-    fn cb_container_type(&mut self, model_type: ModelContainerType) -> ControlFlow<T> {
+    fn cb_container_type(&mut self, model_type: ContainerType) -> ControlFlow<T> {
         ControlFlow::Continue(())
     }
 
-    fn cb_hyper_parameters(&mut self, hparams: FixedHyperparameters) -> ControlFlow<T> {
+    fn cb_hyper_parameters(&mut self, hparams: LlamaHyperparameters) -> ControlFlow<T> {
         ControlFlow::Continue(())
     }
 
@@ -134,30 +121,30 @@ pub fn load_model_from_reader<T>(
     handler: &mut impl LoadHandler<T>,
 ) -> Result<(), LoadError<T>> {
     // Verify magic
-    let container_type: ModelContainerType = match read_u32(&mut reader)? {
-        ggml::FILE_MAGIC_GGMF => ModelContainerType::GGMF,
-        ggml::FILE_MAGIC_GGJT => ModelContainerType::GGJT,
-        ggml::FILE_MAGIC_UNVERSIONED => ModelContainerType::Unversioned,
+    let container_type: ContainerType = match read_u32(&mut reader)? {
+        ggml::FILE_MAGIC_GGMF => ContainerType::GGMF,
+        ggml::FILE_MAGIC_GGJT => ContainerType::GGJT,
+        ggml::FILE_MAGIC_UNVERSIONED => ContainerType::GGML,
         magic => return Err(LoadError::InvalidMagic(magic)),
     };
     retchk(handler.cb_container_type(container_type))?;
 
     // Load format version
     match container_type {
-        ModelContainerType::GGMF | ModelContainerType::GGJT => {
+        ContainerType::GGMF | ContainerType::GGJT => {
             let _version: u32 = match read_u32(&mut reader)? {
                 ggml::FORMAT_VERSION => ggml::FORMAT_VERSION,
                 version => return Err(LoadError::InvalidFormatVersion(version)),
             };
         }
-        ModelContainerType::Unversioned => {}
+        ContainerType::GGML => {}
     }
 
     // Load hyper params
     //
     // NOTE: Field order matters! Data is laid out in the file exactly
     // in this order.
-    let hparams = FixedHyperparameters {
+    let hparams = LlamaHyperparameters {
         n_vocab: read_i32(&mut reader)?.try_into()?,
         n_embd: read_i32(&mut reader)?.try_into()?,
         n_mult: read_i32(&mut reader)?.try_into()?,
@@ -174,8 +161,8 @@ pub fn load_model_from_reader<T>(
         let len = read_u32(&mut reader)?.try_into()?;
         let token = read_bytes_with_len(&mut reader, len)?;
         let token_score = match container_type {
-            ModelContainerType::GGMF | ModelContainerType::GGJT => read_f32(&mut reader)?,
-            ModelContainerType::Unversioned => {
+            ContainerType::GGMF | ContainerType::GGJT => read_f32(&mut reader)?,
+            ContainerType::GGML => {
                 // Legacy model, set empty score
                 0.
             }
@@ -185,12 +172,12 @@ pub fn load_model_from_reader<T>(
 
     // Load tensor data
     match container_type {
-        ModelContainerType::GGMF | ModelContainerType::Unversioned => {
+        ContainerType::GGMF | ContainerType::GGML => {
             let _file_offset = reader.stream_position()?;
             drop(reader);
             todo!()
         }
-        ModelContainerType::GGJT => load_weights_ggjt(&mut reader, handler),
+        ContainerType::GGJT => load_weights_ggjt(&mut reader, handler),
     }
 }
 
@@ -238,23 +225,28 @@ fn load_weights_ggjt<T>(
         }
 
         let tensor_info = TensorInfo {
-            name, dims, n_dims, n_elements, ftype,
+            name,
+            dims,
+            n_dims,
+            n_elements,
+            ftype,
         };
 
         // load tensor weights
         let offset_curr = reader.stream_position()?;
         let offset_aligned: u64 = (offset_curr + 31) & !31;
         reader.seek(SeekFrom::Start(offset_aligned))?;
-    
+
         let type_size = ggml::type_size(ftype);
         let buf = retchk(handler.tensor_buffer(tensor_info))?;
         let buf_len = buf.len();
         if !(buf_len == type_size * n_elements) {
-            return Err(LoadError::InvariantBroken(format!("{buf_len} == {type_size} * {n_elements}")));
+            return Err(LoadError::InvariantBroken(format!(
+                "{buf_len} == {type_size} * {n_elements}"
+            )));
         }
         reader.read_exact(buf)?;
     }
 
     Ok(())
 }
-

llama-loader/src/util.rs

@@ -0,0 +1,33 @@
+use std::io::BufRead;
+
+pub fn read_bytes<const N: usize>(reader: &mut impl BufRead) -> Result<[u8; N], std::io::Error> {
+    let mut bytes = [0u8; N];
+    reader.read_exact(&mut bytes)?;
+    Ok(bytes)
+}
+
+pub fn read_i32(reader: &mut impl BufRead) -> Result<i32, std::io::Error> {
+    Ok(i32::from_le_bytes(read_bytes::<4>(reader)?))
+}
+
+pub fn read_u32(reader: &mut impl BufRead) -> Result<u32, std::io::Error> {
+    Ok(u32::from_le_bytes(read_bytes::<4>(reader)?))
+}
+
+pub fn read_f32(reader: &mut impl BufRead) -> Result<f32, std::io::Error> {
+    Ok(f32::from_le_bytes(read_bytes::<4>(reader)?))
+}
+
+pub fn read_bytes_with_len(
+    reader: &mut impl BufRead,
+    len: usize,
+) -> Result<Vec<u8>, std::io::Error> {
+    let mut bytes = vec![0u8; len];
+    reader.read_exact(&mut bytes)?;
+    Ok(bytes)
+}
+
+// NOTE: Implementation from #![feature(buf_read_has_data_left)]
+pub fn has_data_left(reader: &mut impl BufRead) -> Result<bool, std::io::Error> {
+    reader.fill_buf().map(|b| !b.is_empty())
+}

llama-rs/Cargo.toml

@@ -8,6 +8,7 @@ rust-version = "1.65"
 
 [dependencies]
 ggml = { path = "../ggml" }
+llama-loader = { path = "../llama-loader" }
 
 bytemuck = "1.13.1"
 partial_sort = "0.2.0"

llama-rs/src/convert.rs

@@ -16,7 +16,8 @@ use std::{
     vec,
 };
 
-use crate::{util, Hyperparameters, Vocabulary, loader2::encode_element_type};
+use crate::{util, Hyperparameters, Vocabulary};
+use llama_loader::encode_element_type;
 
 /// Converts a `pth` file to a `ggml` file.
 pub fn convert_pth_to_ggml(model_directory: &Path, element_type: ggml::Type) {