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)