| 73 | print(f"Initial Setup...") |
| 74 | |
| 75 | def recreate_model(best_cycle, first_model): |
| 76 | # Read the merge ratios from the corresponding text file |
| 77 | with open(os.path.join(os.getcwd(), f"{model_nameZ}_mergecycle{best_cycle}.txt"), "r", encoding='utf-8') as file: |
| 78 | lines = file.readlines() |
| 79 | # Parse the merge ratios |
| 80 | merge_ratios = list(map(float, lines[3].strip().strip('[]').split(','))) |
| 81 | # Load the second model |
| 82 | print(f"\nLoading Transient Recipient Parent Model {model_nameY} For Recreation of Model {model_nameZ} from cycle {best_cycle} to RAM...") |
| 83 | second_model = AutoModelForCausalLM.from_pretrained(second_model_path).to('cpu') |
| 84 | second_model.eval() |
| 85 | print("Recipient Loaded. Dtype: " + str(second_model.dtype)) |
| 86 | num_layers = first_model.config.num_hidden_layers |
| 87 | print("Number of Layers:", num_layers) |
| 88 | print("Merge Ratios:", merge_ratios) |
| 89 | # Merge the models according to the stored merge ratios |
| 90 | for i in range(num_layers): |
| 91 | first_ratio = merge_ratios[i] |
| 92 | second_ratio = 1 - first_ratio |
| 93 | merged_layer = (first_model.model.layers[i].state_dict(), second_model.model.layers[i].state_dict()) |
| 94 | for key in merged_layer[0].keys(): |
| 95 | merged_layer[0][key] = first_ratio * merged_layer[0][key] + second_ratio * merged_layer[1][key] |
| 96 | second_model.model.layers[i].load_state_dict(merged_layer[0]) |
| 97 | print("Merging Layer " + str(i)) |
| 98 | # Save the merged model |
| 99 | print(f"{Fore.YELLOW}\nSaving User Preference Cycle {best_cycle} to disk and copying files.{Style.RESET_ALL}") |
| 100 | second_model.save_pretrained(merged_model_path, max_shard_size=max_shard_size) |
| 101 | # List of files to copy to merged model dir |
| 102 | files_to_copy = ["special_tokens_map.json", "tokenizer_config.json", "vocab.json", "tokenizer.model", "generation_config.json", "added_tokens.json", "merges.txt"] |
| 103 | # Check for the existence of 'special_tokens_map.json' in both directories |
| 104 | first_model_has_special_tokens = os.path.exists(os.path.join(first_model_path, "special_tokens_map.json")) |
| 105 | second_model_has_special_tokens = os.path.exists(os.path.join(second_model_path, "special_tokens_map.json")) |
| 106 | # Decide the source directory based on the presence of 'special_tokens_map.json' |
| 107 | if first_model_has_special_tokens and not second_model_has_special_tokens: |
| 108 | src_dir = first_model_path |
| 109 | elif second_model_has_special_tokens or not first_model_has_special_tokens: |
| 110 | src_dir = second_model_path |
| 111 | # Copy each file to the new folder |
| 112 | for filename in files_to_copy: |
| 113 | src_path = os.path.join(src_dir, filename) |
| 114 | dst_path = os.path.join(merged_model_path, filename) |
| 115 | print(f"\nCopying files from dir: {src_path}") |
| 116 | print(f"To dir: {dst_path}") |
| 117 | try: |
| 118 | shutil.copy2(src_path, dst_path) |
| 119 | except FileNotFoundError: |
| 120 | print("\nFile " + filename + " not found in " + src_dir + ". Skipping (likely not important).") |
| 121 | del second_model |
| 122 | |
| 123 | |
| 124 | def review_files(first_model): |