From 942894b71433790241ecdc7d7aadef6ca9b04a6a Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Wed, 9 Sep 2026 13:55:45 -0400
Subject: [PATCH 1/7] test: cover local binary parser success and failure paths

How to test: python -m pytest -q (76 passed on Python 3.11).
---
 tests/test_binary_parsers.py | 129 +++++++++++++++++++++++++++++++++++
 1 file changed, 129 insertions(+)
 create mode 100644 tests/test_binary_parsers.py

diff --git a/tests/test_binary_parsers.py b/tests/test_binary_parsers.py
new file mode 100644
index 0000000..3f44636
--- /dev/null
+++ b/tests/test_binary_parsers.py
@@ -0,0 +1,129 @@
+"""Small generated checkpoint fixtures; no model downloads or torch dependency."""
+
+import json
+import pickle
+import struct
+import zipfile
+from collections import OrderedDict
+
+import pytest
+
+from modelinfo.parsers.pytorch import parse_pytorch_header
+from modelinfo.parsers.safetensors import parse_safetensors_header
+
+
+def write_safetensors(path, header):
+    payload = json.dumps(header).encode()
+    path.write_bytes(struct.pack("<Q", len(payload)) + payload)
+
+
+def test_safetensors_extracts_header_without_reading_tensor_payload(tmp_path):
+    path = tmp_path / "model.safetensors"
+    header = {
+        "__metadata__": {"format": "pt"},
+        "weight": {"shape": [2, 3], "dtype": "F16", "data_offsets": [0, 12]},
+    }
+    write_safetensors(path, header)
+    assert parse_safetensors_header(str(path)) == header
+
+
+@pytest.mark.parametrize("data", [b"", b"1234567", struct.pack("<Q", 20) + b"{}"])
+def test_safetensors_rejects_truncated_headers(tmp_path, data):
+    path = tmp_path / "truncated.safetensors"
+    path.write_bytes(data)
+    with pytest.raises(EOFError, match="Unexpected end of file"):
+        parse_safetensors_header(str(path))
+
+
+def test_safetensors_rejects_oversized_header_before_reading_it(tmp_path):
+    path = tmp_path / "oversized.safetensors"
+    path.write_bytes(struct.pack("<Q", 100 * 1024 * 1024 + 1))
+    with pytest.raises(ValueError, match="exceeds maximum allowed size"):
+        parse_safetensors_header(str(path))
+
+
+def test_safetensors_rejects_malformed_json(tmp_path):
+    path = tmp_path / "invalid.safetensors"
+    path.write_bytes(struct.pack("<Q", 1) + b"{")
+    with pytest.raises(json.JSONDecodeError):
+        parse_safetensors_header(str(path))
+
+
+def test_safetensors_index_deduplicates_shards_and_reports_missing_files(tmp_path):
+    index = tmp_path / "model.safetensors.index.json"
+    shard = tmp_path / "part1.safetensors"
+    header = {
+        "__metadata__": {"format": "pt"},
+        "a": {"shape": [2], "dtype": "F16"},
+        "b": {"shape": [3], "dtype": "F32"},
+    }
+    write_safetensors(shard, header)
+    index.write_text(
+        json.dumps(
+            {
+                "weight_map": {
+                    "a": shard.name,
+                    "b": shard.name,
+                    "c": "missing.safetensors",
+                }
+            }
+        )
+    )
+    result = parse_safetensors_header(str(index))
+    assert result["a"] == header["a"]
+    assert result["b"] == header["b"]
+    assert "c" not in result
+    assert result["__metadata__"] == {
+        "missing_shards": 1,
+        "total_shards": 2,
+        "is_sharded": True,
+        "disk_size": shard.stat().st_size,
+    }
+
+
+def test_pytorch_reads_ordered_state_dict_from_archive(tmp_path):
+    path = tmp_path / "model.pt"
+    state = OrderedDict(
+        [
+            ("weight", {"shape": [2, 3], "dtype": "F16"}),
+            ("step", 12),
+        ]
+    )
+    with zipfile.ZipFile(path, "w") as archive:
+        archive.writestr("checkpoint/data.pkl", pickle.dumps(state))
+    assert parse_pytorch_header(str(path)) == {
+        "weight": {"shape": [2, 3], "dtype": "F16"},
+        "step": {"shape": [], "dtype": "F32"},
+    }
+
+
+def test_pytorch_rejects_non_zip_checkpoint(tmp_path):
+    path = tmp_path / "legacy.pt"
+    path.write_bytes(b"not a zip")
+    with pytest.raises(ValueError, match="not a valid zip archive"):
+        parse_pytorch_header(str(path))
+
+
+def test_pytorch_rejects_archive_without_metadata(tmp_path):
+    path = tmp_path / "empty.pt"
+    with zipfile.ZipFile(path, "w") as archive:
+        archive.writestr("checkpoint/version", b"3")
+    with pytest.raises(ValueError, match="Could not find data.pkl"):
+        parse_pytorch_header(str(path))
+
+
+def test_pytorch_rejects_unapproved_pickle_global(tmp_path):
+    path = tmp_path / "forbidden.pt"
+    # A harmless builtin still must not bypass the explicit unpickler allowlist.
+    with zipfile.ZipFile(path, "w") as archive:
+        archive.writestr("checkpoint/data.pkl", pickle.dumps(len))
+    with pytest.raises(pickle.UnpicklingError, match="forbidden"):
+        parse_pytorch_header(str(path))
+
+
+def test_pytorch_propagates_truncated_pickle_error(tmp_path):
+    path = tmp_path / "truncated.pt"
+    with zipfile.ZipFile(path, "w") as archive:
+        archive.writestr("checkpoint/data.pkl", b"")
+    with pytest.raises(EOFError):
+        parse_pytorch_header(str(path))

