| 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) |