diff --git a/Cargo.toml b/Cargo.toml index 094cb11..8fa7249 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "repugnant-pickle" version = "0.0.1" edition = "2021" -resolve = "2" +resolver = "2" [features] default = [] diff --git a/README.md b/README.md index 4c7ab2c..21e6391 100644 --- a/README.md +++ b/README.md @@ -55,8 +55,8 @@ PyTorch file: Int(430348288), ])), Int(327378944), - Seq(Tuple, [Int(1024), Int(50277)]), - Seq(Tuple, [Int(1), Int(1024)]), + Seq(Tuple, [Int(50277), Int(1024)]), + Seq(Tuple, [Int(1024), Int(1)]), Bool(false), Global(Raw(GLOBAL("collections", "OrderedDict")), [ Seq(Tuple, []) @@ -106,8 +106,8 @@ RepugnantTorchTensors( storage_len: 430348288, storage_offset: 327378944, absolute_offset: 327445248, - shape: [1024, 50277], - stride: [1, 1024], + shape: [50277, 1024], + stride: [1024, 1], requires_grad: false, }, RepugnantTorchTensor { @@ -133,6 +133,10 @@ start at offset `327445248` in the Torch file, or if you access the specific file `archive/data/0` in the Torch ZIP, it would start at `327378944`. +It worked on the `.pth` and `.pt` LLM models I have. + +Source: Trust me bro. + ## Usage Look at the examples in [examples](examples/): diff --git a/src/eval.rs b/src/eval.rs index 28f31d3..83479ef 100644 --- a/src/eval.rs +++ b/src/eval.rs @@ -286,11 +286,11 @@ pub fn evaluate<'a>( stack.push(Value::Seq(SequenceType::Tuple, vec![t1])); } PickleOp::TUPLE2 => { - let (t1, t2) = (stack.pop()?, stack.pop()?); + let (t2, t1) = (stack.pop()?, stack.pop()?); stack.push(Value::Seq(SequenceType::Tuple, vec![t1, t2])); } PickleOp::TUPLE3 => { - let (t1, t2, t3) = (stack.pop()?, stack.pop()?, stack.pop()?); + let (t3, t2, t1) = (stack.pop()?, stack.pop()?, stack.pop()?); stack.push(Value::Seq(SequenceType::Tuple, vec![t1, t2, t3])); } PickleOp::APPEND => { diff --git a/src/torch.rs b/src/torch.rs index c973767..8dc80cb 100644 --- a/src/torch.rs +++ b/src/torch.rs @@ -62,6 +62,26 @@ pub enum TensorType { Unknown(String), } +impl TensorType { + /// Get the item size for this tensor type. However, + /// the type of Unknown tensor types is... well, + /// unknown. So you get 0 back there. + pub fn size(&self) -> usize { + match self { + TensorType::Float64 => 8, + TensorType::Float32 => 4, + TensorType::Float16 => 2, + TensorType::BFloat16 => 2, + TensorType::Int64 => 8, + TensorType::Int32 => 4, + TensorType::Int16 => 2, + TensorType::Int8 => 1, + TensorType::UInt8 => 1, + TensorType::Unknown(_) => 0, + } + } +} + impl FromStr for TensorType { type Err = std::convert::Infallible; @@ -123,6 +143,16 @@ pub struct RepugnantTorchTensor { #[derive(Debug, Clone, PartialEq, Eq)] pub struct RepugnantTorchTensors(pub Vec); +impl IntoIterator for RepugnantTorchTensors { + type Item = RepugnantTorchTensor; + + type IntoIter = std::vec::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + impl RepugnantTorchTensors { pub fn new_from_file>(filename: P) -> Result { let mut zp = zip::ZipArchive::new(File::open(filename)?)?; @@ -239,6 +269,7 @@ impl RepugnantTorchTensors { zf.compression() == zip::CompressionMethod::STORE, "Can't handle compressed files", ); + let offs = offs * stype.size() as u64; tensors.push(RepugnantTorchTensor { name: k.to_string(), device: sdev.to_string(),