Spaces:
Runtime error
Runtime error
File size: 6,385 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 |
import sys
import os
import argparse
import collections
import torch
tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)
from tencentpretrain.utils.config import load_hyperparam
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--input_model_path", type=str, default="models/input_model.bin",
help=".")
parser.add_argument("--output_model_path", type=str, default="models/output_model.bin",
help=".")
parser.add_argument("--config_path", type=str,
help=".")
args = parser.parse_args()
args = load_hyperparam(args)
input_model = torch.load(args.input_model_path)
if "word" in args.embedding:
input_model["embedding.word.embedding.weight"] = input_model["embedding.word_embedding.weight"]
input_model.pop("embedding.word_embedding.weight")
if "pos" in args.embedding:
input_model["embedding.pos.embedding.weight"] = input_model["embedding.position_embedding.weight"]
input_model.pop("embedding.position_embedding.weight")
if "seg" in args.embedding:
input_model["embedding.seg.embedding.weight"] = input_model["embedding.segment_embedding.weight"]
input_model.pop("embedding.segment_embedding.weight")
if "sinusoidalpos" in args.embedding:
input_model["embedding.sinusoidalpos.pe"] = input_model["embedding.pe"]
input_model.pop("embedding.pe")
if hasattr(args, "decoder") and args.decoder is not None:
for n in list(input_model.keys()): # target.decoder -> decoder
if n.split('.')[1] == "decoder":
input_model[".".join(n.split('.')[1:])] = input_model[n]
input_model.pop(n)
if n.split('.')[1] == "embedding":
input_model[".".join(["tgt_embedding"] + n.split('.')[2:])] = input_model[n]
input_model.pop(n)
if "word" in args.embedding:
input_model["tgt_embedding.word.embedding.weight"] = input_model["tgt_embedding.word_embedding.weight"]
input_model.pop("tgt_embedding.word_embedding.weight")
if "pos" in args.embedding:
input_model["tgt_embedding.pos.embedding.weight"] = input_model["tgt_embedding.position_embedding.weight"]
input_model.pop("tgt_embedding.position_embedding.weight")
if "seg" in args.embedding:
input_model["tgt_embedding.seg.embedding.weight"] = input_model["tgt_embedding.segment_embedding.weight"]
input_model.pop("tgt_embedding.segment_embedding.weight")
if "sinusoidalpos" in args.embedding:
input_model["tgt_embedding.sinusoidalpos.pe"] = input_model["tgt_embedding.pe"]
input_model.pop("tgt_embedding.pe")
if "mlm" in args.target:
try:
input_model["target.mlm.linear_1.weight"] = input_model["target.mlm_linear_1.weight"]
input_model.pop("target.mlm_linear_1.weight")
input_model["target.mlm.linear_1.bias"] = input_model["target.mlm_linear_1.bias"]
input_model.pop("target.mlm_linear_1.bias")
input_model["target.mlm.layer_norm.gamma"] = input_model["target.layer_norm.gamma"]
input_model.pop("target.layer_norm.gamma")
input_model["target.mlm.layer_norm.beta"] = input_model["target.layer_norm.beta"]
input_model.pop("target.layer_norm.beta")
input_model["target.mlm.linear_2.weight"] = input_model["target.mlm_linear_2.weight"]
input_model.pop("target.mlm_linear_2.weight")
input_model["target.mlm.linear_2.bias"] = input_model["target.mlm_linear_2.bias"]
input_model.pop("target.mlm_linear_2.bias")
except:
pass
if "sp" in args.target:
try:
input_model["target.sp.linear_1.weight"] = input_model["target.sp_linear_1.weight"]
input_model.pop("target.sp_linear_1.weight")
input_model["target.sp.linear_1.bias"] = input_model["target.sp_linear_1.bias"]
input_model.pop("target.sp_linear_1.bias")
input_model["target.sp.linear_2.weight"] = input_model["target.sp_linear_2.weight"]
input_model.pop("target.sp_linear_2.weight")
input_model["target.sp.linear_2.bias"] = input_model["target.sp_linear_2.bias"]
input_model.pop("target.sp_linear_2.bias")
except:
pass
try:
input_model["target.sp.linear_1.weight"] = input_model["target.nsp_linear_1.weight"]
input_model.pop("target.nsp_linear_1.weight")
input_model["target.sp.linear_1.bias"] = input_model["target.nsp_linear_1.bias"]
input_model.pop("target.nsp_linear_1.bias")
input_model["target.sp.linear_2.weight"] = input_model["target.nsp_linear_2.weight"]
input_model.pop("target.nsp_linear_2.weight")
input_model["target.sp.linear_2.bias"] = input_model["target.nsp_linear_2.bias"]
input_model.pop("target.nsp_linear_2.bias")
except:
pass
try:
input_model["target.sp.linear_1.weight"] = input_model["target.sop_linear_1.weight"]
input_model.pop("target.sop_linear_1.weight")
input_model["target.sp.linear_1.bias"] = input_model["target.sop_linear_1.bias"]
input_model.pop("target.sop_linear_1.bias")
input_model["target.sp.linear_2.weight"] = input_model["target.sop_linear_2.weight"]
input_model.pop("target.sop_linear_2.weight")
input_model["target.sp.linear_2.bias"] = input_model["target.sop_linear_2.bias"]
input_model.pop("target.sop_linear_2.bias")
except:
pass
if "lm" in args.target:
try:
input_model["target.lm.output_layer.weight"] = input_model["target.output_layer.weight"]
input_model.pop("target.output_layer.weight")
if args.has_lmtarget_bias:
input_model["target.lm.output_layer.bias"] = input_model["target.output_layer.bias"]
input_model.pop("target.output_layer.bias")
except:
pass
torch.save(input_model, args.output_model_path)
if __name__ == "__main__":
main()
|