Update bin2safetensors/convert.py

bin2safetensors/convert.py

@@ -2,6 +2,7 @@ import argparse
 import json
 import os
 import shutil
+import re
 from collections import defaultdict
 from inspect import signature
 from tempfile import TemporaryDirectory
@@ -311,14 +312,59 @@ def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitIn
         return new_pr, errors
 
 
-if __name__ == "__main__":
-    DESCRIPTION = """
-    Simple utility tool to convert automatically some weights on the hub to `safetensors` format.
-    It is PyTorch exclusive for now.
-    It works by downloading the weights (PT), converting them locally, and uploading them back
-    as a PR on the hub.
-    """
-    for i in range(1, 16):  # Range starts at 1 and ends at 15
-        input_filename = f"jondurbin_airoboros-l2-70b-gpt4-1.4.1/pytorch_model-{i:05d}-of-00015.bin"
-        output_filename = f"pytorch_model-{i:05d}-of-00015.safetensors"
+def main(input_directory, output_directory):
+    # Get a list of all files in the input directory
+    files = os.listdir(input_directory)
+    
+    # Filter the list to get only the relevant files
+    model_files = [file for file in files if re.match(r'pytorch_model-\d{5}-of-\d{5}\.bin', file)]
+    
+    # Determine the range for the loop based on the number of model files
+    num_models = len(model_files)
+    
+    if num_models == 0:
+        print("No model files found in the input directory.")
+        return
+    
+    # Extract yyyyy from the first model filename
+    match = re.search(r'pytorch_model-\d{5}-of-(\d{5})\.bin', model_files[0])
+    if match:
+        yyyyy = int(match.group(1))
+    else:
+        print("Unable to determine the number of shards from the filename.")
+        return
+    
+    if num_models != yyyyy:
+        print("Error: Number of shards mismatch.")
+        return
+    
+    # Copy *.json files (except pytorch_model.bin.index.json) from input to output directory
+    for file in files:
+        if file.endswith('.json') and not file == 'pytorch_model.bin.index.json':
+            src = os.path.join(input_directory, file)
+            dest = os.path.join(output_directory, file)
+            shutil.copy2(src, dest)
+            print(f"Copied {src} to {dest}")
+    
+    # Copy *.model files from input to output directory
+    for file in files:
+        if file.endswith('.model'):
+            src = os.path.join(input_directory, file)
+            dest = os.path.join(output_directory, file)
+            shutil.copy2(src, dest)
+            print(f"Copied {src} to {dest}")
+    
+    # Convert and rename model files
+    for i in range(1, num_models + 1):
+        input_filename = os.path.join(input_directory, f"pytorch_model-{i:05d}-of-{yyyyy:05d}.bin")
+        output_filename = os.path.join(output_directory, f"model-{i:05d}-of-{yyyyy:05d}.safetensors")
         convert_file(input_filename, output_filename)
+        print(f"Converted {input_filename} to {output_filename}")
+    
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert pytorch_model model to safetensor and copy JSON and .model files.")
+    parser.add_argument("input_directory", help="Path to the input directory containing pytorch_model files")
+    parser.add_argument("output_directory", help="Path to the output directory for converted safetensor files")
+    args = parser.parse_args()
+    
+    main(args.input_directory, args.output_directory)