| import torch | |
| from safetensors.torch import load_file, save_file | |
| # List of safetensor files | |
| safetensor_files = [ | |
| "model-00001-of-00003.safetensors", | |
| "model-00002-of-00003.safetensors", | |
| "model-00003-of-00003.safetensors" | |
| ] | |
| # Load weights from safetensor files | |
| merged_weights = {} | |
| for file in safetensor_files: | |
| weights = load_file(file) | |
| merged_weights.update(weights) | |
| # Save the merged weights into a single pytorch_model.bin file | |
| torch.save(merged_weights, "pytorch_model.bin") | |
| print("Merged weights saved to pytorch_model.bin") | |