temp

Signed-off-by: YdrMaster <ydrml@hotmail.com>

ggus/src/lib.rs

@@ -11,12 +11,12 @@ mod writer;
 pub use header::GGufFileHeader;
 pub use metadata::{
     utok, GGufArray, GGufFileType, GGufMetaDataValueType, GGufMetaKV, GGufMetaKVPairs,
-    GGufTokenType, GENERAL_ALIGNMENT,
+    GGufTokenType, DEFAULT_ALIGNMENT, GENERAL_ALIGNMENT,
 };
 pub use name::GGufFileName;
 pub use reader::{GGufReadError, GGufReader};
 pub use tensor::{GGmlType, GGufTensorInfo, GGufTensors};
-pub use writer::GGufWriter;
+pub use writer::{GGufMetaWriter, GGufSimulator, GGufTensorWriter};
 
 macro_rules! sizeof {
     ($ty:ty) => {

ggus/src/metadata/general.rs

@@ -1,6 +1,7 @@
 use super::{GGufFileType, GGufMetaDataValueType as Ty, GGufMetaKVPairs};
 
 pub const GENERAL_ALIGNMENT: &str = "general.alignment";
+pub const DEFAULT_ALIGNMENT: usize = 32;
 
 impl<'a> GGufMetaKVPairs<'a> {
     pub fn architecture(&self) -> Option<&'a str> {
@@ -17,7 +18,9 @@ impl<'a> GGufMetaKVPairs<'a> {
     #[inline]
     pub fn alignment(&self) -> usize {
         self.get_typed(GENERAL_ALIGNMENT, Ty::U32)
-            .map_or(32, |mut reader| reader.read::<u32>().unwrap() as _)
+            .map_or(DEFAULT_ALIGNMENT, |mut reader| {
+                reader.read::<u32>().unwrap() as _
+            })
     }
 
     #[inline]

ggus/src/metadata/mod.rs

@@ -11,7 +11,7 @@ use crate::{
 use indexmap::IndexMap;
 use std::{hash::Hash, slice::from_raw_parts};
 
-pub use general::GENERAL_ALIGNMENT;
+pub use general::{DEFAULT_ALIGNMENT, GENERAL_ALIGNMENT};
 pub use tokenizer::{utok, GGufArray, GGufTokenType};
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]

ggus/src/writer.rs

@@ -1,90 +1,211 @@
-use crate::{sizeof, tensor::GGmlType, GGufFileHeader, GGufMetaDataValueType, GENERAL_ALIGNMENT};
+use crate::{
+    GGufFileHeader, GGufMetaDataValueType, GGufTensorInfo, DEFAULT_ALIGNMENT, GENERAL_ALIGNMENT,
+};
 use internal::Internal;
 use std::{
     io::{Result, Write},
     mem::size_of_val,
     slice::from_raw_parts,
 };
 
-pub struct GGufWriter<T: Write>(Internal<T>);
+pub struct GGufMetaWriter<T: Write> {
+    writer: GGufWriter<T>,
+    alignment: usize,
+}
 
-impl<T: Write> GGufWriter<T> {
+pub struct GGufTensorWriter<'t, T: Write> {
+    writer: GGufWriter<T>,
+    alignment: usize,
+    data: Vec<&'t [u8]>,
+    offset: usize,
+}
+
+pub struct GGufSimulator(GGufWriter<NWrite>);
+
+impl<T: Write> GGufMetaWriter<T> {
     #[inline]
     pub fn new(writer: T, header: GGufFileHeader) -> Result<Self> {
-        let mut ans = Self(Internal::new(writer));
-        ans.write_bytes(as_slice(&header))?;
-        Ok(ans)
+        let mut internal = Internal::new(writer);
+        internal.write_bytes(as_slice(&header))?;
+        Ok(Self {
+            writer: GGufWriter(internal),
+            alignment: DEFAULT_ALIGNMENT,
+        })
     }
 
     #[inline]
-    pub const fn written_bytes(&self) -> usize {
-        self.0.written_bytes()
+    pub fn write_alignment(&mut self, align: usize) -> Result<()> {
+        self.write_meta_kv(
+            GENERAL_ALIGNMENT,
+            GGufMetaDataValueType::U32,
+            (align as u32).to_le_bytes(),
+        )
     }
 
-    #[inline]
-    pub fn write_bytes(&mut self, val: &[u8]) -> Result<()> {
-        self.0.write_bytes(val)
+    pub fn write_meta_kv(
+        &mut self,
+        key: impl AsRef<str>,
+        ty: GGufMetaDataValueType,
+        val: impl AsRef<[u8]>,
+    ) -> Result<()> {
+        if key.as_ref() == GENERAL_ALIGNMENT {
+            let &[a, b, c, d] = val.as_ref() else {
+                use std::io::{Error, ErrorKind::InvalidData};
+                return Err(Error::new(InvalidData, "general.alignment must be an u32"));
+            };
+            self.alignment = u32::from_le_bytes([a, b, c, d]) as _;
+        }
+
+        self.writer.write_str(key)?;
+        self.writer.write(ty)?;
+        self.writer.write_bytes(val.as_ref())
     }
 
     #[inline]
-    pub fn write<U: Copy>(&mut self, val: U) -> Result<()> {
-        self.write_bytes(as_slice(&val))
+    pub fn finish<'t>(self) -> GGufTensorWriter<'t, T> {
+        GGufTensorWriter {
+            writer: self.writer,
+            alignment: self.alignment,
+            data: Vec::new(),
+            offset: 0,
+        }
+    }
+}
+
+impl<'t, T: Write> GGufTensorWriter<'t, T> {
+    pub fn write_tensor(&mut self, info: &GGufTensorInfo, data: &'t [u8]) -> Result<()> {
+        self.writer.write_str(info.name())?;
+
+        let shape = info.shape();
+        self.writer.write(shape.len() as u32)?;
+        self.writer.write_bytes(as_slice(shape))?;
+
+        self.writer.write(info.ggml_type())?;
+        self.writer.write(self.offset as u64)?;
+
+        self.offset += data.len();
+        self.offset = self.offset.div_ceil(self.alignment) * self.alignment;
+        self.data.push(&data[info.offset()..][..info.nbytes()]);
+
+        Ok(())
+    }
+
+    fn pad(&mut self) -> Result<()> {
+        let padding =
+            (self.alignment - self.writer.written_bytes() % self.alignment) % self.alignment;
+        for _ in 0..padding {
+            self.writer.write(0u8)?;
+        }
+        Ok(())
+    }
+
+    pub fn finish(mut self) -> Result<usize> {
+        for data in std::mem::take(&mut self.data) {
+            self.pad()?;
+            self.writer.write_bytes(data)?;
+        }
+        Ok(self.writer.written_bytes())
     }
+}
 
+impl GGufSimulator {
     #[inline]
-    pub fn write_bool(&mut self, val: bool) -> Result<()> {
-        self.write_bytes(if val { &[1] } else { &[0] })
+    pub fn new(with_header: bool) -> Self {
+        let mut internal = Internal::new(NWrite);
+        if with_header {
+            internal
+                .write_bytes(as_slice(&GGufFileHeader::default()))
+                .unwrap();
+        }
+        Self(GGufWriter(internal))
     }
 
     #[inline]
-    pub fn write_str(&mut self, val: impl AsRef<str>) -> Result<()> {
-        let val = val.as_ref();
-        self.write_bytes(as_slice(&(val.len() as u64)))?;
-        self.write_bytes(val.as_bytes())
+    pub fn write_alignment(&mut self) -> usize {
+        self.write_meta_kv(
+            GENERAL_ALIGNMENT,
+            GGufMetaDataValueType::U32,
+            DEFAULT_ALIGNMENT.to_le_bytes(),
+        )
     }
 
     pub fn write_meta_kv(
         &mut self,
         key: impl AsRef<str>,
         ty: GGufMetaDataValueType,
         val: impl AsRef<[u8]>,
-    ) -> Result<()> {
-        self.write_str(key)?;
-        self.write(ty)?;
-        self.write_bytes(val.as_ref())
+    ) -> usize {
+        self.0.write_str(key).unwrap();
+        self.0.write(ty).unwrap();
+        self.0.write_bytes(val.as_ref()).unwrap();
+        self.0.written_bytes()
     }
 
-    pub fn write_alignment(&mut self, align: usize) -> Result<()> {
-        self.write_meta_kv(
-            GENERAL_ALIGNMENT,
-            GGufMetaDataValueType::U32,
-            (align as u32).to_le_bytes(),
-        )
+    pub fn write_tensor(&mut self, info: &GGufTensorInfo) -> usize {
+        self.0.write_str(info.name()).unwrap();
+
+        let shape = info.shape();
+        self.0.write(shape.len() as u32).unwrap();
+        self.0.write_bytes(as_slice(shape)).unwrap();
+
+        self.0.write(info.ggml_type()).unwrap();
+        self.0.write(0u64).unwrap();
+
+        self.0.write(info.nbytes()).unwrap();
+
+        self.0.written_bytes()
     }
+}
 
-    pub fn write_tensor_info(
-        &mut self,
-        name: impl AsRef<str>,
-        shape: &[u64],
-        ty: GGmlType,
-        offset: usize,
-    ) -> Result<()> {
-        self.write_str(name)?;
-        self.write(shape.len() as u32)?;
-        self.write_bytes(unsafe { from_raw_parts(shape.as_ptr().cast(), size_of_val(shape)) })?;
-        self.write(ty)?;
-        self.write(offset as u64)
+struct NWrite;
+
+impl Write for NWrite {
+    #[inline(always)]
+    fn write(&mut self, buf: &[u8]) -> Result<usize> {
+        Ok(buf.len())
+    }
+    #[inline(always)]
+    fn flush(&mut self) -> Result<()> {
+        Ok(())
+    }
+}
+
+struct GGufWriter<T: Write>(Internal<T>);
+
+impl<T: Write> GGufWriter<T> {
+    // #[inline]
+    // pub fn new(writer: T, header: GGufFileHeader) -> Result<Self> {
+    //     let mut ans = Self(Internal::new(writer));
+    //     ans.write_bytes(as_slice(&header))?;
+    //     Ok(ans)
+    // }
+
+    #[inline]
+    pub const fn written_bytes(&self) -> usize {
+        self.0.written_bytes()
     }
 
     #[inline]
-    pub fn flush(&mut self) -> Result<()> {
-        self.0.flush()
+    pub fn write_bytes(&mut self, val: &[u8]) -> Result<()> {
+        self.0.write_bytes(val)
+    }
+
+    #[inline]
+    pub fn write<U: Copy>(&mut self, val: U) -> Result<()> {
+        self.write_bytes(as_slice(&val))
+    }
+
+    #[inline]
+    pub fn write_str(&mut self, val: impl AsRef<str>) -> Result<()> {
+        let val = val.as_ref();
+        self.write_bytes(as_slice(&(val.len() as u64)))?;
+        self.write_bytes(val.as_bytes())
     }
 }
 
 #[inline(always)]
-fn as_slice<T: Sized>(val: &T) -> &[u8] {
-    unsafe { from_raw_parts(val as *const _ as *const _, sizeof!(T)) }
+fn as_slice<T: ?Sized>(val: &T) -> &[u8] {
+    unsafe { from_raw_parts(val as *const _ as *const _, size_of_val(val)) }
 }
 
 mod internal {
@@ -108,10 +229,5 @@ mod internal {
             self.1 += val.len();
             self.0.write_all(val.as_ref())
         }
-
-        #[inline]
-        pub fn flush(&mut self) -> Result<()> {
-            self.0.flush()
-        }
     }
 }

xtask/src/filter.rs

@@ -1,8 +1,5 @@
-use crate::{
-    file_info::print_file_info, gguf_file::GGufFile, name_pattern::compile_patterns, write_tensors,
-    YES,
-};
-use ggus::{GGufFileHeader, GGufWriter};
+use crate::{file_info::print_file_info, gguf_file::GGufFile, name_pattern::compile_patterns, YES};
+use ggus::{GGufFileHeader, GGufMetaWriter};
 use std::{fs::File, path::PathBuf};
 
 #[derive(Args, Default)]
@@ -68,15 +65,19 @@ impl FilterArgs {
         out = out.with_extension("part.gguf");
 
         let header = GGufFileHeader::new(3, tensors.len() as _, meta_kvs.len() as _);
-        let mut writer = GGufWriter::new(File::create(&out).unwrap(), header).unwrap();
+        let mut writer = GGufMetaWriter::new(File::create(&out).unwrap(), header).unwrap();
         for kv in &meta_kvs {
             writer
                 .write_meta_kv(kv.key(), kv.ty(), kv.value_bytes())
                 .unwrap();
         }
-        write_tensors(&mut writer, &tensors, file.meta_kvs.alignment(), file.data);
+        let mut writer = writer.finish();
+        for t in &tensors {
+            writer.write_tensor(t, file.data).unwrap();
+        }
+        let n_bytes = writer.finish().unwrap();
 
         println!("{YES}Shard is written to: \"{}\"", out.display());
-        print_file_info(tensors.len(), meta_kvs.len(), writer.written_bytes());
+        print_file_info(tensors.len(), meta_kvs.len(), n_bytes);
     }
 }