Add safetensors merge and split helper files

merge.py

@@ -0,0 +1,26 @@
+import os
+import math
+import json
+
+OUTPUT_FILE_NAME = "consolidated.safetensors" # Merge output file name
+CHUNK_PATHS_FILE = "chunk_paths.json"
+
+def merge(chunk_paths): 
+    output_path = os.path.join(os.path.dirname(chunk_paths[0]), OUTPUT_FILE_NAME)
+    
+    with open(output_path, "wb") as f_out:
+        for filepath in chunk_paths:
+            with open(filepath, "rb") as f_in:
+                f_out.write(f_in.read())
+                
+    print(f"Merged file saved to {output_path}")
+	
+if __name__ == "__main__":
+    
+    if os.path.exists(CHUNK_PATHS_FILE):
+        with open(CHUNK_PATHS_FILE) as f:
+            chunk_paths = json.load(f) 
+    else:
+        chunk_paths = split(main_filepath)
+          
+    merge(chunk_paths)	

split.py

@@ -0,0 +1,43 @@
+import os
+import math
+import json
+
+CHUNK_SIZE = 2 * 1024**3 # 40GB
+
+CHUNK_PATHS_FILE = "chunk_paths.json" 
+
+def split(filepath, chunk_size=CHUNK_SIZE):
+    basename = os.path.basename(filepath)
+    dirname = os.path.dirname(filepath)
+    extension = basename.split(".")[-1]
+    
+    filename_no_ext = basename.split(".")[-2] 
+    file_size = os.path.getsize(filepath)
+    
+    num_chunks = math.ceil(file_size / chunk_size)
+    digit_count = len(str(num_chunks))
+    
+    chunk_paths = []
+    
+    for i in range(1, num_chunks+1):
+        start = (i-1) * chunk_size
+        
+        chunk_filename = f"{filename_no_ext}-{str(i).zfill(digit_count)}-of-{str(num_chunks).zfill(digit_count)}.{extension}"
+        split_path = os.path.join(dirname, chunk_filename)
+        
+        with open(filepath, "rb") as f_in:
+            f_in.seek(start)
+            chunk = f_in.read(chunk_size)
+            
+        with open(split_path, "wb") as f_out:
+            f_out.write(chunk)
+            
+        chunk_paths.append(split_path)
+        
+    with open(CHUNK_PATHS_FILE, 'w') as f:
+        json.dump(chunk_paths, f)
+    
+    return chunk_paths
+
+main_filepath = "consolidated.safetensors" # File to be split
+chunk_paths = split(main_filepath)