Repository: lm-sys/FastChat
File: fastchat/model/apply_delta.py (Python)

Code before:
import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import LlamaTokenizer, AutoConfig
import gc
import shutil

def apply_delta(base_model_path, target_model_path, delta_path):

    print(f"Loading the delta from {delta_path}")
    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path, use_fast=False)


    DEFAULT_PAD_TOKEN = "[PAD]"
    base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=False)
    num_new_tokens = base_tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))
    print("num_new_tokens: ",num_new_tokens)

    delta_num = 0
    delta_state_dict = torch.load(f"{delta_path}/pytorch_model-0000{str(delta_num%3+1)}-of-00003.bin", map_location=torch.device('cpu'))
    config = AutoConfig.from_pretrained(base_model_path)

    

    print("Applying the delta")
    for i in range(42):
        if i < 10:
            file_name = f"pytorch_model-0000{str(i)}-of-00041.bin"
        else:
            file_name = f"pytorch_model-000{str(i)}-of-00041.bin"
        state_dict_path = f"{base_model_path}/{file_name}"

        state_dict = torch.load(state_dict_path, map_location=torch.device('cpu'))

        for name, param in tqdm(state_dict.items(), desc="Applying delta"):

            while(name not in delta_state_dict):
                delta_num = (delta_num + 1) % 3
                delta_state_dict = torch.load(f"{delta_path}/pytorch_model-0000{str(delta_num%3+1)}-of-00003.bin", map_location=torch.device('cpu'))
                gc.collect()

            if param.shape != delta_state_dict[name].shape:
                new_embeddings = torch.zeros(len(base_tokenizer),state_dict[name].shape[1])
                new_embeddings[:len(base_tokenizer)-num_new_tokens] = state_dict[name]
                state_dict[name] = new_embeddings
                
            state_dict[name] += delta_state_dict[name]

        # 保存更新后的 state_dict
        torch.save(state_dict, f"{target_model_path}/{file_name}")
        state_dict = None 

Review comment (suggestion):
Do not use a hardcoded number. Please
1. Ues the latest main branch of huggingface/transformers
2. Redo the llama weight conversion by following https://huggingface.co/docs/transformers/main/model_doc/llama
3. Test the following deltas
```
--delta lmsys/vicuna-7b-delta-v1.1
--delta lmsys/vicuna-13b-delta-v1.1
```

Code after:
import glob
import json
import os
import shutil
import tempfile

from huggingface_hub import snapshot_download
import torch
from torch import nn
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig


GB = 1 << 30


def split_files(model_path, tmp_path, split_size):
    if not os.path.exists(model_path):
        model_path = snapshot_download(repo_id=model_path)
    if not os.path.exists(tmp_path):
        os.makedirs(tmp_path)

    file_pattern = os.path.join(model_path, "pytorch_model-*.bin")
    files = glob.glob(file_pattern)

    part = 0
    for file_path in tqdm(files):
        state_dict = torch.load(file_path)
        new_state_dict = {}

        current_size = 0

        for name, param in state_dict.items():
            param_size = param.numel() * param.element_size()

            if current_size + param_size > split_size:
                new_file_name = f"pytorch_model-{part}.bin"
                new_file_path = os.path.join(tmp_path, new_file_name)
                torch.save(new_state_dict, new_file_path)
                current_size = 0
                new_state_dict = None
                gc.collect()
                new_state_dict = {}
                part += 1

            new_state_dict[name] = param
            current_size += param_size

        new_file_name = f"pytorch_model-{part}.bin"
        new_file_path = os.path.join(tmp_path, new_file_name)
        torch.save(new_state_dict, new_file_path)