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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

# pt-loader User Skill

Use this package to safely read Torch zip `.pt` checkpoints and export tensors to safetensors.

This skill is for end users consuming the package APIs. It is not a contributor/developer workflow guide.

## Python

Install:

```bash
pip install pt-safe-loader
```

Basic usage:

```python
from pt_loader import PtCheckpoint

ckpt = PtCheckpoint.load("model.pt")

# Inspect metadata
meta = ckpt.metadata()
print(meta["tensor_count"])
print(meta["security"]["objects"])
print(meta["security"]["calls"])

# Export safetensors + yaml
result = ckpt.export(format="safetensors", dir="out")
print(result["weights_path"])
print(result["metadata_path"])

# Load tensors in memory
tensors = ckpt.state_dict(backend="numpy")
print(next(iter(tensors.items()))[0], next(iter(tensors.values())).shape)
```

Advanced root-key extraction:

```python
ckpt = PtCheckpoint.load(
  "model.pt",
  state_dict_root_keys=["model", "optimizer", "ema.model"],  # deep keys supported
  state_dict_root_strict=False,  # skip missing keys; fail if none succeed
)

result = ckpt.export(format="safetensors", dir="out")
print(result["weights_paths"])  # per-root files, e.g. *.model.safetensors
```

Notes:
- `state_dict_root_key="model"` is still accepted for backward compatibility.
- `state_dict_root_keys=[]` means extract from root only.

## Rust

Add dependency:

```toml
[dependencies]
pt-loader = "0.1"
```

Basic usage:

```rust
use pt_loader::{ExportFormat, ExportOptions, LoadOptions, PtCheckpoint};
use std::path::Path;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = "model.pt";
    let ckpt = PtCheckpoint::load(input, LoadOptions::default())?;

    let meta = ckpt.metadata();
    println!("tensor_count={}", meta.tensor_count);
    println!("objects={:?}", meta.security.objects);
    println!("calls={:?}", meta.security.calls);

    let out = ckpt.export(
        "out",
        ExportOptions::new(ExportFormat::Safetensors, Some(Path::new(input))),
    )?;
    println!("weights={}", out.weights_path.display());
    println!("metadata={}", out.metadata_path.unwrap().display());
    Ok(())
}
```

Root-key extraction options:

```rust
use pt_loader::LoadOptions;

let opts = LoadOptions {
    state_dict_root_keys: vec![
        "model".to_string(),
        "optimizer".to_string(),
        "ema.model".to_string(), // deep key
    ],
    state_dict_root_strict: false,
    ..LoadOptions::default()
};
```

Notes:
- Empty `state_dict_root_keys` extracts from root only.
- Non-empty keys extract only those roots (no fallback scan).
- Strict mode fails fast on missing/invalid roots.