mod-trace 0.4.0__tar.gz → 0.4.2__tar.gz

{mod_trace-0.4.0 → mod_trace-0.4.2}/Cargo.lock

@@ -16,7 +16,7 @@ checksum = "6b947ae49db0d222b1dbc6b113ce7248a3fc3a6ca21b696717bfc000ba4484d8"
 
 [[package]]
 name = "mod-trace"
-version = "0.4.0"
+version = "0.4.2"
 dependencies = [
  "serde",
  "serde_json",

{mod_trace-0.4.0 → mod_trace-0.4.2}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mod-trace"
-version = "0.4.0"
+version = "0.4.2"
 edition = "2024"
 description = "Rust CLI for inspecting ML model artifacts without loading the framework"
 license = "MIT"

{mod_trace-0.4.0 → mod_trace-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mod-trace
-Version: 0.4.0
+Version: 0.4.2
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
@@ -500,9 +500,27 @@ cargo run -- inspect models/tiny-distilbert-base-cased/model_fixed.onnx
 
 Fixed shapes such as `[1, 8]` produce better numeric estimates than symbolic shapes such as `[batch, sequence]`.
 
+## PyTorch
+
+mod-trace reads PyTorch `torch.save` files (`.pt`, `.pth`, `.bin`, `.ckpt`) **natively — no torch, no Python**:
+
+```sh
+mod-trace inspect      model.pt
+mod-trace diff         old.pt new.pt
+mod-trace explain-diff old.pt new.pt
+mod-trace check --max-parameter-growth 30% old.pt new.pt
+```
+
+It parses the `torch.save` zip (pickled structure + raw tensor storages) and reports file size, tensor count, **estimated parameter count** (storage bytes ÷ dtype), **dominant dtype**, **recovered parameter/layer names**, and a sampled weight fingerprint that changes on a retrain/finetune. ZIP64 is handled for models over 4 GB.
+
+Limits, by design:
+- **Tensor shapes are not decoded** (that needs a full pickle interpreter) — you get counts, names, and dtype, not per-tensor shapes.
+- **Legacy (pre-1.6) pickle `.pt`** and older Hugging Face `pytorch_model.bin` recover names + a fingerprint but not parameter sizes (tensors aren't stored as zip entries there). Modern zip-format saves get the full report.
+- **`.safetensors` is a different format** and is not read by mod-trace.
+
 ### Exporting any PyTorch model to ONNX
 
-mod-trace does not read native PyTorch `.pt`/`.pth` files (those are Python pickles / TorchScript archives). The supported path is to export to ONNX, which is the usual serving format anyway. For a plain `nn.Module` the export is a single call:
+For richer graph-level detail (operators, attention layers), or if you only have a `.safetensors`/legacy file, export to ONNX — the usual serving format — which mod-trace reads fully. For a plain `nn.Module` the export is a single call:
 
 ```python
 import torch

{mod_trace-0.4.0 → mod_trace-0.4.2}/README.md

@@ -483,9 +483,27 @@ cargo run -- inspect models/tiny-distilbert-base-cased/model_fixed.onnx
 
 Fixed shapes such as `[1, 8]` produce better numeric estimates than symbolic shapes such as `[batch, sequence]`.
 
+## PyTorch
+
+mod-trace reads PyTorch `torch.save` files (`.pt`, `.pth`, `.bin`, `.ckpt`) **natively — no torch, no Python**:
+
+```sh
+mod-trace inspect      model.pt
+mod-trace diff         old.pt new.pt
+mod-trace explain-diff old.pt new.pt
+mod-trace check --max-parameter-growth 30% old.pt new.pt
+```
+
+It parses the `torch.save` zip (pickled structure + raw tensor storages) and reports file size, tensor count, **estimated parameter count** (storage bytes ÷ dtype), **dominant dtype**, **recovered parameter/layer names**, and a sampled weight fingerprint that changes on a retrain/finetune. ZIP64 is handled for models over 4 GB.
+
+Limits, by design:
+- **Tensor shapes are not decoded** (that needs a full pickle interpreter) — you get counts, names, and dtype, not per-tensor shapes.
+- **Legacy (pre-1.6) pickle `.pt`** and older Hugging Face `pytorch_model.bin` recover names + a fingerprint but not parameter sizes (tensors aren't stored as zip entries there). Modern zip-format saves get the full report.
+- **`.safetensors` is a different format** and is not read by mod-trace.
+
 ### Exporting any PyTorch model to ONNX
 
-mod-trace does not read native PyTorch `.pt`/`.pth` files (those are Python pickles / TorchScript archives). The supported path is to export to ONNX, which is the usual serving format anyway. For a plain `nn.Module` the export is a single call:
+For richer graph-level detail (operators, attention layers), or if you only have a `.safetensors`/legacy file, export to ONNX — the usual serving format — which mod-trace reads fully. For a plain `nn.Module` the export is a single call:
 
 ```python
 import torch

{mod_trace-0.4.0 → mod_trace-0.4.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "mod-trace"
-version = "0.4.0"
+version = "0.4.2"
 description = "Rust CLI for inspecting ML model artifacts without loading the framework"
 readme = "README.md"
 requires-python = ">=3.9"

{mod_trace-0.4.0 → mod_trace-0.4.2}/src/main.rs

@@ -1815,7 +1815,7 @@ fn artifact_kind(path: &str) -> ArtifactKind {
         Some("cbm") => return ArtifactKind::CatBoost,
         Some("lgb") => return ArtifactKind::LightGbm,
         Some("onnx") => return ArtifactKind::Onnx,
-        Some("pt") | Some("pth") => return ArtifactKind::PyTorch,
+        Some("pt") | Some("pth") | Some("bin") | Some("ckpt") => return ArtifactKind::PyTorch,
         Some("json") => return ArtifactKind::Json,
         _ => {}
     }
@@ -1824,11 +1824,22 @@ fn artifact_kind(path: &str) -> ArtifactKind {
         ArtifactKind::CatBoost
     } else if has_lgbm_signature(path) {
         ArtifactKind::LightGbm
+    } else if has_pickle_signature(path) {
+        // torch.save legacy pickle (0x80 proto marker) with an unknown extension
+        ArtifactKind::PyTorch
     } else {
         ArtifactKind::Unknown
     }
 }
 
+fn has_pickle_signature(path: &str) -> bool {
+    let mut header = [0u8; 1];
+    fs::File::open(path)
+        .and_then(|mut file| file.read_exact(&mut header))
+        .map(|_| header[0] == 0x80)
+        .unwrap_or(false)
+}
+
 fn has_cbm_signature(path: &str) -> bool {
     let mut header = [0u8; 4];
     fs::File::open(path)
@@ -3169,15 +3180,21 @@ fn print_pt_report(report: &pt::PtReport, limit: usize) {
     }
     println!();
     println!("Structure:");
-    println!("  Tensors (storages): {}", report.tensor_count);
-    println!(
-        "  Parameters (est):   {}",
-        format_count_human(report.estimated_parameter_count as usize)
-    );
-    println!(
-        "  Parameter bytes:    {}",
-        format_bytes(report.total_parameter_bytes as usize)
-    );
+    if report.is_zip {
+        println!("  Tensors (storages): {}", report.tensor_count);
+        println!(
+            "  Parameters (est):   {}",
+            format_count_human(report.estimated_parameter_count as usize)
+        );
+        println!(
+            "  Parameter bytes:    {}",
+            format_bytes(report.total_parameter_bytes as usize)
+        );
+    } else {
+        println!("  Tensors (storages): unknown (legacy pickle)");
+        println!("  Parameters (est):   unknown (legacy pickle does not store sizes as entries)");
+        println!("  Parameter bytes:    unknown (legacy pickle)");
+    }
     print_optional("Dominant dtype", report.dominant_dtype.as_deref());
     println!();
     println!("Parameter-like Internals:");

{mod_trace-0.4.0 → mod_trace-0.4.2}/src/pt.rs

@@ -35,10 +35,6 @@ impl PtReport {
     }
 }
 
-pub fn looks_like_pt(head: &[u8]) -> bool {
-    head.starts_with(b"PK\x03\x04") || head.first() == Some(&0x80)
-}
-
 const SAMPLE_PER_STORAGE: u64 = 1 << 20; // 1 MiB sampled per tensor for the weight fingerprint
 
 pub fn inspect(path: &str) -> Result<PtReport, String> {
@@ -231,7 +227,8 @@ fn inspect_zip(file: &mut File, path: &str, total: u64) -> Result<PtReport, Stri
 }
 
 fn inspect_legacy(file: &mut File, path: &str, total: u64) -> Result<PtReport, String> {
-    // Legacy (non-zip) pickle: we can only scan strings + fingerprint the bytes.
+    // Legacy (non-zip) pickle: tensors aren't stored as separate entries, so we
+    // recover names + dtype + a count of storage references, but not sizes.
     let bytes = read_at(file, 0, total as usize)?;
     let param_names = recover_param_names(&bytes);
     let (dtype, _) = dominant_dtype(&bytes);
@@ -250,6 +247,7 @@ fn inspect_legacy(file: &mut File, path: &str, total: u64) -> Result<PtReport, S
     })
 }
 
+
 fn file_name(path: &str) -> String {
     Path::new(path)
         .file_name()