From 578fd84266a5a5a8d727fb1ec94f59d47ce46c3f Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Fri, 11 Sep 2026 16:22:17 -0400
Subject: [PATCH 2/7] Use fixed pickle protocol fixtures in parser tests

---
 tests/test_binary_parsers.py | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/tests/test_binary_parsers.py b/tests/test_binary_parsers.py
index 3f44636..9ec0bf4 100644
--- a/tests/test_binary_parsers.py
+++ b/tests/test_binary_parsers.py
@@ -1,10 +1,9 @@
 """Small generated checkpoint fixtures; no model downloads or torch dependency."""
 
 import json
-import pickle
 import struct
 import zipfile
-from collections import OrderedDict
+from pickle import UnpicklingError
 
 import pytest
 
@@ -83,14 +82,11 @@ def test_safetensors_index_deduplicates_shards_and_reports_missing_files(tmp_pat
 
 def test_pytorch_reads_ordered_state_dict_from_archive(tmp_path):
     path = tmp_path / "model.pt"
-    state = OrderedDict(
-        [
-            ("weight", {"shape": [2, 3], "dtype": "F16"}),
-            ("step", 12),
-        ]
-    )
+    # Fixed protocol-0 OrderedDict fixture: inspect metadata without generating
+    # executable serialization during the test.
+    state = b"ccollections\nOrderedDict\np0\n(tRp1\nVweight\np2\n(dp3\nVshape\np4\n(lp5\nI2\naI3\nasVdtype\np6\nVF16\np7\nssVstep\np8\nI12\ns."
     with zipfile.ZipFile(path, "w") as archive:
-        archive.writestr("checkpoint/data.pkl", pickle.dumps(state))
+        archive.writestr("checkpoint/data.pkl", state)
     assert parse_pytorch_header(str(path)) == {
         "weight": {"shape": [2, 3], "dtype": "F16"},
         "step": {"shape": [], "dtype": "F32"},
@@ -116,8 +112,8 @@ def test_pytorch_rejects_unapproved_pickle_global(tmp_path):
     path = tmp_path / "forbidden.pt"
     # A harmless builtin still must not bypass the explicit unpickler allowlist.
     with zipfile.ZipFile(path, "w") as archive:
-        archive.writestr("checkpoint/data.pkl", pickle.dumps(len))
-    with pytest.raises(pickle.UnpicklingError, match="forbidden"):
+        archive.writestr("checkpoint/data.pkl", b"cbuiltins\nlen\n.")
+    with pytest.raises(UnpicklingError, match="forbidden"):
         parse_pytorch_header(str(path))
 
 

From 9f3375d8908edd211db7a47c6cf729fda9cb2586 Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Tue, 15 Sep 2026 00:08:23 -0400
Subject: [PATCH 3/7] Read supported PyTorch tensor storage references without
 loading payloads

---
 src/modelinfo/parsers/pytorch.py | 13 +++++++++++++
 tests/test_pytorch_storage.py    | 29 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 tests/test_pytorch_storage.py

diff --git a/src/modelinfo/parsers/pytorch.py b/src/modelinfo/parsers/pytorch.py
index 3926f2b..a04f3ff 100644
--- a/src/modelinfo/parsers/pytorch.py
+++ b/src/modelinfo/parsers/pytorch.py
@@ -49,6 +49,19 @@ class RestrictedUnpickler(pickle.Unpickler):
         },
     }
 
