EADST

Merge Safetensors to Bin File

from safetensors.torch import load_file
import torch
import os

safetensors_dir = "./your_model_folder_path/"

# get the number of safetensors
number = 0
for root, _, files in os.walk(safetensors_dir):
    for file in files:
        if "-of-" in file:
            number = file.split("-of-")[-1].strip(".safetensors")
            break

# merge safetensors to bin
model_data = {}
for i in range(1, int(number)+1):
    file_path = os.path.join(safetensors_dir, f"model-0000{i}-of-{number}.safetensors")
    data = load_file(file_path)
    model_data.update(data)

# save bin
torch.save(model_data, "./your_model_folder_path/pytorch_model.bin")