From cc440e32f0a7e0e5f0b142b650637ad02ec85118 Mon Sep 17 00:00:00 2001 From: YdrMaster Date: Mon, 12 Aug 2024 18:20:51 +0800 Subject: [PATCH] =?UTF-8?q?feat(ggus):=20GGuf=20=E6=96=87=E4=BB=B6?= =?UTF-8?q?=E6=A0=BC=E5=BC=8F=E7=A7=BB=E5=8A=A8=E5=88=B0=20ggus=20?= =?UTF-8?q?=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: YdrMaster --- Cargo.lock | 16 ++ xtask/src/gguf_file.rs => ggus/src/file.rs | 16 +- ggus/src/header.rs | 5 +- ggus/src/lib.rs | 19 +- ggus/src/metadata/general.rs | 2 +- ggus/src/metadata/llm.rs | 2 +- ggus/src/metadata/mod.rs | 7 +- ggus/src/metadata/tokenizer.rs | 5 +- ggus/src/name.rs | 3 +- ggus/src/{reader.rs => read.rs} | 8 +- ggus/src/tensor.rs | 27 +- ggus/src/write/internal.rs | 118 +++++++++ ggus/src/write/mod.rs | 6 + ggus/src/write/simulator.rs | 91 +++++++ ggus/src/write/writer.rs | 86 ++++++ ggus/src/writer.rs | 287 --------------------- xtask/Cargo.toml | 5 +- xtask/src/convert/mod.rs | 4 +- xtask/src/convert/read.rs | 6 +- xtask/src/file_info.rs | 67 +++-- xtask/src/filter.rs | 8 +- xtask/src/main.rs | 1 - xtask/src/merge.rs | 10 +- xtask/src/split.rs | 10 +- 24 files changed, 421 insertions(+), 388 deletions(-) rename xtask/src/gguf_file.rs => ggus/src/file.rs (83%) rename ggus/src/{reader.rs => read.rs} (85%) create mode 100644 ggus/src/write/internal.rs create mode 100644 ggus/src/write/mod.rs create mode 100644 ggus/src/write/simulator.rs create mode 100644 ggus/src/write/writer.rs delete mode 100644 ggus/src/writer.rs diff --git a/Cargo.lock b/Cargo.lock index 23a2758..100be70 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -127,6 +127,12 @@ dependencies = [ "half", ] +[[package]] +name = "either" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" + [[package]] name = "equivalent" version = "1.0.1" @@ -178,6 +184,15 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "libc" version = "0.2.155" @@ -355,6 +370,7 @@ dependencies = [ "clap", "digit-layout", "ggus", + "itertools", "memmap2", "regex", ] diff --git a/xtask/src/gguf_file.rs b/ggus/src/file.rs similarity index 83% rename from xtask/src/gguf_file.rs rename to ggus/src/file.rs index 1ca2e22..7a7928f 100644 --- a/xtask/src/gguf_file.rs +++ b/ggus/src/file.rs @@ -1,9 +1,8 @@ -use ggus::{GGufFileHeader, GGufMetaKVPairs, GGufReadError, GGufTensors}; +use crate::{pad, GGufFileHeader, GGufMetaKVPairs, GGufReadError, GGufTensors}; use std::{error::Error, fmt}; #[derive(Clone)] -pub(crate) struct GGufFile<'a> { - #[allow(unused)] +pub struct GGuf<'a> { pub header: GGufFileHeader, pub meta_kvs: GGufMetaKVPairs<'a>, pub tensors: GGufTensors<'a>, @@ -11,7 +10,7 @@ pub(crate) struct GGufFile<'a> { } #[derive(Debug)] -pub(crate) enum GGufError { +pub enum GGufError { MagicMismatch, EndianNotSupport, VersionNotSupport, @@ -31,8 +30,8 @@ impl fmt::Display for GGufError { impl Error for GGufError {} -impl<'a> GGufFile<'a> { - pub fn new(data: &'a [u8]) -> Result { +impl<'a> GGuf<'a> { + pub fn scan(data: &'a [u8]) -> Result { let header = unsafe { data.as_ptr().cast::().read() }; if !header.is_magic_correct() { return Err(GGufError::MagicMismatch); @@ -66,8 +65,3 @@ impl<'a> GGufFile<'a> { }) } } - -#[inline(always)] -pub(crate) const fn pad(pos: usize, align: usize) -> usize { - (align - pos % align) % align -} diff --git a/ggus/src/header.rs b/ggus/src/header.rs index 8e2a5e3..26ef136 100644 --- a/ggus/src/header.rs +++ b/ggus/src/header.rs @@ -1,5 +1,4 @@ -use crate::sizeof; -use std::str::Utf8Error; +use std::str::Utf8Error; #[derive(Clone, Default, Debug)] #[repr(C)] @@ -45,6 +44,6 @@ impl GGufFileHeader { #[inline] pub const fn nbytes(&self) -> usize { - sizeof!(Self) + size_of::() } } diff --git a/ggus/src/lib.rs b/ggus/src/lib.rs index e959c8d..85cb877 100644 --- a/ggus/src/lib.rs +++ b/ggus/src/lib.rs @@ -1,27 +1,26 @@ #![doc = include_str!("../README.md")] #![deny(warnings)] +mod file; mod header; mod metadata; mod name; -mod reader; +mod read; mod tensor; -mod writer; +mod write; +pub use file::{GGuf, GGufError}; pub use header::GGufFileHeader; pub use metadata::{ utok, GGufArray, GGufFileType, GGufMetaDataValueType, GGufMetaKV, GGufMetaKVPairs, GGufTokenType, DEFAULT_ALIGNMENT, GENERAL_ALIGNMENT, }; pub use name::GGufFileName; -pub use reader::{GGufReadError, GGufReader}; +pub use read::{GGufReadError, GGufReader}; pub use tensor::{GGmlType, GGufTensorInfo, GGufTensors}; -pub use writer::{GGufMetaWriter, GGufSimulator, GGufTensorWriter}; +pub use write::{GGufMetaWriter, GGufSimulator, GGufTensorWriter}; -macro_rules! sizeof { - ($ty:ty) => { - std::mem::size_of::<$ty>() - }; +#[inline(always)] +const fn pad(pos: usize, align: usize) -> usize { + (align - pos % align) % align } - -use sizeof; diff --git a/ggus/src/metadata/general.rs b/ggus/src/metadata/general.rs index e3d8fb6..4a820b1 100644 --- a/ggus/src/metadata/general.rs +++ b/ggus/src/metadata/general.rs @@ -1,4 +1,4 @@ -use super::{GGufFileType, GGufMetaDataValueType as Ty, GGufMetaKVPairs}; +use crate::{GGufFileType, GGufMetaDataValueType as Ty, GGufMetaKVPairs}; pub const GENERAL_ALIGNMENT: &str = "general.alignment"; pub const DEFAULT_ALIGNMENT: usize = 32; diff --git a/ggus/src/metadata/llm.rs b/ggus/src/metadata/llm.rs index 1a71073..3417a11 100644 --- a/ggus/src/metadata/llm.rs +++ b/ggus/src/metadata/llm.rs @@ -1,4 +1,4 @@ -use super::{GGufMetaDataValueType, GGufMetaKVPairs}; +use crate::{GGufMetaDataValueType, GGufMetaKVPairs}; pub struct LlmMeta<'a>(GGufMetaKVPairs<'a>, &'a str); diff --git a/ggus/src/metadata/mod.rs b/ggus/src/metadata/mod.rs index 8c047b1..652d158 100644 --- a/ggus/src/metadata/mod.rs +++ b/ggus/src/metadata/mod.rs @@ -4,10 +4,7 @@ mod general; mod llm; mod tokenizer; -use crate::{ - reader::{GGufReadError, GGufReader}, - sizeof, -}; +use crate::{GGufReadError, GGufReader}; use indexmap::IndexMap; use std::{hash::Hash, slice::from_raw_parts}; @@ -234,7 +231,7 @@ impl<'a> GGufMetaKV<'a> { self.key .as_ptr() .add(self.key.len()) - .add(sizeof!(GGufMetaDataValueType)), + .add(size_of::()), self.len, ) } diff --git a/ggus/src/metadata/tokenizer.rs b/ggus/src/metadata/tokenizer.rs index 084fcc6..cb77d29 100644 --- a/ggus/src/metadata/tokenizer.rs +++ b/ggus/src/metadata/tokenizer.rs @@ -1,5 +1,4 @@ -use super::{GGufMetaDataValueType, GGufMetaKVPairs}; -use crate::reader::GGufReader; +use crate::{GGufMetaDataValueType, GGufMetaKVPairs, GGufReader}; use std::marker::PhantomData; #[repr(transparent)] @@ -78,7 +77,7 @@ impl<'a> GGufArray<'a, str> { } } -impl<'a, T: Copy> Iterator for GGufArray<'a, T> { +impl<'a, T: Copy + 'static> Iterator for GGufArray<'a, T> { type Item = T; #[inline] fn next(&mut self) -> Option { diff --git a/ggus/src/name.rs b/ggus/src/name.rs index 5c62f0e..8b1edb5 100644 --- a/ggus/src/name.rs +++ b/ggus/src/name.rs @@ -1,5 +1,4 @@ -use core::fmt; -use std::str::FromStr; +use std::{fmt, str::FromStr}; #[derive(Clone, Debug)] pub struct GGufFileName { diff --git a/ggus/src/reader.rs b/ggus/src/read.rs similarity index 85% rename from ggus/src/reader.rs rename to ggus/src/read.rs index 2ce8db1..9080179 100644 --- a/ggus/src/reader.rs +++ b/ggus/src/read.rs @@ -1,4 +1,4 @@ -use crate::{sizeof, GGufMetaDataValueType}; +use crate::GGufMetaDataValueType; use std::str::Utf8Error; pub struct GGufReader<'a> { @@ -25,8 +25,8 @@ impl<'a> GGufReader<'a> { self.cursor } - pub(crate) fn skip(&mut self, len: usize) -> Result<(), GGufReadError> { - let len = len * sizeof!(T); + pub(crate) fn skip(&mut self, len: usize) -> Result<(), GGufReadError> { + let len = len * size_of::(); let data = &self.data[self.cursor..]; if data.len() >= len { self.cursor += len; @@ -36,7 +36,7 @@ impl<'a> GGufReader<'a> { } } - pub fn read(&mut self) -> Result { + pub fn read(&mut self) -> Result { let ptr = self.data[self.cursor..].as_ptr().cast::(); self.skip::(1)?; Ok(unsafe { ptr.read_unaligned() }) diff --git a/ggus/src/tensor.rs b/ggus/src/tensor.rs index b4db449..0acc3aa 100644 --- a/ggus/src/tensor.rs +++ b/ggus/src/tensor.rs @@ -1,9 +1,8 @@ -use crate::{ - reader::{GGufReadError, GGufReader}, - sizeof, -}; +use crate::{GGufReadError, GGufReader}; use indexmap::IndexMap; -use std::{hash::Hash, marker::PhantomData, slice::from_raw_parts, str::from_utf8_unchecked}; +use std::{ + hash::Hash, marker::PhantomData, mem::size_of, slice::from_raw_parts, str::from_utf8_unchecked, +}; #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] #[repr(u32)] @@ -45,7 +44,7 @@ pub enum GGmlType { impl GGmlType { fn nbytes(self) -> usize { match self { - Self::F32 => sizeof!(f32), + Self::F32 => size_of::(), Self::F16 => 2, Self::Q4_0 => todo!(), Self::Q4_1 => todo!(), @@ -67,11 +66,11 @@ impl GGmlType { Self::IQ3S => todo!(), Self::IQ2S => todo!(), Self::IQ4XS => todo!(), - Self::I8 => sizeof!(i8), - Self::I16 => sizeof!(i16), - Self::I32 => sizeof!(i32), - Self::I64 => sizeof!(i64), - Self::F64 => sizeof!(f64), + Self::I8 => size_of::(), + Self::I16 => size_of::(), + Self::I32 => size_of::(), + Self::I64 => size_of::(), + Self::F64 => size_of::(), Self::IQ1M => todo!(), _ => unimplemented!(), } @@ -162,13 +161,13 @@ impl<'a> GGufTensorInfo<'a> { let ptr = name.as_ptr().add(name.len()); let ndim = ptr.cast::().read_unaligned() as usize; - let ptr = ptr.add(sizeof!(u32)); + let ptr = ptr.add(size_of::()); let shape = ptr; - let ptr = ptr.add(ndim * sizeof!(u64)); + let ptr = ptr.add(ndim * size_of::()); let ggml_type = ptr.cast::().read_unaligned(); - let ptr = ptr.add(sizeof!(u32)); + let ptr = ptr.add(size_of::()); let offset = ptr.cast::().read_unaligned(); let mut body = vec![0u64; 4 + ndim].into_boxed_slice(); diff --git a/ggus/src/write/internal.rs b/ggus/src/write/internal.rs new file mode 100644 index 0000000..9a488b5 --- /dev/null +++ b/ggus/src/write/internal.rs @@ -0,0 +1,118 @@ +use crate::{pad, GGmlType, GGufFileHeader, GGufMetaDataValueType, GENERAL_ALIGNMENT}; +use internal::Internal; +use std::{ + io::{Result, Write}, + slice::from_raw_parts, +}; + +pub(super) struct GGufWriter(Internal); + +impl GGufWriter { + #[inline] + fn write(&mut self, val: &[U]) -> Result<()> { + self.0 + .write_bytes(unsafe { from_raw_parts(val.as_ptr().cast(), size_of_val(val)) }) + } + + #[inline] + fn write_str(&mut self, val: impl AsRef) -> Result<()> { + let val = val.as_ref().as_bytes(); + self.write(&[val.len() as u64])?; + self.write(val) + } + + #[inline] + pub fn new(writer: T, header: GGufFileHeader) -> Result { + let mut ans = Self(Internal::new(writer)); + ans.write(unsafe { + from_raw_parts( + &header as *const _ as *const u8, + size_of::(), + ) + })?; + Ok(ans) + } + + #[inline] + pub const fn written_bytes(&self) -> usize { + self.0.written_bytes() + } + + #[inline] + pub fn write_alignment(&mut self, align: usize) -> Result { + self.write_meta_kv( + GENERAL_ALIGNMENT, + GGufMetaDataValueType::U32, + (align as u32).to_le_bytes(), + ) + .map(Option::unwrap) + } + + pub fn write_meta_kv( + &mut self, + key: impl AsRef, + ty: GGufMetaDataValueType, + val: impl AsRef<[u8]>, + ) -> Result> { + let key = key.as_ref(); + let val = val.as_ref(); + + self.write_str(key)?; + self.write(&[ty])?; + self.write(val)?; + + Ok(if key == GENERAL_ALIGNMENT { + let &[a, b, c, d] = val else { + panic!("general.alignment must be an u32") + }; + Some(u32::from_le_bytes([a, b, c, d]) as _) + } else { + None + }) + } + + pub fn write_tensor_info( + &mut self, + name: &str, + shape: &[u64], + data_type: GGmlType, + offset: u64, + ) -> Result<()> { + self.write_str(name)?; + self.write(&[shape.len() as u32])?; + self.write(shape)?; + self.write(&[data_type])?; + self.write(&[offset]) + } + + pub fn write_data(&mut self, data: &[u8], align: usize) -> Result<()> { + for _ in 0..pad(self.written_bytes(), align) { + self.write(&[0u8])?; + } + self.write(data) + } +} + +mod internal { + use std::io::{BufWriter, Result, Write}; + + pub(super) struct Internal(BufWriter, usize); + + impl Internal { + #[inline] + pub fn new(writer: T) -> Self { + Self(BufWriter::new(writer), 0) + } + + #[inline] + pub const fn written_bytes(&self) -> usize { + self.1 + } + + #[inline] + pub fn write_bytes(&mut self, val: &[u8]) -> Result<()> { + self.1 += val.len(); + self.0.write_all(val.as_ref()) + } + } +} diff --git a/ggus/src/write/mod.rs b/ggus/src/write/mod.rs new file mode 100644 index 0000000..17b153f --- /dev/null +++ b/ggus/src/write/mod.rs @@ -0,0 +1,6 @@ +mod internal; +mod simulator; +mod writer; + +pub use simulator::GGufSimulator; +pub use writer::{GGufMetaWriter, GGufTensorWriter}; diff --git a/ggus/src/write/simulator.rs b/ggus/src/write/simulator.rs new file mode 100644 index 0000000..866e510 --- /dev/null +++ b/ggus/src/write/simulator.rs @@ -0,0 +1,91 @@ +use super::internal::GGufWriter; +use crate::{pad, GGufFileHeader, GGufMetaDataValueType, GGufTensorInfo, DEFAULT_ALIGNMENT}; +use std::io::{Result, Write}; + +pub struct GGufSimulator { + writer: GGufWriter, + alignment: usize, + data: Vec, + offset: usize, +} + +impl Default for GGufSimulator { + #[inline] + fn default() -> Self { + Self::new() + } +} + +impl GGufSimulator { + #[inline] + pub fn new() -> Self { + Self { + writer: GGufWriter::new(NWrite, GGufFileHeader::default()).unwrap(), + alignment: DEFAULT_ALIGNMENT, + data: Vec::new(), + offset: 0, + } + } + + #[inline] + pub fn with_alignment(align: usize) -> Self { + let mut ans = Self::new(); + ans.write_alignment(align); + ans + } + + #[inline] + pub fn write_alignment(&mut self, align: usize) { + self.alignment = self.writer.write_alignment(align).unwrap(); + } + + #[inline] + pub fn write_meta_kv( + &mut self, + key: impl AsRef, + ty: GGufMetaDataValueType, + val: impl AsRef<[u8]>, + ) { + if let Some(align) = self.writer.write_meta_kv(key, ty, val).unwrap() { + self.alignment = align; + } + } + + pub fn write_tensor(&mut self, info: &GGufTensorInfo) { + self.offset += pad(self.offset, self.alignment); + self.writer + .write_tensor_info( + info.name(), + info.shape(), + info.ggml_type(), + self.offset as _, + ) + .unwrap(); + + let len = info.nbytes(); + self.offset += len; + self.data.push(len); + } + + pub fn written_bytes(&self) -> usize { + let mut total = self.writer.written_bytes(); + for len in &self.data { + total += pad(total, self.alignment); + total += len; + } + total + } +} + +struct NWrite; + +impl Write for NWrite { + #[inline(always)] + fn write(&mut self, buf: &[u8]) -> Result { + Ok(buf.len()) + } + #[inline(always)] + fn flush(&mut self) -> Result<()> { + Ok(()) + } +} diff --git a/ggus/src/write/writer.rs b/ggus/src/write/writer.rs new file mode 100644 index 0000000..ef66c6f --- /dev/null +++ b/ggus/src/write/writer.rs @@ -0,0 +1,86 @@ +use super::internal::GGufWriter; +use crate::{pad, GGufFileHeader, GGufMetaDataValueType, GGufTensorInfo, DEFAULT_ALIGNMENT}; +use std::io::{Result, Write}; + +pub struct GGufMetaWriter { + writer: GGufWriter, + alignment: usize, +} + +pub struct GGufTensorWriter<'t, T: Write> { + writer: GGufWriter, + alignment: usize, + data: Vec<&'t [u8]>, + offset: usize, +} + +impl GGufMetaWriter { + #[inline] + pub fn new(writer: T, header: GGufFileHeader) -> Result { + Ok(Self { + writer: GGufWriter::new(writer, header)?, + alignment: DEFAULT_ALIGNMENT, + }) + } + + #[inline] + pub fn write_alignment(&mut self, align: usize) -> Result<()> { + self.alignment = self.writer.write_alignment(align)?; + Ok(()) + } + + #[inline] + pub fn write_meta_kv( + &mut self, + key: impl AsRef, + ty: GGufMetaDataValueType, + val: impl AsRef<[u8]>, + ) -> Result<()> { + if let Some(align) = self.writer.write_meta_kv(key, ty, val)? { + self.alignment = align; + } + Ok(()) + } + + #[inline] + 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.offset += pad(self.offset, self.alignment); + self.writer.write_tensor_info( + info.name(), + info.shape(), + info.ggml_type(), + self.offset as _, + )?; + + let data = &data[info.offset()..][..info.nbytes()]; + self.offset += data.len(); + self.data.push(data); + + Ok(()) + } + + pub fn finish(self) -> Result { + let Self { + mut writer, + alignment, + data, + .. + } = self; + + for data in data { + writer.write_data(data, alignment)?; + } + Ok(writer.written_bytes()) + } +} diff --git a/ggus/src/writer.rs b/ggus/src/writer.rs deleted file mode 100644 index 3098784..0000000 --- a/ggus/src/writer.rs +++ /dev/null @@ -1,287 +0,0 @@ -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 GGufMetaWriter { - writer: GGufWriter, - alignment: usize, -} - -pub struct GGufTensorWriter<'t, T: Write> { - writer: GGufWriter, - alignment: usize, - data: Vec<&'t [u8]>, - offset: usize, -} - -pub struct GGufSimulator { - writer: GGufWriter, - alignment: usize, - direct: usize, - data: Vec, -} - -impl GGufMetaWriter { - #[inline] - pub fn new(writer: T, header: GGufFileHeader) -> Result { - Ok(Self { - writer: GGufWriter::new(writer, header)?, - alignment: DEFAULT_ALIGNMENT, - }) - } - - #[inline] - pub fn write_alignment(&mut self, align: usize) -> Result<()> { - self.alignment = self.writer.write_alignment(align)?; - Ok(()) - } - - #[inline] - pub fn write_meta_kv( - &mut self, - key: impl AsRef, - ty: GGufMetaDataValueType, - val: impl AsRef<[u8]>, - ) -> Result<()> { - if let Some(align) = self.writer.write_meta_kv(key, ty, val)? { - self.alignment = align; - } - Ok(()) - } - - #[inline] - 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.offset += pad(self.offset, self.alignment); - self.writer.write(self.offset as u64)?; - - let len = info.nbytes(); - self.offset += len; - self.data.push(&data[info.offset()..][..len]); - - Ok(()) - } - - pub fn finish(self) -> Result { - let Self { - mut writer, - alignment, - data, - .. - } = self; - - for data in data { - for _ in 0..pad(writer.written_bytes(), alignment) { - writer.write(0u8)?; - } - writer.write_bytes(data)?; - } - Ok(writer.written_bytes()) - } -} - -impl Default for GGufSimulator { - #[inline] - fn default() -> Self { - Self::new() - } -} - -impl GGufSimulator { - #[inline] - pub fn new() -> Self { - Self { - writer: GGufWriter::new(NWrite, GGufFileHeader::default()).unwrap(), - alignment: DEFAULT_ALIGNMENT, - direct: 0, - data: Vec::new(), - } - } - - #[inline] - pub fn with_alignment(align: usize) -> Self { - let mut ans = Self::new(); - ans.write_alignment(align); - ans - } - - #[inline] - pub fn write(&mut self, n: usize) { - self.direct += n; - } - - #[inline] - pub fn write_alignment(&mut self, align: usize) { - self.alignment = self.writer.write_alignment(align).unwrap(); - } - - #[inline] - pub fn write_meta_kv( - &mut self, - key: impl AsRef, - ty: GGufMetaDataValueType, - val: impl AsRef<[u8]>, - ) { - if let Some(align) = self.writer.write_meta_kv(key, ty, val).unwrap() { - self.alignment = align; - } - } - - pub fn write_tensor(&mut self, info: &GGufTensorInfo) { - self.writer.write_str(info.name()).unwrap(); - - let shape = info.shape(); - self.writer.write(shape.len() as u32).unwrap(); - self.writer.write_bytes(as_slice(shape)).unwrap(); - - self.writer.write(info.ggml_type()).unwrap(); - self.writer.write(0u64).unwrap(); - - self.data.push(info.nbytes()); - } - - pub fn written_bytes(&self) -> usize { - let mut total = self.writer.written_bytes() + self.direct; - for len in &self.data { - total += pad(total, self.alignment); - total += len; - } - total - } -} - -struct NWrite; - -impl Write for NWrite { - #[inline(always)] - fn write(&mut self, buf: &[u8]) -> Result { - Ok(buf.len()) - } - #[inline(always)] - fn flush(&mut self) -> Result<()> { - Ok(()) - } -} - -struct GGufWriter(Internal); - -impl GGufWriter { - #[inline] - pub fn new(writer: T, header: GGufFileHeader) -> Result { - 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 write_bytes(&mut self, val: &[u8]) -> Result<()> { - self.0.write_bytes(val) - } - - #[inline] - pub fn write(&mut self, val: U) -> Result<()> { - self.write_bytes(as_slice(&val)) - } - - #[inline] - pub fn write_str(&mut self, val: impl AsRef) -> Result<()> { - let val = val.as_ref(); - self.write_bytes(as_slice(&(val.len() as u64)))?; - self.write_bytes(val.as_bytes()) - } - - #[inline] - pub fn write_alignment(&mut self, align: usize) -> Result { - self.write_meta_kv( - GENERAL_ALIGNMENT, - GGufMetaDataValueType::U32, - (align as u32).to_le_bytes(), - ) - .map(Option::unwrap) - } - - pub fn write_meta_kv( - &mut self, - key: impl AsRef, - ty: GGufMetaDataValueType, - val: impl AsRef<[u8]>, - ) -> Result> { - let key = key.as_ref(); - let val = val.as_ref(); - - self.write_str(key)?; - self.write(ty)?; - self.write_bytes(val)?; - - Ok(if key == GENERAL_ALIGNMENT { - let &[a, b, c, d] = val else { - panic!("general.alignment must be an u32") - }; - Some(u32::from_le_bytes([a, b, c, d]) as _) - } else { - None - }) - } -} - -#[inline(always)] -fn as_slice(val: &T) -> &[u8] { - unsafe { from_raw_parts(val as *const _ as *const _, size_of_val(val)) } -} - -#[inline(always)] -const fn pad(pos: usize, alignment: usize) -> usize { - (alignment - pos % alignment) % alignment -} - -mod internal { - use std::io::{BufWriter, Result, Write}; - - pub(super) struct Internal(BufWriter, usize); - - impl Internal { - #[inline] - pub fn new(writer: T) -> Self { - Self(BufWriter::new(writer), 0) - } - - #[inline] - pub const fn written_bytes(&self) -> usize { - self.1 - } - - #[inline] - pub fn write_bytes(&mut self, val: &[u8]) -> Result<()> { - self.1 += val.len(); - self.0.write_all(val.as_ref()) - } - } -} diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index 064cd74..a0f6f5c 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -7,7 +7,8 @@ authors = ["YdrMaster "] [dependencies] ggus = { version = "*", path = "../ggus" } -clap = { version = "4.5", features = ["derive"] } +digit-layout = "0.0" memmap2 = "0.9" regex = "1.10" -digit-layout = "0.0" +clap = { version = "4.5", features = ["derive"] } +itertools = "0.13.0" diff --git a/xtask/src/convert/mod.rs b/xtask/src/convert/mod.rs index 0109093..b88f2e9 100644 --- a/xtask/src/convert/mod.rs +++ b/xtask/src/convert/mod.rs @@ -1,9 +1,9 @@ mod read; mod write; -use crate::{file_info::FileInfo, gguf_file::GGufError}; +use crate::file_info::FileInfo; use digit_layout::DigitLayout; -use ggus::GENERAL_ALIGNMENT; +use ggus::{GGufError, GENERAL_ALIGNMENT}; use memmap2::Mmap; use read::read_files; use regex::Regex; diff --git a/xtask/src/convert/read.rs b/xtask/src/convert/read.rs index 5ef9250..6eb253e 100644 --- a/xtask/src/convert/read.rs +++ b/xtask/src/convert/read.rs @@ -1,12 +1,12 @@ -use crate::gguf_file::{GGufError, GGufFile}; +use ggus::{GGuf, GGufError}; pub(super) fn read_files<'a>( files: impl IntoIterator + 'a, -) -> Result>, GGufError> { +) -> Result>, GGufError> { std::thread::scope(|s| { files .into_iter() - .map(|data| s.spawn(|| GGufFile::new(data))) + .map(|data| s.spawn(|| GGuf::scan(data))) .collect::>() .into_iter() .map(|t| t.join().unwrap()) diff --git a/xtask/src/file_info.rs b/xtask/src/file_info.rs index 2e3ab0d..7fbd40e 100644 --- a/xtask/src/file_info.rs +++ b/xtask/src/file_info.rs @@ -1,5 +1,4 @@ -use crate::YES; -use std::{fmt, path::PathBuf}; +use std::{cmp::max, fmt, path::PathBuf}; pub struct FileInfo { pub path: PathBuf, @@ -8,30 +7,56 @@ pub struct FileInfo { pub n_bytes: usize, } -impl fmt::Display for FileInfo { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let Self { - path, - n_tensors, - n_meta_kvs, - n_bytes, - } = self; - write!( - f, - "\ -{YES}Shard is written to: \"{}\" - Number of tensors: {n_tensors} - Number of meta kvs: {n_meta_kvs} - File size: {}", - path.display(), - MemSize(*n_bytes), - ) +pub fn show_file_info(file_info: &[FileInfo]) { + if file_info.is_empty() { + return; + } + + const PATH: &str = "Path"; + const SIZE: &str = "Size"; + const META_KVS: &str = "Meta KVs"; + const TENSORS: &str = "Tensors"; + + let mut max_path_len = PATH.len(); + let mut max_size_len = SIZE.len(); + let mut max_meta_kvs_len = META_KVS.len(); + let mut max_tensors_len = TENSORS.len(); + + let mut path = Vec::with_capacity(file_info.len()); + let mut size = Vec::with_capacity(file_info.len()); + let mut meta_kvs = Vec::with_capacity(file_info.len()); + let mut tensors = Vec::with_capacity(file_info.len()); + + for info in file_info { + let path_ = info.path.display().to_string(); + let size_ = MemSize(info.n_bytes).to_string(); + let meta_kvs_ = info.n_meta_kvs.to_string(); + let tensors_ = info.n_tensors.to_string(); + + max_path_len = max(max_path_len, path_.len()); + max_size_len = max(max_size_len, size_.len()); + max_meta_kvs_len = max(max_meta_kvs_len, meta_kvs_.len()); + max_tensors_len = max(max_tensors_len, tensors_.len()); + + path.push(path_); + size.push(size_); + meta_kvs.push(meta_kvs_); + tensors.push(tensors_); + } + + let line = format!("+-{0:-max_size_len$} | {m:>max_meta_kvs_len$} | {t:>max_tensors_len$} |"); } + println!("{line}"); } #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] #[repr(transparent)] -pub struct MemSize(pub usize); +struct MemSize(pub usize); impl fmt::Display for MemSize { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { diff --git a/xtask/src/filter.rs b/xtask/src/filter.rs index b94a305..3e95e4f 100644 --- a/xtask/src/filter.rs +++ b/xtask/src/filter.rs @@ -1,4 +1,4 @@ -use crate::{convert::ConvertArgs, name_pattern::compile_patterns}; +use crate::{convert::ConvertArgs, file_info::show_file_info, name_pattern::compile_patterns}; use std::path::PathBuf; #[derive(Args, Default)] @@ -39,10 +39,6 @@ impl FilterArgs { } .convert() .unwrap(); - - let [file] = &*files else { - unreachable!(); - }; - println!("{file}"); + show_file_info(&files); } } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 17b81a5..4c52029 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -1,7 +1,6 @@ mod convert; mod file_info; mod filter; -mod gguf_file; mod merge; mod name_pattern; mod shards; diff --git a/xtask/src/merge.rs b/xtask/src/merge.rs index 2070d75..eb05133 100644 --- a/xtask/src/merge.rs +++ b/xtask/src/merge.rs @@ -1,4 +1,6 @@ -use crate::{convert::ConvertArgs, name_pattern::compile_patterns, shards::Shards}; +use crate::{ + convert::ConvertArgs, file_info::show_file_info, name_pattern::compile_patterns, shards::Shards, +}; use std::path::PathBuf; #[derive(Args, Default)] @@ -34,10 +36,6 @@ impl MergeArgs { } .convert() .unwrap(); - - let [file] = &*files else { - unreachable!(); - }; - println!("{file}"); + show_file_info(&files); } } diff --git a/xtask/src/split.rs b/xtask/src/split.rs index 24b9e73..984a660 100644 --- a/xtask/src/split.rs +++ b/xtask/src/split.rs @@ -1,4 +1,6 @@ -use crate::{convert::ConvertArgs, name_pattern::compile_patterns, shards::Shards}; +use crate::{ + convert::ConvertArgs, file_info::show_file_info, name_pattern::compile_patterns, shards::Shards, +}; use std::{path::PathBuf, str::from_utf8}; #[derive(Args, Default)] @@ -62,10 +64,6 @@ impl SplitArgs { } .convert() .unwrap(); - - for file in files { - println!("{file}"); - println!(); - } + show_file_info(&files); } }