from pathlib import Path from time import sleep from tqdm import tqdm import argparse import requests import git import merge import os import shutil import sys import yaml from huggingface_hub import snapshot_download from huggingface_hub import HfApi, hf_hub_download def parse_arguments(): parser = argparse.ArgumentParser(description="Merge HuggingFace models") parser.add_argument('repo_list', type=str, help='File containing list of repositories to merge, supports mergekit yaml or txt') parser.add_argument('output_dir', type=str, help='Directory for the merged models') parser.add_argument('-staging', type=str, default='./staging', help='Path to staging folder') 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') parser.add_argument('--dry', action='store_true', help='Run in dry mode without making any changes') return parser.parse_args() def repo_list_generator(file_path, default_p, default_lambda_val): _, file_extension = os.path.splitext(file_path) # Branching based on file extension if file_extension.lower() == '.yaml' or file_extension.lower() == ".yml": with open(file_path, 'r') as file: data = yaml.safe_load(file) for model_info in data['models']: model_name = model_info['model'] p = model_info.get('parameters', {}).get('weight', default_p) lambda_val = 1 / model_info.get('parameters', {}).get('density', default_lambda_val) yield model_name, p, lambda_val else: # Defaulting to txt file processing with open(file_path, "r") as file: repos_to_process = file.readlines() for repo in repos_to_process: yield repo.strip(), default_p, default_lambda_val def reset_directories(directories, dry_run): for directory in directories: if os.path.exists(directory): if dry_run: print(f"[DRY RUN] Would delete directory {directory}") else: # Check if the directory is a symlink if os.path.islink(directory): os.unlink(directory) # Remove the symlink else: shutil.rmtree(directory, ignore_errors=False) print(f"Directory {directory} deleted successfully.") def do_merge(tensor_map, staging_path, p, lambda_val, dry_run=False): if dry_run: print(f"[DRY RUN] Would merge with {staging_path}") else: try: print(f"Merge operation for {staging_path}") tensor_map = merge.merge_folder(tensor_map, staging_path, p, lambda_val) print("Merge operation completed successfully.") except Exception as e: print(f"Error during merge operation: {e}") return tensor_map def download_repo(repo_name, path, dry_run=False): if not os.path.exists(path): os.makedirs(path) api = HfApi() # Get the list of all files in the repository using HfApi repo_files = api.list_repo_files(repo_name) if dry_run: print(f"[DRY RUN] Would download the following files from {repo_name} to {path}:") for file_path in repo_files: print(file_path) else: print(f"Downloading the entire repository {repo_name} directly to {path}.") for file_path in repo_files: print(f"Downloading {path}/{file_path}...") # Download each file directly to the specified path hf_hub_download( repo_id=repo_name, filename=file_path, cache_dir=path, local_dir=path, # Store directly in the target directory local_dir_use_symlinks=False # Ensure symlinks are not used ) print(f"Repository {repo_name} downloaded successfully to {path}.") def should_create_symlink(repo_name): if os.path.exists(repo_name): return True, os.path.isfile(repo_name) return False, False def download_or_link_repo(repo_name, path, dry_run=False): symlink, is_file = should_create_symlink(repo_name) if symlink and is_file: os.makedirs(path, exist_ok=True) symlink_path = os.path.join(path, os.path.basename(repo_name)) os.symlink(repo_name, symlink_path) elif symlink: os.symlink(repo_name, path) else: download_repo(repo_name, path, dry_run) def delete_repo(path, dry_run=False): if dry_run: print(f"[DRY RUN] Would delete repository at {path}") else: try: shutil.rmtree(path) print(f"Repository at {path} deleted successfully.") except Exception as e: print(f"Error deleting repository at {path}: {e}") def get_max_vocab_size(repo_list): max_vocab_size = 0 repo_with_max_vocab = None for repo in repo_list: repo_name = repo[0].strip() url = f"https://huggingface.co/{repo_name}/raw/main/config.json" try: response = requests.get(url) response.raise_for_status() config = response.json() vocab_size = config.get("vocab_size", 0) if vocab_size > max_vocab_size: max_vocab_size = vocab_size repo_with_max_vocab = repo_name except requests.RequestException as e: print(f"Error fetching data from {url}: {e}") return max_vocab_size, repo_with_max_vocab def download_json_files(repo_name, file_paths, output_dir): base_url = f"https://huggingface.co/{repo_name}/raw/main/" for file_path in file_paths: url = base_url + file_path response = requests.get(url) if response.status_code == 200: with open(os.path.join(output_dir, os.path.basename(file_path)), 'wb') as file: file.write(response.content) else: print(f"Failed to download {file_path}") def process_repos(output_dir, base_model, staging_model, repo_list_file, p, lambda_val, dry_run=False): # Check if output_dir exists if os.path.exists(output_dir): sys.exit(f"Output directory '{output_dir}' already exists. Exiting to prevent data loss.") # Reset base and staging directories reset_directories([base_model, staging_model], dry_run) repo_list_gen = repo_list_generator(repo_list_file, p, lambda_val) repos_to_process = list(repo_list_gen) # Initial download for 'base_model' download_or_link_repo(repos_to_process[0][0].strip(), base_model, dry_run) tensor_map = merge.map_tensors_to_files(base_model) for i, repo in enumerate(tqdm(repos_to_process[1:], desc='Merging Repos')): repo_name = repo[0].strip() repo_p = repo[1] repo_lambda = repo[2] delete_repo(staging_model, dry_run) download_or_link_repo(repo_name, staging_model, dry_run) tensor_map = do_merge(tensor_map, staging_model, repo_p, repo_lambda, dry_run) os.makedirs(output_dir, exist_ok=True) merge.copy_nontensor_files(base_model, output_dir) # Handle LLMs that add tokens by taking the largest if os.path.exists(os.path.join(output_dir, 'config.json')): max_vocab_size, repo_name = get_max_vocab_size(repos_to_process) if max_vocab_size > 0: file_paths = ['config.json', 'special_tokens_map.json', 'tokenizer.json', 'tokenizer_config.json'] download_json_files(repo_name, file_paths, output_dir) reset_directories([base_model, staging_model], dry_run) merge.save_tensor_map(tensor_map, output_dir) if __name__ == "__main__": args = parse_arguments() staging_path = Path(args.staging) os.makedirs(args.staging, exist_ok=True) process_repos(args.output_dir, staging_path / 'base_model', staging_path / 'staging_model', args.repo_list, args.p, args.lambda_val, args.dry)