codingutf8 2023 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 convert gptsanjapanese checkpoints from the original repository to pytorch model import argparse import json import os from collections import ordereddict import numpy as np import tensorflow as tf import torch def converttfgptsantoptargs parameterfile os path joinargs tfmodeldir parameters json params json loadsopenparameterfile read if not params raise valueerror fit seems that the json file at parameterfile is empty make sure you have a correct json file if not args output endswith pt args output args output pt newstate ordereddict with tf devicecpu 0 reader tf train loadcheckpointargs tfmodeldir shapes reader getvariabletoshapemap for keyname in shapes keys vnp reader gettensorkeyname astypenp float16 if keyname endswithadamm or keyname endswithadamv continue if keyname startswithpasts if keyname startswithpastsmlp player intkeyname9 elif keyname startswithpastsout player 8 name model sqout d weight player 2 enter to nn sequencial with tanh so 2 at a time state vnp transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname startswithmodelmoe player intkeyname9 split0 if keyname endswithswitchgatingkernel name model blocks d feedforward mlp router classifier weight player state vnp transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname endswithsoftmlpkernel name model blocks d feedforward softbypassmlp weight player state vnp transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname endswithwokernel or keyname endswithwikernel nlayer keyname9 7 for i in range16 name model blocks d feedforward mlp experts expertd s weight player i nlayer state vnpi transpose1 0 copy in meshtensorflow it is one array so it is divided newstatename torch tensorstate elif keyname startswithmodelmlp player intkeyname9 split0 if keyname endswithp1kernel name model blocks d feedforward mlp wi weight player state vnp transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname endswithp1bias name model blocks d feedforward mlp wi bias player state vnp copy same because it is one dimensional newstatename torch tensorstate elif keyname endswithp2kernel name model blocks d feedforward mlp wo weight player state vnp transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname endswithp2bias name model blocks d feedforward mlp wo bias player state vnp copy same because it is one dimensional newstatename torch tensorstate elif keyname startswithmodelln player intkeyname8 split0 if keyname endswithb name model blocks d feedforward norm bias player state vnp copy same because it is one dimensional newstatename torch tensorstate elif keyname endswithg name model blocks d feedforward norm weight player state vnp copy same because it is one dimensional newstatename torch tensorstate elif keyname startswithmodelatt player intkeyname9 split0 if keyname endswithqkvkernel state vnp copy compute same dimension as meshtensorflow using einsum stateq state 0 statek state 1 statev state 2 stateq stateq reshapestateq shape0 stateq shape1 stateq shape2 transpose1 0 copy meshtensorflow is a diagonal matrix statek statek reshapestatek shape0 statek shape1 statek shape2 transpose1 0 copy meshtensorflow is a diagonal matrix statev statev reshapestatev shape0 statev shape1 statev shape2 transpose1 0 copy meshtensorflow is a diagonal matrix name model blocks d selfattn selfattn qproj weight player newstatename torch tensorstateq name model blocks d selfattn selfattn kproj weight player newstatename torch tensorstatek name model blocks d selfattn selfattn vproj weight player newstatename torch tensorstatev elif keyname endswithokernel name model blocks d selfattn selfattn outproj weight player state vnp reshapevnp shape0 vnp shape1 vnp shape2 transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname startswithmodelan player intkeyname8 split0 if keyname endswithb name model blocks d selfattn norm bias player state vnp copy same because it is one dimensional newstatename torch tensorstate elif keyname endswithg name model blocks d selfattn norm weight player state vnp copy same because it is one dimensional newstatename torch tensorstate elif keyname startswithmodelwte or keyname startswithmodelwpe or keyname startswithmodelete nlayer wte embedtokens wpe positionembeddings ete extrapositionembeddings keyname3 name model s weight nlayer state vnp copy same in embedded newstatename torch tensorstate if keyname startswithmodelwte name lmhead weight state vnp copy same in embedded newstatename torch tensorstate elif keyname startswithmodelwob name finallogitsbias state vnp copy same in embedded state state reshape1 1 newstatename torch tensorstate elif keyname modeldensekernel name model lastproject weight state vnp transpose1 0 copy meshtensorflow is a diagonal matrix newstatename torch tensorstate elif keyname modeldense1bias name model lastproject bias state vnp copy same because it is one dimensional newstatename torch tensorstate torch savenewstate args output if name main parser argparse argumentparser descriptionmodel converter formatterclassargparse argumentdefaultshelpformatter parser addargumenttfmodeldir metavarpath typestr requiredtrue helpimport model parser addargumentoutput metavarpath typestr requiredtrue helpoutput model args parser parseargs converttfgptsantoptargs coding utf 8 2023 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 convert gptsanjapanese checkpoints from the original repository to pytorch model enter to nn sequencial with tanh so 2 at a time mesh tensorflow is a diagonal matrix mesh tensorflow is a diagonal matrix mesh tensorflow is a diagonal matrix in mesh tensorflow it is one array so it is divided mesh tensorflow is a diagonal matrix same because it is one dimensional mesh tensorflow is a diagonal matrix same because it is one dimensional same because it is one dimensional same because it is one dimensional compute same dimension as mesh tensorflow using einsum mesh tensorflow is a diagonal matrix mesh tensorflow is a diagonal matrix mesh tensorflow is a diagonal matrix mesh tensorflow is a diagonal matrix same because it is one dimensional same because it is one dimensional same in embedded same in embedded same in embedded mesh tensorflow is a diagonal matrix same because it is one dimensional
import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch


