update style and ban mutiple split

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;
@@ -141,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>,

xtask/src/gguf_file.rs

@@ -51,15 +51,15 @@ impl<'a> GGufFile<'a> {
         })
     }
 
-    pub(crate) fn get_header(&self) -> &GGufFileHeader {
+    pub fn header(&self) -> &GGufFileHeader {
         &self.header
     }
 
-    pub(crate) fn get_meta_kvs(&self) -> &GGufMetaKVPairs<'a> {
+    pub fn meta_kvs(&self) -> &GGufMetaKVPairs<'a> {
         &self.meta_kvs
     }
 
-    pub(crate) fn get_tensors_as_indexmap(&self) -> IndexMap<ggus::GGufTensorInfo, &[u8]> {
+    pub fn tensors_as_indexmap(&self) -> IndexMap<ggus::GGufTensorInfo, &[u8]> {
         self.tensors
             .iter()
             .map(move |t| (t, self.data))

xtask/src/merge.rs

@@ -45,15 +45,15 @@ impl MergeArgs {
 
         let kvs = files
             .iter()
-            .flat_map(|file| file.get_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.get_tensors_as_indexmap())
+            .flat_map(|file| file.tensors_as_indexmap())
             .collect::<IndexMap<_, _>>();
 
         let out = File::create(shards.single_file()).unwrap();
@@ -62,7 +62,7 @@ impl MergeArgs {
 
         let align = files
             .iter()
-            .map(|file| file.get_meta_kvs().alignment())
+            .map(|file| file.meta_kvs().alignment())
             .max()
             .unwrap();
 

xtask/src/split.rs

@@ -1,4 +1,7 @@
-use crate::gguf_file::{pad, GGufError, GGufFile};
+use crate::{
+    gguf_file::{pad, GGufError, GGufFile},
+    loose_shards::LooseShards,
+};
 use ggus::{GGufFileHeader, GGufMetaDataValueType, GGufMetaKVPairs, GGufWriter};
 use std::{fs::File, path::PathBuf};
 
@@ -50,6 +53,11 @@ impl<'a> GGufFileInfo<'a> {
 
 impl SplitArgs {
     pub fn split(self) {
+        let shards = LooseShards::from(&*self.input);
+        if shards.count() > 1 {
+            println!("Model has already been splited");
+            return;
+        }
         let file = File::open(&self.input)
             .map_err(|e| {
                 println!("Failed to open");
@@ -60,10 +68,10 @@ impl SplitArgs {
         let mmap = unsafe { memmap2::Mmap::map(&file).unwrap() };
         let ctx_gguf: GGufFile = GGufFile::new(&mmap).unwrap();
 
-        let align = ctx_gguf.get_meta_kvs().alignment();
+        let align = ctx_gguf.meta_kvs().alignment();
         let ggufs = self.split_strategy(ctx_gguf.clone()).unwrap();
 
-        let tensors = ctx_gguf.get_tensors_as_indexmap();
+        let tensors = ctx_gguf.tensors_as_indexmap();
         let mut tensor_iter: indexmap::map::Iter<ggus::GGufTensorInfo, &[u8]> = tensors.iter();
 
         for gguf in ggufs {
@@ -75,12 +83,6 @@ impl SplitArgs {
 
             let kvs = gguf.meta_kvs.kvs();
             for kv in kvs {
-                if kv.key() == LLM_KV_SPLIT_TENSORS_COUNT
-                    || kv.key() == LLM_KV_SPLIT_COUNT
-                    || kv.key() == LLM_KV_SPLIT_NO
-                {
-                    continue;
-                }
                 writer
                     .write_meta_kv(kv.key(), kv.ty(), kv.value_bytes())
                     .unwrap();
@@ -197,10 +199,10 @@ impl SplitArgs {
             },
         }
 
-        let tensors = ctx_gguf.get_tensors_as_indexmap();
+        let tensors = ctx_gguf.tensors_as_indexmap();
 
         let mut ggufs: Vec<GGufFileInfo> = Vec::new();
-        let n_tensors: u64 = ctx_gguf.get_header().tensor_count;
+        let n_tensors: u64 = ctx_gguf.header().tensor_count;
 
         let setup_gguf_file = |i_split: u64, n_tensors: u64| {
             let mut gguf_file = GGufFileInfo::new_empty();
@@ -221,20 +223,20 @@ impl SplitArgs {
 
         let mut i_split: u64 = 1;
         let mut gguf_file = setup_gguf_file(i_split, n_tensors);
-        gguf_file.meta_kvs = ctx_gguf.get_meta_kvs().clone();
+        gguf_file.meta_kvs = ctx_gguf.meta_kvs().clone();
+
+        gguf_file.header.metadata_kv_count += ctx_gguf.header().metadata_kv_count;
 
-        if gguf_file.meta_kvs.get(LLM_KV_SPLIT_NO).is_some() {
+        if gguf_file.meta_kvs.remove(LLM_KV_SPLIT_NO) {
             gguf_file.header.metadata_kv_count -= 1;
         }
-        if gguf_file.meta_kvs.get(LLM_KV_SPLIT_COUNT).is_some() {
+        if gguf_file.meta_kvs.remove(LLM_KV_SPLIT_COUNT) {
             gguf_file.header.metadata_kv_count -= 1;
         }
-        if gguf_file.meta_kvs.get(LLM_KV_SPLIT_TENSORS_COUNT).is_some() {
+        if gguf_file.meta_kvs.remove(LLM_KV_SPLIT_TENSORS_COUNT) {
             gguf_file.header.metadata_kv_count -= 1;
         }
 
-        gguf_file.header.metadata_kv_count += ctx_gguf.get_header().metadata_kv_count;
-
         if self.no_tensor_first_split {
             ggufs.push(gguf_file);
             gguf_file = setup_gguf_file(i_split, n_tensors);
@@ -267,7 +269,12 @@ impl SplitArgs {
         let output_path = &self.output.unwrap();
         let mut index = 0;
         while index < ggufs.len() {
-            ggufs[index].output_path = split_path(output_path, index + 1, ggufs.len());
+            ggufs[index].output_path = format!(
+                "{}-{:05}-of-{:05}.gguf",
+                output_path,
+                index + 1,
+                ggufs.len()
+            );
             ggufs[index].new_kv_tuples[1] = (LLM_KV_SPLIT_COUNT.to_string(), ty::U16, tensor_count);
             index += 1;
         }
@@ -283,28 +290,3 @@ impl SplitArgs {
         }
     }
 }
-
-fn split_path(path_prefix: &String, split_no: usize, split_count: usize) -> String {
-    format!("{}-{:05}-of-{:05}.gguf", path_prefix, split_no, split_count)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    #[test]
-    fn it_works() {
-        let input = PathBuf::from("/home/qinyiqun/gguf/xtask/src/test/rust/rust_ori.gguf");
-        let output = Some("/home/qinyiqun/gguf/xtask/src/test/rust/rust_ori_oi".to_string());
-        let split_args = SplitArgs {
-            input,
-            output,
-            split_max_tensors: None,
-            split_max_size: Some("300M".to_string()),
-            no_tensor_first_split: true,
-            n_bytes_split: 0,
-            n_split_tensors: 0,
-        };
-
-        split_args.split();
-    }
-}