+    def persistent_load(self, saved_id: Any) -> DummyStorage:
+        # ZIP checkpoints refer to tensor storage via a five-element ID.
+        # Only reconstruct metadata using our existing dummy allowlist; never
+        # load storage bytes, import torch, or invoke a pickle-provided callable.
+        if not isinstance(saved_id, tuple) or len(saved_id) != 5:
+            raise pickle.UnpicklingError("Invalid storage persistent ID")
+        kind, storage_type, _, _, _ = saved_id
+        if (kind != "storage" or not isinstance(storage_type, type)
+                or not issubclass(storage_type, DummyStorage)
+                or storage_type is DummyStorage):
+            raise pickle.UnpicklingError("Unsupported storage persistent ID")
+        return storage_type()
+
     def find_class(self, module: str, name: str) -> Any:
         if module in self.ALLOWED_MODULES and name in self.ALLOWED_MODULES[module]:
             if name == "OrderedDict":
diff --git a/tests/test_pytorch_storage.py b/tests/test_pytorch_storage.py
new file mode 100644
index 0000000..72dc65a
--- /dev/null
+++ b/tests/test_pytorch_storage.py
@@ -0,0 +1,29 @@
+"""Fixed protocol-2 fixture using torch.save's storage persistent ID layout."""
+import io
+import pickle
+import zipfile
+
+import pytest
+
+from modelinfo.parsers.pytorch import RestrictedUnpickler, parse_pytorch_header
+
+
+@pytest.mark.parametrize('storage,dtype', [('FloatStorage', 'F32'), ('HalfStorage', 'F16'),
+    ('BFloat16Storage', 'BF16'), ('IntStorage', 'I32'), ('LongStorage', 'I64')])
+def test_tensor_with_persistent_storage_id(tmp_path, storage, dtype):
+    payload = (b'\x80\x02}X\x06\x00\x00\x00weightctorch._utils\n_rebuild_tensor_v2\n'
+               b'((X\x07\x00\x00\x00storagectorch\n' + storage.encode() + b'\n'
+               b'X\x01\x00\x00\x000X\x03\x00\x00\x00cpuK\x06tQK\x00K\x02K\x03\x86'
+               b'K\x03K\x01\x86\x89ccollections\nOrderedDict\n)RtRs.')
+    path = tmp_path / 'model.pt'
+    with zipfile.ZipFile(path, 'w') as archive:
+        archive.writestr('model/data.pkl', payload)
+    assert parse_pytorch_header(str(path)) == {'weight': {'shape': [2, 3], 'dtype': dtype}}
+
+
+@pytest.mark.parametrize('pid', [None, (), ('module',), ('storage', int, '0', 'cpu', 6),
+    ('storage', None, '0', 'cpu', 6)])
+def test_invalid_persistent_ids_are_rejected(pid):
+    unpickler = RestrictedUnpickler(io.BytesIO())
+    with pytest.raises(pickle.UnpicklingError):
+        unpickler.persistent_load(pid)

