Update merge.py

merge.py

@@ -5,21 +5,43 @@ import shutil
 import torch
 import torch.nn.functional as F
 from safetensors.torch import safe_open, save_file
+import logging
 
-def merge_tensors(tensor1, tensor2, p):
-    # Calculate the delta of the weights
+# 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
-    # Generate the mask m^t from Bernoulli distribution
     m = torch.from_numpy(np.random.binomial(1, p, delta.shape)).to(tensor1.dtype)
-    # Apply the mask to the delta to get δ̃^t
     delta_tilde = m * delta
-    # Scale the masked delta by the dropout rate to get δ̂^t
     delta_hat = delta_tilde / (1 - p)
     return delta_hat
 
-def merge_safetensors(file_path1, file_path2, p, lambda_val):
-    merged_tensors = {}
+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 = {}
     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())
@@ -30,18 +52,31 @@ def merge_safetensors(file_path1, file_path2, p, lambda_val):
             tensor2 = f2.get_tensor(key)
             tensor1, tensor2 = resize_tensors(tensor1, tensor2)
             merged_tensors[key] = tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)
-            print("merging", key)
+            logging.info(f"Merging {key}")
 
     return merged_tensors
 
-class BinDataHandler():
-    def __init__(self, data):
+class BinDataHandler:
+    """
+    A handler for binary data files.
+    """
+    def __init__(self, data: dict):
         self.data = data
 
-    def get_tensor(self, key):
+    def get_tensor(self, key: str) -> torch.Tensor:
         return self.data[key]
 
-def read_tensors(file_path, ext):
+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())
@@ -51,11 +86,20 @@ def read_tensors(file_path, ext):
         return f, set(data.keys())
     return None, None
 
-def resize_tensors(tensor1, tensor2):
+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
 
-    # Pad along the last dimension (width)
     if tensor1.shape[-1] < tensor2.shape[-1]:
         padding_size = tensor2.shape[-1] - tensor1.shape[-1]
         tensor1 = F.pad(tensor1, (0, padding_size, 0, 0))
@@ -63,7 +107,6 @@ def resize_tensors(tensor1, tensor2):
         padding_size = tensor1.shape[-1] - tensor2.shape[-1]
         tensor2 = F.pad(tensor2, (0, padding_size, 0, 0))
 
-    # Pad along the first dimension (height)
     if tensor1.shape[0] < tensor2.shape[0]:
         padding_size = tensor2.shape[0] - tensor1.shape[0]
         tensor1 = F.pad(tensor1, (0, 0, 0, padding_size))
@@ -73,18 +116,28 @@ def resize_tensors(tensor1, tensor2):
 
     return tensor1, tensor2
 
-def merge_folder(tensor_map, directory_path, p, lambda_val):
+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())
-    # Some repos have both bin and safetensors, choose safetensors if so
     ext = None
     for filename in os.listdir(directory_path):
-        # Default to safetensors
         if filename.endswith(".safetensors"):
             ext = ".safetensors"
         if filename.endswith(".bin") and ext is None:
             ext = ".bin"
     if ext is None:
-        raise "Could not find model files"
+        raise FileNotFoundError("Could not find model files")
 
     for filename in os.listdir(directory_path):
         file_path = os.path.join(directory_path, filename)
@@ -95,7 +148,7 @@ def merge_folder(tensor_map, directory_path, p, lambda_val):
                 if "block_sparse_moe.gate" in key:
                     tensor1 = tensor_map[key]['tensor']
                     tensor2 = f.get_tensor(key)
-                    tensor_map[key]['tensor'] = (tensor1 + tensor2) /2.0
+                    tensor_map[key]['tensor'] = (tensor1 + tensor2) / 2.0
                     continue
                 tensor1 = tensor_map[key]['tensor']
                 tensor2 = f.get_tensor(key)
@@ -103,27 +156,48 @@ def merge_folder(tensor_map, directory_path, p, lambda_val):
                 tensor_map[key]['tensor'] = tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)
     return tensor_map
 
-def map_tensors_to_files(directory_path):
-    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}
-
+                tensor_map[key] = {'filename': filename, 'shape': tensor.shape, 'tensor': tensor}
     return tensor_map
 
-def copy_nontensor_files(from_path, to_path):
+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):
-            print(f"Copying {file_path} to {to_path}")
-            shutil.copyfile(file_path, to_path+'/'+filename)
+            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.
 
-def save_tensor_map(tensor_map, output_folder):
+    Args:
+        tensor_map (dict): A dictionary of merged tensors.
+        output_folder (str): Path to the output directory.
+    """
     metadata = {'format': 'pt'}
     by_filename = {}
 
@@ -135,12 +209,14 @@ def save_tensor_map(tensor_map, output_folder):
         by_filename[filename][key] = tensor
 
     for filename in sorted(by_filename.keys()):
-        output_file = output_folder+'/'+filename
-        print("Saving:", output_file)
+        output_file = output_folder + '/' + filename
+        logging.info(f"Saving: {output_file}")
         save_file(by_filename[filename], output_file, metadata=metadata)
 
 def main():
-    # Parse command-line arguments
+    """
+    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')
@@ -162,4 +238,4 @@ def main():
         save_file(merged, args.output_model)
 
 if __name__ == '__main__':
-    main()
+    main()