Upload export_apochat_litert.py with huggingface_hub

export_apochat_litert.py

@@ -89,12 +89,28 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--shard-size",
         type=int,
-        default=5_000_000_000,
+        default=10_000_000_000,
         help="Target size in bytes per PyTorch safetensors shard",
     )
     return parser.parse_args()
 
 
+def _write_sharded_index(checkpoint_dir: Path) -> None:
+    """Write model.safetensors.index.json from shards in a directory."""
+    from safetensors import safe_open
+
+    weight_map: dict[str, str] = {}
+    shards = sorted(checkpoint_dir.glob("model-?????-of-?????.safetensors"))
+    for shard in shards:
+        with safe_open(str(shard), framework="np") as f:
+            for key in f.keys():
+                weight_map[key] = shard.name
+    index = {"metadata": {"total_size": sum(s.stat().st_size for s in shards)}, "weight_map": weight_map}
+    (checkpoint_dir / "model.safetensors.index.json").write_text(
+        json.dumps(index, indent=2, sort_keys=True), encoding="utf-8"
+    )
+
+
 def download_repo_files(repo_id: str, revision: str | None, local_dir: Path) -> None:
     """Download all non-weight files from the MLX repo into local_dir."""
     print(f"Downloading aux files from {repo_id} ...")
@@ -104,6 +120,9 @@ def download_repo_files(repo_id: str, revision: str | None, local_dir: Path) ->
     for fname in files:
         if fname.endswith(".safetensors"):
             continue
+        if fname.endswith(".safetensors.index.json"):
+            # We will regenerate the index if sharding, or use a single file.
+            continue
         print(f"  {fname}")
         hf_hub_download(
             repo_id=repo_id,
@@ -210,9 +229,15 @@ def dequantize_mlx_to_pytorch(
     # Rewrite the final shard names with the actual count.
     shards = sorted(output_dir.glob("model-?????-of-?????.safetensors"))
     total = len(shards)
-    for i, old in enumerate(shards):
-        new = old.with_name(f"model-{i:05d}-of-{total:05d}.safetensors")
-        old.rename(new)
+    if total == 1:
+        # Transformers / litert expect a single "model.safetensors" for unsharded checkpoints.
+        shards[0].rename(output_dir / "model.safetensors")
+    else:
+        for i, old in enumerate(shards):
+            new = old.with_name(f"model-{i:05d}-of-{total:05d}.safetensors")
+            old.rename(new)
+        # Generate a fresh index so transformers can load the sharded checkpoint.
+        _write_sharded_index(output_dir)
 
     print(f"Wrote {total} safetensors shard(s) to {output_dir}")