Upload merge_sharded_safetensors.py

merge_sharded_safetensors.py

@@ -0,0 +1,58 @@
+import os
+from safetensors import safe_open
+from safetensors.torch import save_file
+import torch  # Needed for torch.cat
+
+def merge_safetensor_files(sftsr_files, output_file="model.safetensors"):
+    slices_dict = {}
+    metadata = {}
+
+    for idx, file in enumerate(sftsr_files):
+        with safe_open(file, framework="pt") as sf_tsr:
+            if idx == 0:
+                metadata = sf_tsr.metadata()
+            for key in sf_tsr.keys():
+                tensor = sf_tsr.get_tensor(key)
+                if key not in slices_dict:
+                    slices_dict[key] = []
+                slices_dict[key].append(tensor)
+    
+    merged_tensors = {}
+    for key, slices in slices_dict.items():
+        if len(slices) == 1:
+            merged_tensors[key] = slices[0]
+        else:
+            # Simple heuristic: find dim with mismatched size
+            ref_shape = slices[0].shape
+            concat_dim = None
+            for dim in range(len(ref_shape)):
+                dim_sizes = [s.shape[dim] for s in slices]
+                if len(set(dim_sizes)) > 1:
+                    concat_dim = dim
+                    break
+            if concat_dim is None:
+                concat_dim = 0  # fallback
+            merged_tensors[key] = torch.cat(slices, dim=concat_dim)
+            print(f"Merged key '{key}' from {len(slices)} slices along dim {concat_dim}")
+
+    os.makedirs(os.path.dirname(output_file), exist_ok=True)
+    save_file(merged_tensors, output_file, metadata)
+    print(f"Merged {len(sftsr_files)} shards into {output_file}")
+
+def get_safetensor_files(directory):
+    safetensors_files = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".safetensors"):
+                safetensors_files.append(os.path.join(root, file))
+    return safetensors_files
+
+if __name__ == "__main__":
+    safetensor_files = get_safetensor_files("./shards")
+    print(f"The following shards/chunks will be merged: {safetensor_files}")
+
+    default_output = "./output/merged_model.safetensors"
+    user_output = input(f"Enter output file path [{default_output}]: ").strip()
+    output_file = user_output if user_output else default_output
+
+    merge_safetensor_files(safetensor_files, output_file=output_file)