File size: 9,890 Bytes
7900c16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import argparse
import collections
import torch


def convert_transformer_encoder_from_huggingface_to_tencentpretrain(input_model, output_model, layers_num):
    for i in range(layers_num):
        output_model['encoder.transformer.' + str(i) + '.self_attn.linear_layers.0.weight'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.q_proj.weight']
        output_model['encoder.transformer.' + str(i) + '.self_attn.linear_layers.0.bias'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.q_proj.bias']
        output_model['encoder.transformer.' + str(i) + '.self_attn.linear_layers.1.weight'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.k_proj.weight']
        output_model['encoder.transformer.' + str(i) + '.self_attn.linear_layers.1.bias'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.k_proj.bias']
        output_model['encoder.transformer.' + str(i) + '.self_attn.linear_layers.2.weight'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.v_proj.weight']
        output_model['encoder.transformer.' + str(i) + '.self_attn.linear_layers.2.bias'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.v_proj.bias']
        output_model['encoder.transformer.' + str(i) + '.self_attn.final_linear.weight'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.out_proj.weight']
        output_model['encoder.transformer.' + str(i) + '.self_attn.final_linear.bias'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn.out_proj.bias']

        output_model['encoder.transformer.' + str(i) + '.layer_norm_1.gamma'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn_layer_norm.weight']
        output_model['encoder.transformer.' + str(i) + '.layer_norm_1.beta'] = \
            input_model['model.encoder.layers.' + str(i) + '.self_attn_layer_norm.bias']
     
        output_model['encoder.transformer.' + str(i) + '.feed_forward.linear_1.weight'] = \
            input_model['model.encoder.layers.' + str(i) + '.fc1.weight']
        output_model['encoder.transformer.' + str(i) + '.feed_forward.linear_1.bias'] = \
            input_model['model.encoder.layers.' + str(i) + '.fc1.bias']  
        output_model['encoder.transformer.' + str(i) + '.feed_forward.linear_2.weight'] = \
            input_model['model.encoder.layers.' + str(i) + '.fc2.weight']
        output_model['encoder.transformer.' + str(i) + '.feed_forward.linear_2.bias'] = \
            input_model['model.encoder.layers.' + str(i) + '.fc2.bias']      

        output_model['encoder.transformer.' + str(i) + '.layer_norm_2.gamma'] = \
            input_model['model.encoder.layers.' + str(i) + '.final_layer_norm.weight']
        output_model['encoder.transformer.' + str(i) + '.layer_norm_2.beta'] = \
            input_model['model.encoder.layers.' + str(i) + '.final_layer_norm.bias']    

def convert_transformer_decoder_from_huggingface_to_tencentpretrain(input_model, output_model, layers_num):
    for i in range(layers_num):
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.linear_layers.0.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.q_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.linear_layers.0.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.q_proj.bias']            
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.linear_layers.1.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.k_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.linear_layers.1.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.k_proj.bias']    
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.linear_layers.2.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.v_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.linear_layers.2.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.v_proj.bias']    

        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.final_linear.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.out_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.self_attn.final_linear.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn.out_proj.bias']
        output_model['decoder.transformer_decoder.' + str(i) + '.layer_norm_1.gamma'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn_layer_norm.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.layer_norm_1.beta'] = \
            input_model['model.decoder.layers.' + str(i) + '.self_attn_layer_norm.bias']

        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.linear_layers.0.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.q_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.linear_layers.0.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.q_proj.bias']
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.linear_layers.1.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.k_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.linear_layers.1.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.k_proj.bias']
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.linear_layers.2.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.v_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.linear_layers.2.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.v_proj.bias']
 
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.final_linear.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.out_proj.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.context_attn.final_linear.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn.out_proj.bias']     
        output_model['decoder.transformer_decoder.' + str(i) + '.layer_norm_2.gamma'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn_layer_norm.weight']     
        output_model['decoder.transformer_decoder.' + str(i) + '.layer_norm_2.beta'] = \
            input_model['model.decoder.layers.' + str(i) + '.encoder_attn_layer_norm.bias']    

        output_model['decoder.transformer_decoder.' + str(i) + '.feed_forward.linear_1.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.fc1.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.feed_forward.linear_1.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.fc1.bias']
        output_model['decoder.transformer_decoder.' + str(i) + '.feed_forward.linear_2.weight'] = \
            input_model['model.decoder.layers.' + str(i) + '.fc2.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.feed_forward.linear_2.bias'] = \
            input_model['model.decoder.layers.' + str(i) + '.fc2.bias']

        output_model['decoder.transformer_decoder.' + str(i) + '.layer_norm_3.gamma'] = \
            input_model['model.decoder.layers.' + str(i) + '.final_layer_norm.weight']
        output_model['decoder.transformer_decoder.' + str(i) + '.layer_norm_3.beta'] = \
            input_model['model.decoder.layers.' + str(i) + '.final_layer_norm.bias']            

def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--input_model_path", type=str, default="models/input_model.pt",
                        help=".")
    parser.add_argument("--output_model_path", type=str, default="models/output_model.bin",
                        help=".")
    parser.add_argument("--layers_num", type=int, default=12, help=".")
    parser.add_argument("--decoder_layers_num", type=int, default=6, help=".")

    args = parser.parse_args()

    input_model = torch.load(args.input_model_path)

    output_model = collections.OrderedDict()

    for i in range(2):
        output_model["embedding.speech.conv.conv_layers." + str(i) + ".0.weight"] = \
            input_model["model.encoder.conv.conv_layers." + str(i) + ".weight"]
        output_model["embedding.speech.conv.conv_layers." + str(i) + ".0.bias"] = \
            input_model["model.encoder.conv.conv_layers." + str(i) + ".bias"]

    output_model['tgt_embedding.word.embedding.weight'] = input_model['model.decoder.embed_tokens.weight']

    convert_transformer_encoder_from_huggingface_to_tencentpretrain(input_model, output_model, args.layers_num)
    convert_transformer_decoder_from_huggingface_to_tencentpretrain(input_model, output_model, args.decoder_layers_num)
    output_model['encoder.layer_norm.gamma'] = input_model['model.encoder.layer_norm.weight']
    output_model['encoder.layer_norm.beta'] = input_model['model.encoder.layer_norm.bias']
    output_model['decoder.layer_norm.gamma'] = input_model['model.decoder.layer_norm.weight']
    output_model['decoder.layer_norm.beta'] = input_model['model.decoder.layer_norm.bias']

    output_model['target.lm.output_layer.weight'] = input_model["lm_head.weight"]

    torch.save(output_model, args.output_model_path)

if __name__ == "__main__":
    main()