From 314757b5b87ddd6cbac32ffade419764cf7b0a2d Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Tue, 15 Sep 2026 00:15:40 -0400
Subject: [PATCH 4/7] Reject truncated GGUF headers and tolerate partial stream
 reads

---
 src/modelinfo/parsers/gguf.py | 66 ++++++++++++++++++++---------------
 tests/test_gguf_truncation.py | 35 +++++++++++++++++++
 2 files changed, 73 insertions(+), 28 deletions(-)
 create mode 100644 tests/test_gguf_truncation.py

diff --git a/src/modelinfo/parsers/gguf.py b/src/modelinfo/parsers/gguf.py
index 3af2fb4..c33eb06 100644
--- a/src/modelinfo/parsers/gguf.py
+++ b/src/modelinfo/parsers/gguf.py
@@ -10,36 +10,46 @@
     29: "IQ1_M", 30: "BF16", 31: "Q4_0_4_4", 32: "Q4_0_4_8", 33: "Q4_0_8_8",
 }
 
+def _read_exact(f: Any, size: int) -> bytes:
+    data = bytearray()
+    while len(data) < size:
+        chunk = f.read(size - len(data))
+        if not chunk:
+            raise EOFError("Unexpected end of GGUF header")
+        data.extend(chunk)
+    return bytes(data)
+
+
 def _read_gguf_value(f: Any, val_type: int) -> Any:
     if val_type == 0:
-        return struct.unpack("<B", f.read(1))[0]
+        return struct.unpack("<B", _read_exact(f, 1))[0]
     elif val_type == 1:
-        return struct.unpack("<b", f.read(1))[0]
+        return struct.unpack("<b", _read_exact(f, 1))[0]
     elif val_type == 2:
-        return struct.unpack("<H", f.read(2))[0]
+        return struct.unpack("<H", _read_exact(f, 2))[0]
     elif val_type == 3:
-        return struct.unpack("<h", f.read(2))[0]
+        return struct.unpack("<h", _read_exact(f, 2))[0]
     elif val_type == 4:
-        return struct.unpack("<I", f.read(4))[0]
+        return struct.unpack("<I", _read_exact(f, 4))[0]
     elif val_type == 5:
-        return struct.unpack("<i", f.read(4))[0]
+        return struct.unpack("<i", _read_exact(f, 4))[0]
     elif val_type == 6:
-        return struct.unpack("<f", f.read(4))[0]
+        return struct.unpack("<f", _read_exact(f, 4))[0]
     elif val_type == 7:
-        return struct.unpack("<?", f.read(1))[0]
+        return struct.unpack("<?", _read_exact(f, 1))[0]
     elif val_type == 8:
-        slen = struct.unpack("<Q", f.read(8))[0]
-        return f.read(slen).decode("utf-8")
+        slen = struct.unpack("<Q", _read_exact(f, 8))[0]
+        return _read_exact(f, slen).decode("utf-8")
     elif val_type == 9:
-        arr_type = struct.unpack("<I", f.read(4))[0]
-        arr_len = struct.unpack("<Q", f.read(8))[0]
+        arr_type = struct.unpack("<I", _read_exact(f, 4))[0]
+        arr_len = struct.unpack("<Q", _read_exact(f, 8))[0]
         return [_read_gguf_value(f, arr_type) for _ in range(arr_len)]
     elif val_type == 10:
-        return struct.unpack("<Q", f.read(8))[0]
+        return struct.unpack("<Q", _read_exact(f, 8))[0]
     elif val_type == 11:
-        return struct.unpack("<q", f.read(8))[0]
+        return struct.unpack("<q", _read_exact(f, 8))[0]
     elif val_type == 12:
-        return struct.unpack("<d", f.read(8))[0]
+        return struct.unpack("<d", _read_exact(f, 8))[0]
     else:
         raise ValueError(f"Unknown GGUF value type: {val_type}")
 
