Also convert .index.json file (#35)

The conversion to safetensors wasn't taking into account the
pytorch_model.bin.index.json file that needs to be converted to
model.safetensors.index.json. This also converts `bin.index.json` files

Signed-off-by: Joe Runde <[email protected]>
Co-authored-by: Maximilien Philippe Marie de Bayser <[email protected]>
Co-authored-by: Nick Hill <[email protected]>

Author: 3 people

server/tests/utils/test_convert.py

@@ -13,7 +13,7 @@ def test_convert_files():
     local_pt_files = download_weights(model_id, extension=".bin")
     local_pt_files = [Path(p) for p in local_pt_files]
     local_st_files = [
-        p.parent / f"{p.stem.lstrip('pytorch_')}.safetensors" for p in local_pt_files
+        p.parent / f"{p.stem.removeprefix('pytorch_')}.safetensors" for p in local_pt_files
     ]
     convert_files(local_pt_files, local_st_files, discard_names=[])
 

server/text_generation_server/cli.py

@@ -138,16 +138,21 @@ def convert_to_safetensors(
     # Get local pytorch file paths
     model_path = utils.get_model_path(model_name, revision)
     local_pt_files = utils.local_weight_files(model_path, ".bin")
+    local_pt_index_files = utils.local_index_files(model_path, ".bin")
+    if len(local_pt_index_files) > 1:
+        print(f"Found more than one .bin.index.json file: {local_pt_index_files}")
+        return
 
     if not local_pt_files:
         print("No pytorch .bin files found to convert")
         return
 
     local_pt_files = [Path(f) for f in local_pt_files]
+    local_pt_index_file = local_pt_index_files[0] if local_pt_index_files else None
 
     # Safetensors final filenames
     local_st_files = [
-        p.parent / f"{p.stem.lstrip('pytorch_')}.safetensors"
+        p.parent / f"{p.stem.removeprefix('pytorch_')}.safetensors"
         for p in local_pt_files
     ]
 
@@ -173,6 +178,16 @@ def convert_to_safetensors(
     except Exception:
         discard_names = []
 
+    if local_pt_index_file:
+        local_pt_index_file = Path(local_pt_index_file)
+        local_st_index_file = local_pt_index_file.parent / f"{local_pt_index_file.stem.removeprefix('pytorch_').rstrip('.bin.index')}.safetensors.index.json"
+
+        if os.path.exists(local_st_index_file):
+            print("Existing .safetensors.index.json file found, remove it first to reconvert")
+            return
+
+        utils.convert_index_file(local_pt_index_file, local_st_index_file, local_pt_files, local_st_files)
+
     # Convert pytorch weights to safetensors
     utils.convert_files(local_pt_files, local_st_files, discard_names)
 

server/text_generation_server/utils/__init__.py

@@ -1,4 +1,4 @@
-from text_generation_server.utils.convert import convert_file, convert_files
+from text_generation_server.utils.convert import convert_file, convert_files, convert_index_file
 from text_generation_server.utils.dist import (
     initialize_torch_distributed,
     run_rank_n,
@@ -10,6 +10,7 @@
 from text_generation_server.utils.hub import (
     get_model_path,
     local_weight_files,
+    local_index_files,
     weight_files,
     weight_hub_files,
     download_weights,
@@ -41,13 +42,15 @@
 __all__ = [
     "convert_file",
     "convert_files",
+    "convert_index_file",
     "initialize_torch_distributed",
     "run_rank_n",
     "print_rank_n",
     "get_torch_dtype",
     "RANK",
     "get_model_path",
     "local_weight_files",
+    "local_index_files",
     "weight_files",
     "weight_hub_files",
     "download_weights",

server/text_generation_server/utils/convert.py

@@ -1,6 +1,7 @@
 import datetime
 import torch
 import os
+import json
 
 from loguru import logger
 from pathlib import Path
@@ -87,12 +88,25 @@ def convert_file(pt_file: Path, sf_file: Path, discard_names: List[str]):
             raise RuntimeError(f"The output tensors do not match for key {k}")
 
 
+def convert_index_file(source_file: Path, dest_file: Path, pt_files: List[Path], sf_files: List[Path]):
+    weight_file_map = {s.name: d.name for s, d in zip(pt_files, sf_files)}
+
+    logger.info(f"Converting pytorch .bin.index.json files to .safetensors.index.json")
+    with open(source_file, "r") as f:
+        index = json.load(f)
+
+    index["weight_map"] = {k: weight_file_map[v] for k, v in index["weight_map"].items()}
+    
+    with open(dest_file, "w") as f:
+        json.dump(index, f, indent=4)
+
+
 def convert_files(pt_files: List[Path], sf_files: List[Path], discard_names: List[str] = None):
     assert len(pt_files) == len(sf_files)
 
     # Filter non-inference files
     pairs = [p for p in zip(pt_files, sf_files) if not any(
-        s in p[0].name for s in ["arguments", "args", "training", "optimizer", "scheduler"]
+        s in p[0].name for s in ["arguments", "args", "training", "optimizer", "scheduler", "index"]
     )]
 
     N = len(pairs)

server/text_generation_server/utils/hub.py

@@ -100,3 +100,9 @@ def local_weight_files(model_path: str, extension=".safetensors"):
     """Get the local safetensors filenames"""
     ext = "" if extension is None else extension
     return glob.glob(f"{model_path}/*{ext}")
+
+
+def local_index_files(model_path: str, extension=".safetensors"):
+    """Get the local .index.json filename"""
+    ext = "" if extension is None else extension
+    return glob.glob(f"{model_path}/*{ext}.index.json")