Skip to content

Commit

Permalink
fix convert (#397)
Browse files Browse the repository at this point in the history
Co-authored-by: janinezhao <janinezhao@tencent.com>
  • Loading branch information
JINGZIjingzi and janinezhao committed Oct 27, 2023
1 parent 4f9d551 commit efd6e28
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 8 deletions.
8 changes: 3 additions & 5 deletions scripts/convert_pegasus_from_huggingface_to_uer.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,25 +7,23 @@
uer_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, uer_dir)

from scripts.convert_bart_from_huggingface_to_uer import\
from scripts.convert_bart_from_huggingface_to_uer import \
convert_encoder_decoder_transformer_from_huggingface_to_uer

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("--layers_num", type=int, default=6, help=".")
parser.add_argument("--decoder_layers_num", type=int, default=6, help=".")
parser.add_argument("--layers_num", type=int, default=12, help=".")
parser.add_argument("--decoder_layers_num", type=int, default=12, help=".")

args = parser.parse_args()

input_model = torch.load(args.input_model_path, map_location="cpu")

output_model = collections.OrderedDict()

output_model["embedding.sinusoidalpos.pe"] = input_model["model.encoder.embed_positions.weight"].unsqueeze(1)
output_model["tgt_embedding.sinusoidalpos.pe"] = input_model["model.decoder.embed_positions.weight"].unsqueeze(1)
output_model["embedding.word.embedding.weight"] = input_model["model.encoder.embed_tokens.weight"]
output_model["tgt_embedding.word.embedding.weight"] = input_model["model.decoder.embed_tokens.weight"]
output_model["target.lm.output_layer.weight"] = input_model["lm_head.weight"]
Expand Down
21 changes: 18 additions & 3 deletions scripts/convert_pegasus_from_uer_to_huggingface.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,12 @@
import argparse
import collections
import torch
import math

uer_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, uer_dir)

from scripts.convert_bart_from_uer_to_huggingface import\
from scripts.convert_bart_from_uer_to_huggingface import \
convert_encoder_decoder_transformer_from_uer_to_huggingface

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
Expand All @@ -17,6 +18,7 @@
help=".")
parser.add_argument("--layers_num", type=int, default=12, help=".")
parser.add_argument("--decoder_layers_num", type=int, default=12, help=".")
parser.add_argument("--max_seq_length", type=int, default=1024, help=".")

args = parser.parse_args()

Expand All @@ -25,8 +27,21 @@
output_model = collections.OrderedDict()

output_model["model.shared.weight"] = input_model["embedding.word.embedding.weight"]
output_model["model.encoder.embed_positions.weight"] = input_model["embedding.sinusoidalpos.pe"].squeeze(1)
output_model["model.decoder.embed_positions.weight"] = input_model["tgt_embedding.sinusoidalpos.pe"].squeeze(1)

emb_size = input_model["embedding.word.embedding.weight"].shape[1]
pe = torch.zeros(args.max_seq_length, emb_size)
position = torch.arange(0, args.max_seq_length).unsqueeze(1)
div_term = torch.exp(
(
torch.arange(0, emb_size, 2, dtype=torch.float)
*- (math.log(10000.0) / emb_size)
)
)
pe[:, 0::2] = torch.sin(position.float() * div_term)
pe[:, 1::2] = torch.cos(position.float() * div_term)

output_model["model.encoder.embed_positions.weight"] = pe
output_model["model.decoder.embed_positions.weight"] = pe
output_model["model.encoder.embed_tokens.weight"] = input_model["embedding.word.embedding.weight"]
output_model["model.decoder.embed_tokens.weight"] = input_model["tgt_embedding.word.embedding.weight"]
output_model["lm_head.weight"] = input_model["target.lm.output_layer.weight"]
Expand Down

0 comments on commit efd6e28

Please sign in to comment.