Merge pull request #1 from qinyiqun/AddGGuFSplit

add split

ggus/src/header.rs

@@ -1,7 +1,7 @@
 use crate::sizeof;
 use std::str::Utf8Error;
 
-#[derive(Default, Debug)]
+#[derive(Default, Debug, Clone)]
 #[repr(C)]
 pub struct GGufFileHeader {
     magic: [u8; 4],

ggus/src/metadata/mod.rs

@@ -1,4 +1,4 @@
-//! See <https://github.com/ggerganov/ggml/blob/master/docs/gguf.md#standardized-key-value-pairs>.
+//! See <https://github.com/ggerganov/ggml/blob/master/docs/gguf.md#standardized-key-value-pairs>.
 
 mod general;
 mod llm;
@@ -82,6 +82,14 @@ pub struct GGufMetaKVPairs<'a> {
 }
 
 impl<'a> GGufMetaKVPairs<'a> {
+    pub fn new(nbytes_: usize) -> Self {
+        let indices_ = IndexMap::new();
+        Self {
+            indices: indices_,
+            nbytes: nbytes_,
+        }
+    }
+
     pub fn scan(count: u64, data: &'a [u8]) -> Result<Self, GGufReadError<'a>> {
         let mut reader = GGufReader::new(data);
         let mut indices = IndexMap::with_capacity(count as _);
@@ -133,6 +141,16 @@ impl<'a> GGufMetaKVPairs<'a> {
             .map(|(&key, &len)| GGufMetaKV { key, len })
     }
 
+    pub fn remove(&mut self, _key: &str) -> bool {
+        match self.indices.swap_remove_entry(_key) {
+            Some((_, v)) => {
+                self.nbytes -= v;
+                true
+            }
+            None => false,
+        }
+    }
+
     fn get_typed(
         &self,
         name: impl AsRef<str>,

ggus/src/tensor.rs

@@ -93,6 +93,14 @@ pub struct GGufTensors<'a> {
 }
 
 impl<'a> GGufTensors<'a> {
+    pub fn new(nbytes_: usize) -> Self {
+        let indices_: IndexMap<&str, ()> = IndexMap::new();
+        Self {
+            indices: indices_,
+            nbytes: nbytes_,
+        }
+    }
+
     pub fn scan(count: u64, data: &'a [u8]) -> Result<Self, GGufReadError<'a>> {
         let mut reader = GGufReader::new(data);
         let mut indices = IndexMap::with_capacity(count as _);

ggus/src/writer.rs

@@ -10,9 +10,9 @@ pub struct GGufWriter<T: Write>(BufWriter<T>, usize);
 impl<T: Write> GGufWriter<T> {
     #[inline]
     pub fn new(writer: T, header: GGufFileHeader) -> Result<Self> {
-        let mut buf = BufWriter::new(writer);
-        buf.write_all(as_slice(&header))?;
-        Ok(Self(buf, 0))
+        let mut ans = Self(BufWriter::new(writer), 0);
+        ans.write_bytes(as_slice(&header))?;
+        Ok(ans)
     }
 
     #[inline]

xtask/src/gguf_file.rs

@@ -0,0 +1,73 @@
+use ggus::{GGufFileHeader, GGufMetaKVPairs, GGufReadError, GGufTensors};
+use indexmap::IndexMap;
+
+#[derive(Clone)]
+pub(crate) struct GGufFile<'a> {
+    header: GGufFileHeader,
+    meta_kvs: GGufMetaKVPairs<'a>,
+    tensors: GGufTensors<'a>,
+    data: &'a [u8],
+}
+
+#[derive(Debug)]
+pub(crate) enum GGufError<'a> {
+    MagicMismatch,
+    EndianNotSupport,
+    VersionNotSupport,
+    #[allow(dead_code)]
+    Reading(GGufReadError<'a>),
+    FileSizeError,
+    SplitModeRepeated,
+}
+
+impl<'a> GGufFile<'a> {
+    pub(crate) fn new(data: &'a [u8]) -> Result<Self, GGufError<'a>> {
+        let header = unsafe { data.as_ptr().cast::<GGufFileHeader>().read() };
+        if !header.is_magic_correct() {
+            return Err(GGufError::MagicMismatch);
+        }
+        if !header.is_native_endian() {
+            return Err(GGufError::EndianNotSupport);
+        }
+        if header.version != 3 {
+            return Err(GGufError::VersionNotSupport);
+        }
+
+        let cursor = header.nbytes();
+        let meta_kvs = GGufMetaKVPairs::scan(header.metadata_kv_count, &data[cursor..])
+            .map_err(GGufError::Reading)?;
+
+        let cursor = cursor + meta_kvs.nbytes();
+        let tensors =
+            GGufTensors::scan(header.tensor_count, &data[cursor..]).map_err(GGufError::Reading)?;
+
+        let align = meta_kvs.alignment();
+        let cursor = (cursor + tensors.nbytes() + align - 1) / align * align;
+        Ok(Self {
+            header,
+            meta_kvs,
+            tensors,
+            data: &data[cursor..],
+        })
+    }
+
+    pub fn header(&self) -> &GGufFileHeader {
+        &self.header
+    }
+
+    pub fn meta_kvs(&self) -> &GGufMetaKVPairs<'a> {
+        &self.meta_kvs
+    }
+
+    pub fn tensors_as_indexmap(&self) -> IndexMap<ggus::GGufTensorInfo, &[u8]> {
+        self.tensors
+            .iter()
+            .map(move |t| (t, self.data))
+            .collect::<IndexMap<_, _>>()
+    }
+}
+
+#[inline(always)]
+pub(crate) const fn pad(pos: usize, align: usize) -> usize {
+    (align - pos % align) % align
+}

xtask/src/main.rs

@@ -1,3 +1,4 @@
+mod gguf_file;
 mod loose_shards;
 mod merge;
 mod show;

xtask/src/merge.rs

@@ -1,7 +1,8 @@
-use crate::loose_shards::LooseShards;
-use ggus::{
-    GGufFileHeader, GGufMetaDataValueType, GGufMetaKVPairs, GGufReadError, GGufTensors, GGufWriter,
+use crate::{
+    gguf_file::{pad, GGufFile},
+    loose_shards::LooseShards,
 };
+use ggus::{GGufFileHeader, GGufMetaDataValueType, GGufWriter};
 use indexmap::{IndexMap, IndexSet};
 use std::{fs::File, iter::zip, path::PathBuf, thread};
 
@@ -44,32 +45,32 @@ impl MergeArgs {
 
         let kvs = files
             .iter()
-            .flat_map(|file| file.meta_kvs.kvs())
+            .flat_map(|file| file.meta_kvs().kvs())
             .filter(|kv| {
                 let key = kv.key();
                 !key.starts_with("split.") && key != "general.alignment"
             })
             .collect::<IndexSet<_>>();
         let tensors = files
             .iter()
-            .flat_map(|file| file.tensors.iter().map(move |t| (t, file.data)))
+            .flat_map(|file| file.tensors_as_indexmap())
             .collect::<IndexMap<_, _>>();
 
         let out = File::create(shards.single_file()).unwrap();
         let header = GGufFileHeader::new(3, tensors.len() as _, (kvs.len() + 1) as _);
-        let mut writer: GGufWriter<File> = GGufWriter::new(out, header).unwrap();
+        let mut writer = GGufWriter::new(out, header).unwrap();
 
         let align = files
             .iter()
-            .map(|file| file.meta_kvs.alignment())
+            .map(|file| file.meta_kvs().alignment())
             .max()
             .unwrap();
 
         writer
             .write_meta_kv(
                 "general.alignment",
-                GGufMetaDataValueType::U64,
-                (align as u64).to_le_bytes(),
+                GGufMetaDataValueType::U32,
+                (align as u32).to_le_bytes(),
             )
             .unwrap();
 
@@ -108,54 +109,3 @@ impl MergeArgs {
         }
     }
 }
-
-#[inline(always)]
-const fn pad(pos: usize, align: usize) -> usize {
-    (align - pos % align) % align
-}
-
-struct GGufFile<'a> {
-    meta_kvs: GGufMetaKVPairs<'a>,
-    tensors: GGufTensors<'a>,
-    data: &'a [u8],
-}
-
-#[derive(Debug)]
-enum GGufError<'a> {
-    MagicMismatch,
-    EndianNotSupport,
-    VersionNotSupport,
-    #[allow(dead_code)]
-    Reading(GGufReadError<'a>),
-}
-
-impl<'a> GGufFile<'a> {
-    fn new(data: &'a [u8]) -> Result<Self, GGufError<'a>> {
-        let header = unsafe { data.as_ptr().cast::<GGufFileHeader>().read() };
-        if !header.is_magic_correct() {
-            return Err(GGufError::MagicMismatch);
-        }
-        if !header.is_native_endian() {
-            return Err(GGufError::EndianNotSupport);
-        }
-        if header.version != 3 {
-            return Err(GGufError::VersionNotSupport);
-        }
-
-        let cursor = header.nbytes();
-        let meta_kvs = GGufMetaKVPairs::scan(header.metadata_kv_count, &data[cursor..])
-            .map_err(GGufError::Reading)?;
-
-        let cursor = cursor + meta_kvs.nbytes();
-        let tensors =
-            GGufTensors::scan(header.tensor_count, &data[cursor..]).map_err(GGufError::Reading)?;
-
-        let align = meta_kvs.alignment();
-        let cursor = (cursor + tensors.nbytes() + align - 1) / align * align;
-        Ok(Self {
-            meta_kvs,
-            tensors,
-            data: &data[cursor..],
-        })
-    }
-}