def convert_tf_gptsan_to_pt(args):
    parameter_file = os.path.join(args.tf_model_dir, "parameters.json")
    params = json.loads(open(parameter_file).read())
    if not params:
        raise ValueError(
            f"It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file."
        )
    if not args.output.endswith(".pt"):
        args.output = args.output + ".pt"
    new_state = OrderedDict()
    with tf.device("/CPU:0"):
        reader = tf.train.load_checkpoint(args.tf_model_dir)
        shapes = reader.get_variable_to_shape_map()
        for key_name in shapes.keys():
            vnp = reader.get_tensor(key_name).astype(np.float16)
            if key_name.endswith("/adam_m") or key_name.endswith("/adam_v"):
                continue
            if key_name.startswith("pasts/"):
                if key_name.startswith("pasts/mlp"):
                    player = int(key_name[9])
                elif key_name.startswith("pasts/out"):
                    player = 8
                name = "model.sqout.%d.weight" % (player * 2)  
                state = vnp.transpose([1, 0]).copy()  
                new_state[name] = torch.tensor(state)
            elif key_name.startswith("model/moe"):
                player = int(key_name[9:].split("/")[0])
                if key_name.endswith("/switch_gating/kernel"):
                    name = "model.blocks.%d.feed_forward.mlp.router.classifier.weight" % player
                    state = vnp.transpose([1, 0]).copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/softmlp/kernel"):
                    name = "model.blocks.%d.feed_forward.soft_bypass_mlp.weight" % player
                    state = vnp.transpose([1, 0]).copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/wo/kernel") or key_name.endswith("/wi/kernel"):
                    nlayer = key_name[-9:-7]
                    for i in range(16):
                        name = "model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight" % (player, i, nlayer)
                        state = (
                            vnp[i].transpose([1, 0]).copy()
                        )  
                        new_state[name] = torch.tensor(state)
            elif key_name.startswith("model/mlp"):
                player = int(key_name[9:].split("/")[0])
                if key_name.endswith("/p1/kernel"):
                    name = "model.blocks.%d.feed_forward.mlp.wi.weight" % player
                    state = vnp.transpose([1, 0]).copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/p1/bias"):
                    name = "model.blocks.%d.feed_forward.mlp.wi.bias" % player
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/p2/kernel"):
                    name = "model.blocks.%d.feed_forward.mlp.wo.weight" % player
                    state = vnp.transpose([1, 0]).copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/p2/bias"):
                    name = "model.blocks.%d.feed_forward.mlp.wo.bias" % player
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
            elif key_name.startswith("model/ln"):
                player = int(key_name[8:].split("/")[0])
                if key_name.endswith("/b"):
                    name = "model.blocks.%d.feed_forward.norm.bias" % player
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/g"):
                    name = "model.blocks.%d.feed_forward.norm.weight" % player
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
            elif key_name.startswith("model/att"):
                player = int(key_name[9:].split("/")[0])
                if key_name.endswith("/qkv/kernel"):
                    state = vnp.copy()  
                    state_q = state[:, 0, :, :]
                    state_k = state[:, 1, :, :]
                    state_v = state[:, 2, :, :]
                    state_q = (
                        state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]])
                        .transpose([1, 0])
                        .copy()
                    )  
                    state_k = (
                        state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]])
                        .transpose([1, 0])
                        .copy()
                    )  
                    state_v = (
                        state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]])
                        .transpose([1, 0])
                        .copy()
                    )  
                    name = "model.blocks.%d.self_attn.self_attn.q_proj.weight" % player
                    new_state[name] = torch.tensor(state_q)
                    name = "model.blocks.%d.self_attn.self_attn.k_proj.weight" % player
                    new_state[name] = torch.tensor(state_k)
                    name = "model.blocks.%d.self_attn.self_attn.v_proj.weight" % player
                    new_state[name] = torch.tensor(state_v)
                elif key_name.endswith("/o/kernel"):
                    name = "model.blocks.%d.self_attn.self_attn.out_proj.weight" % player
                    state = (
                        vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]).transpose([1, 0]).copy()
                    )  
                    new_state[name] = torch.tensor(state)
            elif key_name.startswith("model/an"):
                player = int(key_name[8:].split("/")[0])
                if key_name.endswith("/b"):
                    name = "model.blocks.%d.self_attn.norm.bias" % player
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
                elif key_name.endswith("/g"):
                    name = "model.blocks.%d.self_attn.norm.weight" % player
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
            elif (
                key_name.startswith("model/wte")
                or key_name.startswith("model/wpe")
                or key_name.startswith("model/ete")
            ):
                nlayer = {"wte": "embed_tokens", "wpe": "position_embeddings", "ete": "extra_position_embeddings"}[
                    key_name[-3:]
                ]
                name = "model.%s.weight" % nlayer
                state = vnp.copy()  
                new_state[name] = torch.tensor(state)
                if key_name.startswith("model/wte"):
                    name = "lm_head.weight"
                    state = vnp.copy()  
                    new_state[name] = torch.tensor(state)
            elif key_name.startswith("model/wob"):
                name = "final_logits_bias"
                state = vnp.copy()  
                state = state.reshape((1, -1))
                new_state[name] = torch.tensor(state)
            elif key_name == "model/dense/kernel":
                name = "model.last_project.weight"
                state = vnp.transpose([1, 0]).copy()  
                new_state[name] = torch.tensor(state)
            elif key_name == "model/dense_1/bias":
                name = "model.last_project.bias"
                state = vnp.copy()  
                new_state[name] = torch.tensor(state)
    torch.save(new_state, args.output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="model converter.", formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--tf_model_dir", metavar="PATH", type=str, required=True, help="import model")
    parser.add_argument("--output", metavar="PATH", type=str, required=True, help="output model")
    args = parser.parse_args()
    convert_tf_gptsan_to_pt(args)