pt-loader 0.1.4

Safe parser-based PyTorch checkpoint converter to safetensors
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

use pt_loader::{LoadOptions, PtCheckpoint};
use std::fs::File;
use std::io::Write;
use tempfile::tempdir;
use zip::write::SimpleFileOptions;
use zip::ZipWriter;

#[test]
fn rejects_non_zip_pt() {
  let dir = tempdir().expect("tempdir");
  let pt = dir.path().join("legacy.pt");
  std::fs::write(&pt, b"not a zip checkpoint").expect("write");

  let err = PtCheckpoint::load(&pt, LoadOptions::default()).expect_err("should fail");
  assert!(err.to_string().contains("torch zip checkpoints"));
}

#[test]
fn enforces_pickle_size_limit() {
  let dir = tempdir().expect("tempdir");
  let pt = dir.path().join("large.pt");

  let file = File::create(&pt).expect("create");
  let mut zip = ZipWriter::new(file);
  zip
    .start_file("archive/data.pkl", SimpleFileOptions::default())
    .expect("start");
  zip.write_all(&vec![0u8; 1024]).expect("write");
  zip
    .start_file("archive/data/0", SimpleFileOptions::default())
    .expect("start data");
  zip.write_all(&[0u8; 16]).expect("write data");
  zip.finish().expect("finish");

  let opts = LoadOptions {
    max_pickle_bytes: 32,
    ..LoadOptions::default()
  };

  let err = PtCheckpoint::load(&pt, opts).expect_err("should fail");
  assert!(err.to_string().contains("data.pkl"));
}

#[test]
fn sample_last_pt_no_longer_fails_on_call_build_state() {
  let crate_root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"));
  let path = crate_root.join("samples").join("last.pt");
  if !path.exists() {
    return;
  }

  let err = match PtCheckpoint::load(&path, LoadOptions::default()) {
    Ok(_) => return,
    Err(err) => err,
  };
  assert!(
    !err.to_string().contains("BUILD with non-empty state is not supported"),
    "regression: parser still fails with unsupported call BUILD state"
  );
}