@@ -55,37 +65,37 @@ def parse_gguf_header(path_or_file: str | Any) -> Dict[str, Any]:
 
 def _parse_gguf_header_from_stream(f: Any) -> Dict[str, Any]:
     tensors: Dict[str, Any] = {}
-    magic = f.read(4)
+    magic = _read_exact(f, 4)
     if magic != b"GGUF":
         raise ValueError("Invalid GGUF file: Magic bytes missing.")
         
-    version = struct.unpack("<I", f.read(4))[0]
+    version = struct.unpack("<I", _read_exact(f, 4))[0]
     if version < 2:
         raise ValueError(f"Unsupported GGUF version: {version}")
         
-    tensor_count = struct.unpack("<Q", f.read(8))[0]
-    kv_count = struct.unpack("<Q", f.read(8))[0]
+    tensor_count = struct.unpack("<Q", _read_exact(f, 8))[0]
+    kv_count = struct.unpack("<Q", _read_exact(f, 8))[0]
     
     metadata = {}
     for _ in range(kv_count):
-        key_len = struct.unpack("<Q", f.read(8))[0]
-        key_name = f.read(key_len).decode("utf-8")
-        val_type = struct.unpack("<I", f.read(4))[0]
+        key_len = struct.unpack("<Q", _read_exact(f, 8))[0]
+        key_name = _read_exact(f, key_len).decode("utf-8")
+        val_type = struct.unpack("<I", _read_exact(f, 4))[0]
         metadata[key_name] = _read_gguf_value(f, val_type)
         
     tensors["__metadata__"] = metadata
         
     for _ in range(tensor_count):
-        name_len = struct.unpack("<Q", f.read(8))[0]
-        name = f.read(name_len).decode("utf-8")
+        name_len = struct.unpack("<Q", _read_exact(f, 8))[0]
+        name = _read_exact(f, name_len).decode("utf-8")
         
-        n_dims = struct.unpack("<I", f.read(4))[0]
+        n_dims = struct.unpack("<I", _read_exact(f, 4))[0]
         shape = []
         for _ in range(n_dims):
-            shape.append(struct.unpack("<Q", f.read(8))[0])
+            shape.append(struct.unpack("<Q", _read_exact(f, 8))[0])
         
-        t_type = struct.unpack("<I", f.read(4))[0]
-        f.read(8)  # skip offset bytes
+        t_type = struct.unpack("<I", _read_exact(f, 4))[0]
+        _read_exact(f, 8)  # skip offset bytes
         
         # Strict GGUF tensor type mapping
         dtype = GGML_TYPE_MAP.get(t_type, "Unknown")
diff --git a/tests/test_gguf_truncation.py b/tests/test_gguf_truncation.py
new file mode 100644
index 0000000..09228cf
--- /dev/null
+++ b/tests/test_gguf_truncation.py
@@ -0,0 +1,35 @@
+import io
+import struct
+
+import pytest
+
+from modelinfo.parsers.gguf import parse_gguf_header
+
+
+def tensor_header():
+    return (b'GGUF' + struct.pack('<IQQ', 3, 1, 0)
+            + struct.pack('<Q', 6) + b'weight'
+            + struct.pack('<IQIQ', 1, 3, 0, 0))
+
+
+@pytest.mark.parametrize('missing', range(1, 9))
+def test_rejects_truncated_final_tensor_offset(missing):
+    with pytest.raises(EOFError):
+        parse_gguf_header(io.BytesIO(tensor_header()[:-missing]))
+
+
+def test_rejects_truncated_final_metadata_string():
+    payload = (b'GGUF' + struct.pack('<IQQ', 3, 0, 1)
+               + struct.pack('<Q', 3) + b'key' + struct.pack('<IQ', 8, 5) + b'ab')
+    with pytest.raises(EOFError):
+        parse_gguf_header(io.BytesIO(payload))
+
+
+class ShortReads(io.BytesIO):
+    def read(self, size=-1):
+        return super().read(min(size, 2))
+
+
+def test_stream_can_return_partial_reads():
+    assert parse_gguf_header(ShortReads(tensor_header()))['weight'] == {
+        'shape': [3], 'dtype': 'F32'}

