codingutf8 2023 microsoft and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license weights conversion script for phi this script downloads both phi1 and phi1 5 checkpoints to checkpointpath and then converts the weights to hugfgingface model s format and saves them in pytorchdumpfolderpath we change names 124 layers023 for phi model layers converting the weights save either the whole model or the converted weights release gpu memory for the 2nd model if cuda was used release gpu memory for the 2nd model if cuda was used required parameters coding utf 8 2023 microsoft and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license weights conversion script for phi this script downloads both phi 1 and phi 1 5 checkpoints to checkpoint_path and then converts the weights to hugfgingface model s format and saves them in pytorch_dump_folder_path we change names 1 24 layers 0 23 for phi model layers phi 1 or phi 1_5 converting the weights save either the whole model or the converted weights release gpu memory for the 2nd model if cuda was used release gpu memory for the 2nd model if cuda was used required parameters
import argparse
import gc
import os

import torch
from huggingface_hub import hf_hub_download

from transformers import PhiConfig, PhiForCausalLM


_MODELS = {
    "microsoft/phi-1": "https://huggingface.co/microsoft/phi-1/blob/main/pytorch_model.bin",
    "microsoft/phi-1_5": "https://huggingface.co/microsoft/phi-1_5/blob/main/pytorch_model.bin",
}


PHI_MAPPING = {
    "layers.0.wte.weight": "model.embed_tokens.weight",
    "layers.25.linear.bias": "lm_head.bias",
    "layers.25.linear.weight": "lm_head.weight",
    "layers.25.ln.bias": "model.final_layernorm.bias",
    "layers.25.ln.weight": "model.final_layernorm.weight",
    "layers": "model.layers",
    "ln": "input_layernorm",
    "mixer": "self_attn",
    "Wqkv": "query_key_value",
    "out_proj": "dense",
}


def convert_weights(original_weights, mapping, config):
    converted_weights = {}
    original_weights_keys = sorted(original_weights.keys())

    
    range_change = {
        f"layers.{k}.": f"layers.{v}."
        for k, v in zip(range(1, config.num_hidden_layers + 1), range(0, config.num_hidden_layers))
    }

    mapping.update(**range_change)

    for original_weights_key in original_weights_keys:
        new_key = original_weights_key

        if "rotary_emb" in new_key:
            continue

        if "Wqkv" in new_key:
            if "weight" in new_key:
                weight = original_weights[new_key]
                weights_shape = weight.shape
                weight = (
                    weight.view(3, config.num_attention_heads, -1, config.hidden_size)
                    .transpose(0, 1)
                    .reshape(*weights_shape)
                )
                original_weights[new_key] = weight
            elif "bias" in new_key:
                bias = original_weights[new_key]
                bias_shape = bias.shape
                bias = bias.view(3, config.num_attention_heads, -1).transpose(0, 1).reshape(*bias_shape)
                original_weights[new_key] = bias

        for k, v in mapping.items():
            if k in new_key:
                new_key = new_key.replace(k, v)

        converted_weights[new_key] = original_weights.pop(original_weights_key)

    return converted_weights


def _download(url: str, root: str):
    repo_id = f"{url.split('/')[3]}/{url.split('/')[4]}"
    filename = f"{url.split('/')[-1]}"
    hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        force_filename=root,
        local_dir_use_symlinks=False,
    )


def convert_phi_weights(checkpoint_path, pytorch_dump_folder_path, use_cuda, save_weights_directly):
    device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
    for each_model_name, each_model_url in _MODELS.items():
        converted_checkpoint = {}

        model_path = os.path.join(checkpoint_path, each_model_name + "_" + each_model_url.split("/")[-1])
        if not os.path.exists(model_path):
            print(f"\n{each_model_name} was not found! Downloading it to {model_path}")
            _download(url=each_model_url, root=model_path)
        model_checkpoint = torch.load(model_path, map_location=device)
        model_type = each_model_name.split("/")[1]  
        config = PhiConfig.from_pretrained(f"susnato/{model_type}_dev")

        
        converted_checkpoint.update(**convert_weights(model_checkpoint, PHI_MAPPING, config))

        
        if save_weights_directly:
            save_weights_path = os.path.join(
                pytorch_dump_folder_path, each_model_name.split("/")[-1] + "_" + each_model_url.split("/")[-1]
            )
            torch.save(converted_checkpoint, save_weights_path)
            print(f"Model weights saved at {save_weights_path}!")

        else:
            model = PhiForCausalLM(config).to(device)
            model.load_state_dict(converted_checkpoint, strict=True)
            save_model_path = os.path.join(pytorch_dump_folder_path, model_type)
            model.save_pretrained(save_model_path)
            print(f"Model saved at {save_model_path}!")

            
            del config, model

        
        del model_checkpoint, converted_checkpoint
        if use_cuda:
            torch.cuda.empty_cache()
        gc.collect()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--checkpoint_path", type=str, help="Path to the folder of downloaded checkpoints. (Please enter full path)"
    )
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        help="Path to the output PyTorch model. (Please enter full path)",
    )
    parser.add_argument(
        "--use_cuda",
        default=False,
        type=bool,
        help="Whether to load the weights on GPU during conversion or not, False by default",
    )
    parser.add_argument(
        "--save_weights_directly",
        default=True,
        type=bool,
        help="Whether to save the weights directly after conversion or load the weight to the Phi model and then save "
        "the Phi model along with weights. True by default",
    )

    args = parser.parse_args()
    convert_phi_weights(args.checkpoint_path, args.pytorch_dump_folder_path, args.use_cuda, args.save_weights_directly)