Spaces:
Running
Running
| import argparse | |
| import numpy as np | |
| import os | |
| import shutil | |
| import torch | |
| import torch.nn.functional as F | |
| from safetensors.torch import safe_open, save_file | |
| import logging | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO) | |
| def merge_tensors(tensor1: torch.Tensor, tensor2: torch.Tensor, p: float) -> torch.Tensor: | |
| """ | |
| Merge two tensors using dropout and scaling. | |
| Args: | |
| tensor1 (torch.Tensor): The first tensor. | |
| tensor2 (torch.Tensor): The second tensor. | |
| p (float): Dropout probability. | |
| Returns: | |
| torch.Tensor: The merged tensor. | |
| """ | |
| delta = tensor2 - tensor1 | |
| m = torch.from_numpy(np.random.binomial(1, p, delta.shape)).to(tensor1.device) | |
| delta_tilde = m * delta | |
| delta_hat = delta_tilde / (1 - p) | |
| return delta_hat | |
| def merge_safetensors(file_path1: str, file_path2: str, p: float, lambda_val: float) -> dict: | |
| """ | |
| Merge two safetensors files. | |
| Args: | |
| file_path1 (str): Path to the first safetensors file. | |
| file_path2 (str): Path to the second safetensors file. | |
| p (float): Dropout probability. | |
| lambda_val (float): Scaling factor. | |
| Returns: | |
| dict: A dictionary of merged tensors. | |
| """ | |
| merged_tensors = {} | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| with safe_open(file_path1, framework="pt", device="cpu") as f1, safe_open(file_path2, framework="pt", device="cpu") as f2: | |
| keys1 = set(f1.keys()) | |
| keys2 = set(f2.keys()) | |
| common_keys = keys1.intersection(keys2) | |
| for key in common_keys: | |
| tensor1 = f1.get_tensor(key).to(device) | |
| tensor2 = f2.get_tensor(key).to(device) | |
| tensor1, tensor2 = resize_tensors(tensor1, tensor2) | |
| merged_tensors[key] = (tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)).cpu() | |
| logging.info(f"Merging {key}") | |
| return merged_tensors | |
| class BinDataHandler: | |
| """ | |
| A handler for binary data files. | |
| """ | |
| def __init__(self, data: dict): | |
| self.data = data | |
| def get_tensor(self, key: str) -> torch.Tensor: | |
| return self.data[key] | |
| def read_tensors(file_path: str, ext: str) -> tuple: | |
| """ | |
| Read tensors from a file. | |
| Args: | |
| file_path (str): Path to the file. | |
| ext (str): File extension. | |
| Returns: | |
| tuple: A tuple containing the file handler and the set of keys. | |
| """ | |
| if ext == ".safetensors" and file_path.endswith(".safetensors"): | |
| f = safe_open(file_path, framework="pt", device="cpu") | |
| return f, set(f.keys()) | |
| if ext == ".bin" and file_path.endswith(".bin"): | |
| data = torch.load(file_path, map_location=torch.device('cpu')) | |
| f = BinDataHandler(data) | |
| return f, set(data.keys()) | |
| return None, None | |
| def resize_tensors(tensor1: torch.Tensor, tensor2: torch.Tensor) -> tuple: | |
| """ | |
| Resize tensors to ensure they have the same shape. | |
| Args: | |
| tensor1 (torch.Tensor): The first tensor. | |
| tensor2 (torch.Tensor): The second tensor. | |
| Returns: | |
| tuple: A tuple containing the resized tensors. | |
| """ | |
| if len(tensor1.shape) not in [1, 2]: | |
| return tensor1, tensor2 | |
| if tensor1.shape[-1] < tensor2.shape[-1]: | |
| padding_size = tensor2.shape[-1] - tensor1.shape[-1] | |
| tensor1 = F.pad(tensor1, (0, padding_size, 0, 0)) | |
| elif tensor2.shape[-1] < tensor1.shape[-1]: | |
| padding_size = tensor1.shape[-1] - tensor2.shape[-1] | |
| tensor2 = F.pad(tensor2, (0, padding_size, 0, 0)) | |
| if tensor1.shape[0] < tensor2.shape[0]: | |
| padding_size = tensor2.shape[0] - tensor1.shape[0] | |
| tensor1 = F.pad(tensor1, (0, 0, 0, padding_size)) | |
| elif tensor2.shape[0] < tensor1.shape[0]: | |
| padding_size = tensor1.shape[0] - tensor2.shape[0] | |
| tensor2 = F.pad(tensor2, (0, 0, 0, padding_size)) | |
| return tensor1, tensor2 | |
| def merge_folder(tensor_map: dict, directory_path: str, p: float, lambda_val: float) -> dict: | |
| """ | |
| Merge tensors from a directory of model files. | |
| Args: | |
| tensor_map (dict): A dictionary mapping tensor keys to their file paths. | |
| directory_path (str): Path to the directory containing model files. | |
| p (float): Dropout probability. | |
| lambda_val (float): Scaling factor. | |
| Returns: | |
| dict: A dictionary of merged tensors. | |
| """ | |
| keys1 = set(tensor_map.keys()) | |
| ext = None | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| for filename in os.listdir(directory_path): | |
| if filename.endswith(".safetensors"): | |
| ext = ".safetensors" | |
| if filename.endswith(".bin") and ext is None: | |
| ext = ".bin" | |
| if ext is None: | |
| raise FileNotFoundError("Could not find model files") | |
| for filename in os.listdir(directory_path): | |
| file_path = os.path.join(directory_path, filename) | |
| f, keys2 = read_tensors(file_path, ext) | |
| if keys2: | |
| common_keys = keys1.intersection(keys2) | |
| for key in common_keys: | |
| if "block_sparse_moe.gate" in key: | |
| tensor1 = tensor_map[key]['tensor'].to(device) | |
| tensor2 = f.get_tensor(key).to(device) | |
| tensor_map[key]['tensor'] = (tensor1 + tensor2) / 2.0 | |
| continue | |
| tensor1 = tensor_map[key]['tensor'].to(device) | |
| tensor2 = f.get_tensor(key).to(device) | |
| tensor1, tensor2 = resize_tensors(tensor1, tensor2) | |
| tensor_map[key]['tensor'] = (tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)).cpu() | |
| return tensor_map | |
| def map_tensors_to_files(directory_path: str) -> dict: | |
| """ | |
| Map tensors to their respective files in a directory. | |
| Args: | |
| directory_path (str): Path to the directory containing model files. | |
| Returns: | |
| dict: A dictionary mapping tensor keys to their file paths. | |
| """ | |
| tensor_map = {} | |
| for filename in os.listdir(directory_path): | |
| file_path = os.path.join(directory_path, filename) | |
| f, keys = read_tensors(file_path, '.safetensors') | |
| if keys: | |
| for key in keys: | |
| tensor = f.get_tensor(key) | |
| tensor_map[key] = {'filename': filename, 'shape': tensor.shape, 'tensor': tensor} | |
| return tensor_map | |
| def copy_nontensor_files(from_path: str, to_path: str): | |
| """ | |
| Copy non-tensor files from one directory to another. | |
| Args: | |
| from_path (str): Path to the source directory. | |
| to_path (str): Path to the destination directory. | |
| """ | |
| for filename in os.listdir(from_path): | |
| file_path = os.path.join(from_path, filename) | |
| if from_path != to_path and not filename.startswith(".") and not filename.startswith("README") and not filename.endswith(".bin") and not filename.endswith(".safetensors") and not filename.endswith(".pt") and not os.path.isdir(file_path): | |
| logging.info(f"Copying {file_path} to {to_path}") | |
| shutil.copyfile(file_path, to_path + '/' + filename) | |
| def save_tensor_map(tensor_map: dict, output_folder: str): | |
| """ | |
| Save the merged tensor map to the output directory. | |
| Args: | |
| tensor_map (dict): A dictionary of merged tensors. | |
| output_folder (str): Path to the output directory. | |
| """ | |
| metadata = {'format': 'pt'} | |
| by_filename = {} | |
| for key, value in tensor_map.items(): | |
| filename = value["filename"] | |
| tensor = value["tensor"] | |
| if filename not in by_filename: | |
| by_filename[filename] = {} | |
| by_filename[filename][key] = tensor | |
| for filename in sorted(by_filename.keys()): | |
| output_file = output_folder + '/' + filename | |
| logging.info(f"Saving: {output_file}") | |
| save_file(by_filename[filename], output_file, metadata=metadata) | |
| def main(): | |
| """ | |
| Main function to parse command-line arguments and orchestrate the merging process. | |
| """ | |
| parser = argparse.ArgumentParser(description='Merge two safetensor model files.') | |
| parser.add_argument('base_model', type=str, help='The base model safetensor file') | |
| parser.add_argument('second_model', type=str, help='The second model safetensor file') | |
| parser.add_argument('output_model', type=str, help='The output merged model safetensor file') | |
| parser.add_argument('-p', type=float, default=0.5, help='Dropout probability') | |
| parser.add_argument('-lambda', dest='lambda_val', type=float, default=1.0, help='Scaling factor for the weight delta') | |
| args = parser.parse_args() | |
| if os.path.isdir(args.base_model): | |
| if not os.path.exists(args.output_model): | |
| os.makedirs(args.output_model) | |
| tensor_map = map_tensors_to_files(args.base_model) | |
| tensor_map = merge_folder(tensor_map, args.second_model, args.p, args.lambda_val) | |
| copy_nontensor_files(args.base_model, args.output_model) | |
| save_tensor_map(tensor_map, args.output_model) | |
| else: | |
| merged = merge_safetensors(args.base_model, args.second_model, args.p, args.lambda_val) | |
| save_file(merged, args.output_model) | |
| if __name__ == '__main__': | |
| main() | |