From 75fc7d56b490e2ed0d774dcb3d8488ffd5e176b7 Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Tue, 15 Sep 2026 00:26:29 -0400
Subject: [PATCH 5/7] Use shard index assignments when collecting tensor
 metadata

---
 src/modelinfo/parsers/safetensors.py |  9 ++++++---
 tests/test_shard_index_authority.py  | 25 +++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_shard_index_authority.py

diff --git a/src/modelinfo/parsers/safetensors.py b/src/modelinfo/parsers/safetensors.py
index 2e7d705..28bb39b 100644
--- a/src/modelinfo/parsers/safetensors.py
+++ b/src/modelinfo/parsers/safetensors.py
@@ -63,9 +63,12 @@ def parse_safetensors_header(path: str) -> dict[str, Any]:
             total_size += os.path.getsize(shard_path)
         try:
             shard_header = _read_single_header(shard_path)
-            for k, v in shard_header.items():
-                if k != "__metadata__":
-                    tensors[k] = v
+            for name, assigned_shard in weight_map.items():
+                if assigned_shard != shard:
+                    continue
+                if name not in shard_header:
+                    raise ValueError(f"Indexed tensor {name!r} is missing from shard {shard!r}")
+                tensors[name] = shard_header[name]
         except FileNotFoundError:
             missing_shards += 1
             
diff --git a/tests/test_shard_index_authority.py b/tests/test_shard_index_authority.py
new file mode 100644
index 0000000..1a57ef0
--- /dev/null
+++ b/tests/test_shard_index_authority.py
@@ -0,0 +1,25 @@
+import json
+import struct
+from modelinfo.parsers.safetensors import parse_safetensors_header
+
+
+def test_shard_index_selects_only_assigned_tensors(tmp_path):
+    index = tmp_path / 'model.safetensors.index.json'
+    index.write_text(json.dumps({'weight_map': {'weight': 'part.safetensors'}}))
+    header = {'weight': {'shape': [3], 'dtype': 'F16'},
+              'unused': {'shape': [1000], 'dtype': 'F32'}}
+    data = json.dumps(header).encode()
+    (tmp_path / 'part.safetensors').write_bytes(struct.pack('<Q', len(data)) + data)
+    result = parse_safetensors_header(str(index))
+    assert result['weight'] == header['weight']
+    assert 'unused' not in result
+
+
+def test_existing_shard_missing_an_indexed_tensor_is_rejected(tmp_path):
+    import pytest
+    index = tmp_path / 'model.safetensors.index.json'
+    index.write_text(json.dumps({'weight_map': {'weight': 'part.safetensors'}}))
+    data = json.dumps({'unrelated': {'shape': [1], 'dtype': 'F16'}}).encode()
+    (tmp_path / 'part.safetensors').write_bytes(struct.pack('<Q', len(data)) + data)
+    with pytest.raises(ValueError, match='weight'):
+        parse_safetensors_header(str(index))

From 4eb16eb27c919025563e8d3bb4508fbc928d711e Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Wed, 16 Sep 2026 10:55:40 -0400
Subject: [PATCH 6/7] Inspect state dictionaries inside training checkpoint
 wrappers

---
 src/modelinfo/parsers/pytorch.py         |  5 +++++
 tests/test_pytorch_checkpoint_wrapper.py | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)
 create mode 100644 tests/test_pytorch_checkpoint_wrapper.py

diff --git a/src/modelinfo/parsers/pytorch.py b/src/modelinfo/parsers/pytorch.py
index a04f3ff..1ce2b7e 100644
--- a/src/modelinfo/parsers/pytorch.py
+++ b/src/modelinfo/parsers/pytorch.py
@@ -97,6 +97,11 @@ def parse_pytorch_header(path: str) -> Dict[str, Any]:
             data = unpickler.load()
             
     if isinstance(data, dict):
