fix: discover checkpoint shards from their weight index

tests/test_checkpoint_index.py

@@ -0,0 +1,44 @@
+import ast
+import json
+from pathlib import Path
+from types import SimpleNamespace
+import tempfile
+import unittest
+
+source = Path(__file__).resolve().parents[1] / "utils/convert-hf-to-gguf-bitnet.py"
+tree = ast.parse(source.read_text())
+model = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "Model")
+method = next(node for node in model.body if isinstance(node, ast.FunctionDef) and node.name == "_get_part_names")
+format_method = next(node for node in model.body if isinstance(node, ast.FunctionDef) and node.name == "_is_model_safetensors")
+ns = {"json": json, "Path": Path, "Model": SimpleNamespace(count_model_parts=lambda root, suffix: len(list(root.glob("*" + suffix))))}
+exec(compile(ast.Module(body=[method, format_method], type_ignores=[]), str(source), "exec"), ns)
+
+class CheckpointIndexTests(unittest.TestCase):
+    def test_index_controls_shard_names_and_deduplicates(self):
+        for safetensors in (True, False):
+            with self.subTest(safetensors=safetensors), tempfile.TemporaryDirectory() as directory:
+                root = Path(directory)
+                index = "model.safetensors.index.json" if safetensors else "pytorch_model.bin.index.json"
+                extension = ".safetensors" if safetensors else ".bin"
+                names = ["custom-a" + extension, "custom-b" + extension]
+                for name in names + ["optimizer" + extension]: (root / name).touch()
+                (root / index).write_text(json.dumps({"weight_map": {"a": names[0], "b": names[1], "c": names[0]}}))
+                parts = ns["_get_part_names"](SimpleNamespace(dir_model=root, is_safetensors=safetensors, num_parts=3))
+                self.assertEqual(list(parts), names)
+                self.assertEqual(list(parts), names)
+
+    def test_pytorch_index_selects_format_over_unrelated_safetensors(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "pytorch_model.bin.index.json").write_text("{}")
+            (root / "optimizer.safetensors").touch()
+            self.assertFalse(ns["_is_model_safetensors"](SimpleNamespace(dir_model=root)))
+
+    def test_missing_indexed_shard_fails_before_conversion(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "model.safetensors.index.json").write_text(json.dumps({"weight_map": {"a": "missing.safetensors"}}))
+            with self.assertRaises(FileNotFoundError):
+                ns["_get_part_names"](SimpleNamespace(dir_model=root, is_safetensors=True, num_parts=2))
+
+if __name__ == "__main__": unittest.main()

utils/convert-hf-to-gguf-bitnet.py

@@ -244,9 +244,30 @@ def from_model_architecture(cls, arch):
             raise NotImplementedError(f'Architecture {arch!r} not supported!') from None
 
     def _is_model_safetensors(self) -> bool:
+        if (self.dir_model / "model.safetensors.index.json").is_file():
+            return True
+        if (self.dir_model / "pytorch_model.bin.index.json").is_file():
+            return False
         return Model.count_model_parts(self.dir_model, ".safetensors") > 0
 
     def _get_part_names(self):
+        index_name = "model.safetensors.index.json" if self.is_safetensors else "pytorch_model.bin.index.json"
+        index_path = self.dir_model / index_name
+        if index_path.is_file():
+            with index_path.open(encoding="utf-8") as f:
+                index = json.load(f)
+            # The index is authoritative: unrelated checkpoint/optimizer files
+            # must not change the shard count or synthesize different filenames.
+            part_names = tuple(dict.fromkeys(index["weight_map"].values()))
+            if not part_names:
+                raise ValueError(f"Checkpoint index has no shards: {index_path}")
+            for name in part_names:
+                path = Path(name)
+                if path.is_absolute() or ".." in path.parts:
+                    raise ValueError(f"Checkpoint shard must be inside model directory: {name}")
+                if not (self.dir_model / path).is_file():
+                    raise FileNotFoundError(f"Checkpoint shard not found: {self.dir_model / path}")
+            return part_names
         if self.is_safetensors:
             if self.num_parts == 1:  # there's only one .safetensors file
                 return ("model.safetensors",)