| 11 | |
| 12 | |
| 13 | def apply_delta(base_model_path, target_model_path, delta_path): |
| 14 | print(f"Loading the base model from {base_model_path}") |
| 15 | base = AutoModelForCausalLM.from_pretrained( |
| 16 | base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True) |
| 17 | |
| 18 | print(f"Loading the delta from {delta_path}") |
| 19 | delta = AutoModelForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True) |
| 20 | delta_tokenizer = AutoTokenizer.from_pretrained(delta_path, use_fast=False) |
| 21 | |
| 22 | DEFAULT_PAD_TOKEN = "<pad>" |
| 23 | base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=False) |
| 24 | num_new_tokens = base_tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN)) |
| 25 | |
| 26 | base.resize_token_embeddings(len(base_tokenizer)) |
| 27 | input_embeddings = base.get_input_embeddings().weight.data |
| 28 | output_embeddings = base.get_output_embeddings().weight.data |
| 29 | input_embeddings[-num_new_tokens:] = 0 |
| 30 | output_embeddings[-num_new_tokens:] = 0 |
| 31 | |
| 32 | print("Applying the delta") |
| 33 | for name, param in tqdm(base.state_dict().items(), desc="Applying delta"): |
| 34 | assert name in delta.state_dict() |
| 35 | param.data += delta.state_dict()[name] |
| 36 | |
| 37 | print(f"Saving the target model to {target_model_path}") |
| 38 | base.save_pretrained(target_model_path) |
| 39 | delta_tokenizer.save_pretrained(target_model_path) |
| 40 | |
| 41 | |
| 42 | if __name__ == "__main__": |