+        for key in ("state_dict", "model_state_dict"):
+            candidate = data.get(key)
+            if isinstance(candidate, dict) and "shape" not in candidate:
+                data = candidate
+                break
         for k, v in data.items():
             if isinstance(v, dict) and "shape" in v:
                 tensors[k] = v
diff --git a/tests/test_pytorch_checkpoint_wrapper.py b/tests/test_pytorch_checkpoint_wrapper.py
new file mode 100644
index 0000000..67e1e2c
--- /dev/null
+++ b/tests/test_pytorch_checkpoint_wrapper.py
@@ -0,0 +1,16 @@
+import pickle
+import zipfile
+
+import pytest
+
+from modelinfo.parsers.pytorch import parse_pytorch_header
+
+
+@pytest.mark.parametrize('wrapper', ['state_dict', 'model_state_dict'])
+def test_wrapped_state_dict_excludes_training_metadata(tmp_path, wrapper):
+    weight = {'shape': [2, 3], 'dtype': 'F32'}
+    data = {wrapper: {'layer.weight': weight}, 'epoch': 10, 'loss': 0.125}
+    path = tmp_path / 'checkpoint.pt'
+    with zipfile.ZipFile(path, 'w') as archive:
+        archive.writestr('checkpoint/data.pkl', pickle.dumps(data, protocol=2))
+    assert parse_pytorch_header(str(path)) == {'layer.weight': weight}

From 3ab619eac8dc1f6c007be1473f69fd937eea7cb8 Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Wed, 16 Sep 2026 10:58:31 -0400
Subject: [PATCH 7/7] Recognize double, small integer and boolean PyTorch
 storage

---
 src/modelinfo/parsers/pytorch.py | 15 +++++++++++++++
 tests/test_pytorch_storage.py    |  4 +++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/modelinfo/parsers/pytorch.py b/src/modelinfo/parsers/pytorch.py
index 1ce2b7e..712bbda 100644
--- a/src/modelinfo/parsers/pytorch.py
+++ b/src/modelinfo/parsers/pytorch.py
@@ -28,6 +28,16 @@ def dummy_rebuild_tensor_v2(
             dtype = "I32"
         elif storage.name == "LongStorage":
             dtype = "I64"
+        elif storage.name == "DoubleStorage":
+            dtype = "F64"
+        elif storage.name == "ShortStorage":
+            dtype = "I16"
+        elif storage.name == "CharStorage":
+            dtype = "I8"
+        elif storage.name == "ByteStorage":
+            dtype = "U8"
+        elif storage.name == "BoolStorage":
+            dtype = "BOOL"
     return {"shape": list(size), "dtype": dtype}
 
 
@@ -46,6 +56,11 @@ class RestrictedUnpickler(pickle.Unpickler):
             "BFloat16Storage",
             "IntStorage",
             "LongStorage",
+            "DoubleStorage",
+            "ShortStorage",
+            "CharStorage",
+            "ByteStorage",
+            "BoolStorage",
         },
     }
 
diff --git a/tests/test_pytorch_storage.py b/tests/test_pytorch_storage.py
index 72dc65a..7adbdb7 100644
--- a/tests/test_pytorch_storage.py
+++ b/tests/test_pytorch_storage.py
@@ -9,7 +9,9 @@
 
 
 @pytest.mark.parametrize('storage,dtype', [('FloatStorage', 'F32'), ('HalfStorage', 'F16'),
-    ('BFloat16Storage', 'BF16'), ('IntStorage', 'I32'), ('LongStorage', 'I64')])
+    ('BFloat16Storage', 'BF16'), ('IntStorage', 'I32'), ('LongStorage', 'I64'),
+    ('DoubleStorage', 'F64'), ('ShortStorage', 'I16'), ('CharStorage', 'I8'),
+    ('ByteStorage', 'U8'), ('BoolStorage', 'BOOL')])
 def test_tensor_with_persistent_storage_id(tmp_path, storage, dtype):
     payload = (b'\x80\x02}X\x06\x00\x00\x00weightctorch._utils\n_rebuild_tensor_v2\n'
                b'((X\x07\x00\x00\x00storagectorch\n' + storage.encode() + b'\n'