| 123 | |
| 124 | |
| 125 | def apply_delta(base_model_path, target_model_path, delta_path): |
| 126 | print(f"Loading the delta weights from {delta_path}") |
| 127 | delta_tokenizer = AutoTokenizer.from_pretrained(delta_path, use_fast=False) |
| 128 | delta = AutoModelForCausalLM.from_pretrained( |
| 129 | delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True |
| 130 | ) |
| 131 | |
| 132 | print(f"Loading the base model from {base_model_path}") |
| 133 | base = AutoModelForCausalLM.from_pretrained( |
| 134 | base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True |
| 135 | ) |
| 136 | |
| 137 | print("Applying the delta") |
| 138 | for name, param in tqdm(base.state_dict().items(), desc="Applying delta"): |
| 139 | assert name in delta.state_dict() |
| 140 | param.data += delta.state_dict()[name] |
| 141 | |
| 142 | print(f"Saving the target model to {target_model_path}") |
| 143 | base.save_pretrained(target_model_path) |
| 144 | delta_tokenizer.save_pretrained(target_model_path) |
| 145 | |
| 146 | |
| 147 | if __name__ == "__main__": |