added llama_tencentpretrain_pytorch

models/vit/base-32-224_config.json

+{
+  "emb_size": 768,
+  "feedforward_size": 3072,
+  "hidden_size": 768,
+  "hidden_act": "gelu",
+  "heads_num": 12,
+  "layers_num": 12,
+  "dropout": 0.1,
+  "max_seq_length": 50,
+  "data_processor": "vit",
+  "embedding": ["patch", "pos"],
+  "remove_embedding_layernorm": true,
+  "encoder": "transformer",
+  "mask": "fully_visible",
+  "layernorm_positioning": "pre",
+  "target": ["cls"],
+  "image_height": 224,
+  "image_width": 224,
+  "patch_size": 32
+}
\ No newline at end of file

models/vit/huge-14-224_config.json

+{
+  "emb_size": 1280,
+  "feedforward_size": 5120,
+  "hidden_size": 1280,
+  "hidden_act": "gelu",
+  "heads_num": 16,
+  "layers_num": 32,
+  "dropout": 0.1,
+  "max_seq_length": 257,
+  "data_processor": "vit",
+  "embedding": ["patch", "pos"],
+  "remove_embedding_layernorm": true,
+  "encoder": "transformer",
+  "mask": "fully_visible",
+  "layernorm_positioning": "pre",
+  "target": ["cls"],
+  "image_height": 224,
+  "image_width": 224,
+  "patch_size": 14
+}
\ No newline at end of file

models/vit/large-16-224_config.json

+{
+  "emb_size": 1024,
+  "feedforward_size": 4096,
+  "hidden_size": 1024,
+  "hidden_act": "gelu",
+  "heads_num": 16,
+  "layers_num": 24,
+  "dropout": 0.1,
+  "max_seq_length": 197,
+  "data_processor": "vit",
+  "embedding": ["patch", "pos"],
+  "remove_embedding_layernorm": true,
+  "encoder": "transformer",
+  "mask": "fully_visible",
+  "layernorm_positioning": "pre",
+  "target": ["cls"],
+  "image_height": 224,
+  "image_width": 224,
+  "patch_size": 16
+}
\ No newline at end of file

models/vit/large-32-224_config.json

+{
+  "emb_size": 1024,
+  "feedforward_size": 4096,
+  "hidden_size": 1024,
+  "hidden_act": "gelu",
+  "heads_num": 16,
+  "layers_num": 24,
+  "dropout": 0.1,
+  "max_seq_length": 50,
+  "data_processor": "vit",
+  "embedding": ["patch", "pos"],
+  "remove_embedding_layernorm": true,
+  "encoder": "transformer",
+  "mask": "fully_visible",
+  "layernorm_positioning": "pre",
+  "target": ["cls"],
+  "image_height": 224,
+  "image_width": 224,
+  "patch_size": 32
+}
\ No newline at end of file

models/xlm-roberta/base_config.json

+{
+  "emb_size": 768,
+  "feedforward_size": 3072,
+  "hidden_size": 768,
+  "hidden_act": "gelu",
+  "heads_num": 12,
+  "layers_num": 12,
+  "max_seq_length": 514,
+  "dropout": 0.1,
+  "data_processor": "mlm",
+  "embedding": ["word", "pos", "seg"],
+  "encoder": "transformer",
+  "mask": "fully_visible",
+  "target": ["mlm"]
+}
\ No newline at end of file

models/xlm-roberta/large_config.json

+{
+  "emb_size": 1024,
+  "feedforward_size": 4096,
+  "hidden_size": 1024,
+  "hidden_act": "gelu",
+  "heads_num": 16,
+  "layers_num": 24,
+  "max_seq_length": 514,
+  "dropout": 0.1,
+  "data_processor": "mlm",
+  "embedding": ["word", "pos", "seg"],
+  "encoder": "transformer",
+  "mask": "fully_visible",
+  "target": ["mlm"]
+}
\ No newline at end of file

models/xlmroberta_special_tokens_map.json

+{
+  "pad_token": "<pad>",
+  "unk_token": "<unk>",
+  "cls_token": "<s>",
+  "sep_token": "</s>",
+  "mask_token": "<mask>"
+}

preprocess.py

+import argparse
+import six
+from packaging import version
+from tencentpretrain.utils import *
+from tencentpretrain.opts import *
+
+
+assert version.parse(six.__version__) >= version.parse("1.12.0")
+
+
+def main():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    # Path options.
+    parser.add_argument("--corpus_path", type=str, default="data/alpaca_gpt4_data_zh.json",
+                        help="Path of the corpus for pretraining.")
+    parser.add_argument("--dataset_path", type=str, default="data/dataset.pt",
+                        help="Path of the preprocessed dataset.")
+    parser.add_argument("--json_format_corpus", action="store_true",
+                        help="Load from a json file corpus, like {\"text\": \"text\"}.")
+
+    # Preprocess options.
+    tokenizer_opts(parser)
+    tgt_tokenizer_opts(parser)
+    parser.add_argument("--processes_num", type=int, default=16,
+                        help="Split the whole dataset into `processes_num` parts, "
+                             "and process them with `processes_num` processes.")
+    parser.add_argument("--data_processor",
+                        choices=["bert", "lm", "mlm", "bilm", "albert", "mt", "t5", "cls", "prefixlm",
+                                 "gsg", "bart", "cls_mlm", "vit", "vilt", "clip", "s2t", "beit", "dalle", "alpaca"], default="alpaca",
+                        help="The data processor of the pretraining model.")
+    parser.add_argument("--docs_buffer_size", type=int, default=100000,
+                        help="The buffer size of documents in memory, specific to targets that require negative sampling.")
+    parser.add_argument("--seq_length", type=int, default=1024, help="Sequence length of instances.")
+    parser.add_argument("--tgt_seq_length", type=int, default=128, help="Target sequence length of instances.")
+    parser.add_argument("--dup_factor", type=int, default=5,
+                        help="Duplicate instances multiple times.")
+    parser.add_argument("--short_seq_prob", type=float, default=0.1,
+                        help="Probability of truncating sequence."
+                             "The larger value, the higher probability of using short (truncated) sequence.")
+    parser.add_argument("--full_sentences", action="store_true", help="Full sentences.")
+    parser.add_argument("--seed", type=int, default=7, help="Random seed.")
+
+    # Masking options.
+    parser.add_argument("--dynamic_masking", action="store_true", help="Dynamic masking.")
+    parser.add_argument("--whole_word_masking", action="store_true", help="Whole word masking.")
+    parser.add_argument("--span_masking", action="store_true", help="Span masking.")
+    parser.add_argument("--span_geo_prob", type=float, default=0.2,
+                        help="Hyperparameter of geometric distribution for span masking.")
+    parser.add_argument("--span_max_length", type=int, default=10,
+                        help="Max length for span masking.")
+
+    # Sentence selection strategy options.
+    parser.add_argument("--sentence_selection_strategy", choices=["lead", "random"], default="lead",
+                        help="Sentence selection strategy for gap-sentences generation task.")
+
+    args = parser.parse_args()
+
+    # Dynamic masking.
+    if args.dynamic_masking:
+        args.dup_factor = 1
+
+    # Build tokenizer.
+    tokenizer = str2tokenizer[args.tokenizer](args)
+    if args.data_processor == "mt":
+        args.tgt_tokenizer = str2tokenizer[args.tgt_tokenizer](args, False)
+
+    # Build and save dataset.
+    dataset = str2dataset[args.data_processor](args, tokenizer.vocab, tokenizer)
+    dataset.build_and_save(args.processes_num)
+
+
+if __name__ == "__main__":
+    main()

pretrain.py

+import argparse
+import torch
+import tencentpretrain.trainer as trainer
+from tencentpretrain.utils.config import load_hyperparam
+from tencentpretrain.opts import *
+
+
+def main():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    # Path options.
+    parser.add_argument("--dataset_path", type=str, default="data/dataset.pt",
+                        help="Path of the preprocessed dataset.")
+    parser.add_argument("--pretrained_model_path", type=str, default=None,
+                        help="Path of the pretrained model.")
+    parser.add_argument("--output_model_path", type=str, required=True,
+                        help="Path of the output model.")
+    parser.add_argument("--config_path", type=str, default="models/llama/13b_config.json",
+                        help="Config file of model hyper-parameters.")
+
+    # Training and saving options.
+    parser.add_argument("--total_steps", type=int, default=100000,
+                        help="Total training steps.")
+    parser.add_argument("--save_checkpoint_steps", type=int, default=10000,
+                        help="Specific steps to save model checkpoint.")
+    parser.add_argument("--report_steps", type=int, default=5,
+                        help="Specific steps to print prompt.")
+    parser.add_argument("--accumulation_steps", type=int, default=1,
+                        help="Specific steps to accumulate gradient.")
+    parser.add_argument("--batch_size", type=int, default=32,
+                        help="Training batch size. The actual batch_size is [batch_size x world_size x accumulation_steps].")
+    parser.add_argument("--instances_buffer_size", type=int, default=25600,
+                        help="The buffer size of instances in memory.")
+    parser.add_argument("--labels_num", type=int, required=False,
+                        help="Number of prediction labels.")
+    parser.add_argument("--dropout", type=float, default=0.1, help="Dropout value.")
+    parser.add_argument("--seed", type=int, default=7, help="Random seed.")
+
+    # Preprocess options.
+    tokenizer_opts(parser)
+    tgt_tokenizer_opts(parser)
+
+    # Model options.
+    model_opts(parser)
+    parser.add_argument("--data_processor",
+                        choices=["bert", "lm", "mlm", "bilm", "albert", "mt", "t5", "cls",
+                                 "prefixlm", "gsg", "bart", "cls_mlm", "vit", "vilt", "clip", "s2t", "beit", "dalle", "alpaca"], default="alpaca",
+                        help="The data processor of the pretraining model.")
+    parser.add_argument("--deep_init", action="store_true",
+                        help="Scaling initialization of projection layers by a "
+                             "factor of 1/sqrt(2N). Necessary to large models.")
+
+    # Masking options.
+    parser.add_argument("--whole_word_masking", action="store_true", help="Whole word masking.")
+    parser.add_argument("--span_masking", action="store_true", help="Span masking.")
+    parser.add_argument("--span_geo_prob", type=float, default=0.2,
+                        help="Hyperparameter of geometric distribution for span masking.")
+    parser.add_argument("--span_max_length", type=int, default=10,
+                        help="Max length for span masking.")
+
+    # Optimizer options.
+    optimization_opts(parser)
+
+    # GPU options.
+    parser.add_argument("--world_size", type=int, default=1, help="Total number of processes (GPUs) for training.")
+    parser.add_argument("--gpu_ranks", default=[], nargs='+', type=int, help="List of ranks of each process."
+                        " Each process has a unique integer rank whose value is in the interval [0, world_size), and runs in a single GPU.")
+    parser.add_argument("--master_ip", default="tcp://localhost:12345", type=str, help="IP-Port of master for training.")
+    parser.add_argument("--backend", choices=["nccl", "gloo"], default="nccl", type=str, help="Distributed backend.")
+
+    # Deepspeed options.
+    deepspeed_opts(parser)
+
+    # lora options.
+    lora_opts(parser)
+
+    # Log options.
+    log_opts(parser)
+
+    args = parser.parse_args()
+
+    # construct lora dict parameters.
+    if args.use_lora:
+        args.lora_params = {
+            "lora_r": args.lora_r,
+            "lora_alpha": args.lora_alpha,
+            "lora_dropout": args.lora_dropout
+        }
+    else:
+        args.lora_params = None
+
+    if "cls" in args.target:
+        assert args.labels_num is not None, "Cls target needs the denotation of the number of labels."
+
+    # Load hyper-parameters from config file.
+    if args.config_path:
+        args = load_hyperparam(args)
+
+    ranks_num = len(args.gpu_ranks)
+
+    if args.deepspeed:
+        if args.world_size > 1:
+            args.dist_train = True
+        else:
+            args.dist_train = False
+    else:
+        if args.world_size > 1:
+            # Multiprocessing distributed mode.
+            assert torch.cuda.is_available(), "No available GPUs."
+            assert ranks_num <= args.world_size, "Started processes exceed `world_size` upper limit."
+            assert ranks_num <= torch.cuda.device_count(), "Started processes exceeds the available GPUs."
+            args.dist_train = True
+            args.ranks_num = ranks_num
+            print("Using distributed mode for training.")
+        elif args.world_size == 1 and ranks_num == 1:
+            # Single GPU mode.
+            assert torch.cuda.is_available(), "No available GPUs."
+            args.gpu_id = args.gpu_ranks[0]
+            assert args.gpu_id < torch.cuda.device_count(), "Invalid specified GPU device."
+            args.dist_train = False
+            args.single_gpu = True
+            print("Using GPU %d for training." % args.gpu_id)
+        else:
+            # CPU mode.
+            assert ranks_num == 0, "GPUs are specified, please check the arguments."
+            args.dist_train = False
+            args.single_gpu = False
+            print("Using CPU mode for training.")
+
+    trainer.train_and_validate(args)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

+argparse
+six>=1.12.0
+packaging
+numpy
+regex

scripts/average_models.py

+import argparse
+import sys
+import os
+import torch
+
+tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(tencentpretrain_dir)
+
+
+def average(model_list_path):
+    for i, model_path in enumerate(model_list_path):
+        model = torch.load(model_path)
+        if i == 0:
+            avg_model = model
+        else:
+            for k, _ in avg_model.items():
+                avg_model[k].mul_(i).add_(model[k]).div_(i+1)
+
+    return avg_model
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument("--model_list_path", nargs="+", required=True,
+                        help="Path of the input model list.")
+    parser.add_argument("--output_model_path", required=True,
+                        help="Path of the output model.")
+    args = parser.parse_args()
+
+    avg_model = average(args.model_list_path)
+    torch.save(avg_model, args.output_model_path)

scripts/beginning.txt

+上海有什么玩的？
\ No newline at end of file

scripts/build_vocab.py

+"""
+Build vocabulary with given tokenizer
+"""
+import sys
+import os
+import argparse
+
+tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(tencentpretrain_dir)
+
+from tencentpretrain.utils import *
+from tencentpretrain.utils.vocab import Vocab
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    parser.add_argument("--corpus_path", required=True)
+    parser.add_argument("--delimiter", choices=["char", "space"], required=True,
+                        help="Tokenizing the corpus in char-level or by the provided spaces.")
+    parser.add_argument("--output_path", required=True,
+                        help="The output path to save the vocabulary.")
+    parser.add_argument("--workers_num", type=int, default=1,
+                        help="The number of processes to build vocabulary.")
+    parser.add_argument("--min_count", type=int, default=1,
+                        help="The minimum count of words retained in the vocabulary.")
+    
+    args = parser.parse_args()
+
+    # Build tokenizer only for char and space.
+    args.vocab_path, args.spm_model_path = "./models/reserved_vocab.txt", None
+    tokenizer = str2tokenizer[args.delimiter](args)
+
+    # Build and save vocabulary using CharTokenizer or SpaceTokenizer.
+    vocab = Vocab()
+    vocab.build(args.corpus_path, tokenizer, args.workers_num, args.min_count)
+    vocab.save(args.output_path)

scripts/cloze_test.py

+"""
+  This script provides an exmaple to wrap TencentPretrain for cloze test.
+  One character in a line is masked.
+  We should use the target that contains MLM.
+"""
+import sys
+import os
+import torch
+import argparse
+import random
+
+tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(tencentpretrain_dir)
+
+from tencentpretrain.embeddings import *
+from tencentpretrain.encoders import *
+from tencentpretrain.targets import *
+from tencentpretrain.utils.constants import *
+from tencentpretrain.utils import *
+from tencentpretrain.utils.config import load_hyperparam
+from tencentpretrain.model_loader import load_model
+from tencentpretrain.opts import infer_opts, tokenizer_opts
+
+
+def mask_token(tokens, seq_length, tokenizer):
+    """
+    Mask a random token for prediction.
+    """
+    start = 1
+    end = len(tokens) if len(tokens) < seq_length else seq_length
+    mask_pos = random.randint(start, end-1)
+    token = tokens[mask_pos]
+    tokens[mask_pos] = tokenizer.convert_tokens_to_ids([MASK_TOKEN])[0]
+    return (tokens, mask_pos, token)
+
+
+def batch_loader(batch_size, src, seg, mask_pos, label):
+    instances_num = src.size(0)
+    for i in range(instances_num // batch_size):
+        src_batch = src[i * batch_size : (i + 1) * batch_size, :]
+        seg_batch = seg[i * batch_size : (i + 1) * batch_size, :]
+        mask_pos_batch = mask_pos[i * batch_size : (i + 1) * batch_size]
+        label_batch = label[i * batch_size : (i + 1) * batch_size]
+        yield src_batch, seg_batch, mask_pos_batch, label_batch
+
+    if instances_num > instances_num // batch_size * batch_size:
+        src_batch = src[instances_num // batch_size * batch_size :, :]
+        seg_batch = seg[instances_num // batch_size * batch_size :, :]
+        mask_pos_batch = mask_pos[instances_num // batch_size * batch_size :]
+        label_batch = label[instances_num // batch_size * batch_size :]
+        yield src_batch, seg_batch, mask_pos_batch, label_batch
+
+
+def read_dataset(args, path):
+    dataset = []
+    PAD_ID = args.tokenizer.vocab.get(PAD_TOKEN)
+    with open(path, mode="r", encoding="utf-8") as f:
+        for line in f:
+            src = args.tokenizer.convert_tokens_to_ids(args.tokenizer.tokenize(line.strip()))
+            if len(src) == 0:
+                continue
+
+            src = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN]) + src + args.tokenizer.convert_tokens_to_ids([SEP_TOKEN])
+            src, mask_pos, label = mask_token(src, args.seq_length, args.tokenizer)
+
+            seg = [1] * len(src)
+            if len(src) > args.seq_length:
+                src = src[:args.seq_length]
+                seg = seg[:args.seq_length]
+            while len(src) < args.seq_length:
+                src.append(PAD_ID)
+                seg.append(PAD_ID)
+            
+            dataset.append((src, seg, mask_pos, label))
+    return dataset
+
+
+class ClozeTest(torch.nn.Module):
+    def __init__(self, args):
+        super(ClozeTest, self).__init__()
+        # self.embedding = str2embedding[args.embedding](args, len(args.tokenizer.vocab))
+        self.embedding = Embedding(args)
+        for embedding_name in args.embedding:
+            tmp_emb = str2embedding[embedding_name](args, len(args.tokenizer.vocab))
+            self.embedding.update(tmp_emb, embedding_name)
+        self.encoder = str2encoder[args.encoder](args)
+        self.target = MlmTarget(args, len(args.tokenizer.vocab))
+        self.act = str2act[args.hidden_act]
+
+    def forward(self, src, seg):
+        emb = self.embedding(src, seg)
+        output = self.encoder(emb, seg)
+        output = self.act(self.target.mlm_linear_1(output))
+        output = self.target.layer_norm(output)
+        output = self.target.mlm_linear_2(output)
+        prob = torch.nn.Softmax(dim=-1)(output)
+        return prob
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    infer_opts(parser)
+
+    tokenizer_opts(parser)
+
+    parser.add_argument("--topn", type=int, default=10,
+                        help="Print top n nearest neighbours.")
+    
+    args = parser.parse_args()
+
+    args.target = "mlm"
+
+    # Load the hyperparameters from the config file.
+    args = load_hyperparam(args)
+
+    args.tokenizer = str2tokenizer[args.tokenizer](args)
+
+    # Build cloze test model.
+    model = ClozeTest(args)
+    model = load_model(model, args.load_model_path)
+
+    # For simplicity, we use DataParallel wrapper to use multiple GPUs.
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    if torch.cuda.device_count() > 1:
+        print("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
+        model = torch.nn.DataParallel(model)
+    model.eval()
+
+    dataset = read_dataset(args, args.test_path)
+
+    src = torch.LongTensor([sample[0] for sample in dataset])
+    seg = torch.LongTensor([sample[1] for sample in dataset])
+
+    mask_pos = [sample[2] for sample in dataset]
+    label = [sample[3] for sample in dataset]
+    
+    f_pred = open(args.prediction_path, mode="w", encoding="utf-8")
+               
+    for i, (src_batch, seg_batch, mask_pos_batch, label_batch) in \
+        enumerate(batch_loader(args.batch_size, src, seg, mask_pos, label)):
+        src_batch = src_batch.to(device)
+        seg_batch = seg_batch.to(device)
+        prob = model(src_batch, seg_batch)
+
+        for j, p in enumerate(mask_pos_batch):
+            topn_ids = (-prob[j][p]).argsort()[:args.topn]
+
+            sentence = "".join([args.tokenizer.convert_ids_to_tokens([token_id.item()])[0] for token_id in src_batch[j] if token_id != 0])
+            pred_tokens = " ".join(args.tokenizer.convert_ids_to_tokens([token_id.item()])[0] for token_id in topn_ids)
+
+            label_token = args.tokenizer.convert_ids_to_tokens([label_batch[j]])[0]
+
+            f_pred.write(sentence + '\n')
+            f_pred.write("Predicted answer: " + pred_tokens + '\n')
+            f_pred.write("Correct answer: " + label_token + '\n')
+            f_pred.write("\n")
+    
+    f_pred.close()

scripts/convert_albert_from_huggingface_to_tencentpretrain.py

+import argparse
+import collections
+import torch
+
+
+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=".")
+
+args = parser.parse_args()
+
+input_model = torch.load(args.input_model_path, map_location="cpu")
+
+output_model = collections.OrderedDict()
+
+output_model["embedding.word.embedding.weight"] = \
+    input_model["albert.embeddings.word_embeddings.weight"]
+output_model["embedding.pos.embedding.weight"] = \
+    input_model["albert.embeddings.position_embeddings.weight"]
+output_model["embedding.seg.embedding.weight"] = \
+    torch.cat((torch.Tensor([[0]*input_model["albert.embeddings.token_type_embeddings.weight"].size()[1]]),
+               input_model["albert.embeddings.token_type_embeddings.weight"]), dim=0)
+output_model["embedding.layer_norm.gamma"] = \
+    input_model["albert.embeddings.LayerNorm.weight"]
+output_model["embedding.layer_norm.beta"] = \
+    input_model["albert.embeddings.LayerNorm.bias"]
+
+output_model["encoder.linear.weight"] = \
+    input_model["albert.encoder.embedding_hidden_mapping_in.weight"]
+output_model["encoder.linear.bias"] = \
+    input_model["albert.encoder.embedding_hidden_mapping_in.bias"]
+output_model["encoder.transformer.layer_norm_2.gamma"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.weight"]
+output_model["encoder.transformer.layer_norm_2.beta"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm.bias"]
+output_model["encoder.transformer.self_attn.linear_layers.0.weight"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.weight"]
+output_model["encoder.transformer.self_attn.linear_layers.0.bias"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.query.bias"]
+output_model["encoder.transformer.self_attn.linear_layers.1.weight"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.weight"]
+output_model["encoder.transformer.self_attn.linear_layers.1.bias"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.key.bias"]
+output_model["encoder.transformer.self_attn.linear_layers.2.weight"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.weight"]
+output_model["encoder.transformer.self_attn.linear_layers.2.bias"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.value.bias"]
+output_model["encoder.transformer.self_attn.final_linear.weight"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.weight"]
+output_model["encoder.transformer.self_attn.final_linear.bias"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.dense.bias"]
+output_model["encoder.transformer.layer_norm_1.gamma"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.weight"]
+output_model["encoder.transformer.layer_norm_1.beta"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm.bias"]
+output_model["encoder.transformer.feed_forward.linear_1.weight"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.weight"]
+output_model["encoder.transformer.feed_forward.linear_1.bias"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.ffn.bias"]
+output_model["encoder.transformer.feed_forward.linear_2.weight"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.weight"]
+output_model["encoder.transformer.feed_forward.linear_2.bias"] = \
+    input_model["albert.encoder.albert_layer_groups.0.albert_layers.0.ffn_output.bias"]
+
+output_model["target.sp.linear_1.weight"] = input_model["albert.pooler.weight"]
+output_model["target.sp.linear_1.bias"] = input_model["albert.pooler.bias"]
+output_model["target.sp.linear_2.weight"] = input_model["sop_classifier.classifier.weight"]
+output_model["target.sp.linear_2.bias"] = input_model["sop_classifier.classifier.bias"]
+output_model["target.mlm.linear_1.weight"] = input_model["predictions.dense.weight"]
+output_model["target.mlm.linear_1.bias"] = input_model["predictions.dense.bias"]
+output_model["target.mlm.linear_2.weight"] = input_model["predictions.decoder.weight"]
+output_model["target.mlm.linear_2.bias"] = input_model["predictions.bias"]
+output_model["target.mlm.layer_norm.gamma"] = input_model["predictions.LayerNorm.weight"]
+output_model["target.mlm.layer_norm.beta"] = input_model["predictions.LayerNorm.bias"]
+
+torch.save(output_model, args.output_model_path)

scripts/convert_albert_from_original_tf_to_tencentpretrain.py

+import argparse
+import collections
+import tensorflow as tf
+import torch
+from tensorflow.python import pywrap_tensorflow
+
+
+tensors_to_transopse = (
+    "kernel",
+)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    # Path options.
+    parser.add_argument("--input_model_path", type=str, default="models/input_model.ckpt",
+                        help="Path of the input model.")
+    parser.add_argument("--output_model_path", type=str, default="models/output_model.bin",
+                        help="Path of the output model.")
+
+    args = parser.parse_args()
+
+    reader = pywrap_tensorflow.NewCheckpointReader(args.input_model_path)
+    var_to_shape_map = reader.get_variable_to_shape_map()
+
+    input_model = collections.OrderedDict()
+
+    for key in var_to_shape_map:
+        torch_tensor = reader.get_tensor(key)
+        
+        if any([x in key for x in tensors_to_transopse]):
+            torch_tensor = torch_tensor.T
+        if key == "bert/embeddings/token_type_embeddings":
+            col_dim = torch_tensor.shape[1]
+            sess = tf.Session()
+            zeros_var = tf.Variable(tf.zeros([1, col_dim], dtype=tf.float32), name="zeros_var")
+            sess.run(zeros_var.initializer)
+            torch_tensor = sess.run(tf.concat([sess.run(zeros_var), torch_tensor], 0))
+        input_model[key] = torch.Tensor(torch_tensor)
+
+    output_model = collections.OrderedDict()
+
+    output_model["embedding.word.embedding.weight"] = input_model["bert/embeddings/word_embeddings"]
+    output_model["encoder.linear.weight"] = input_model["bert/encoder/embedding_hidden_mapping_in/kernel"]
+    output_model["encoder.linear.bias"] = input_model["bert/encoder/embedding_hidden_mapping_in/bias"]
+    output_model["embedding.pos.embedding.weight"] = input_model["bert/embeddings/position_embeddings"][:512]
+    output_model["embedding.seg.embedding.weight"] = input_model["bert/embeddings/token_type_embeddings"]
+    output_model["embedding.layer_norm.gamma"] = input_model["bert/embeddings/LayerNorm/gamma"]
+    output_model["embedding.layer_norm.beta"] = input_model["bert/embeddings/LayerNorm/beta"]
+
+    output_model["encoder.transformer.self_attn.linear_layers.0.weight"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/self/query/kernel"]
+    output_model["encoder.transformer.self_attn.linear_layers.0.bias"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/self/query/bias"]
+    output_model["encoder.transformer.self_attn.linear_layers.1.weight"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/self/key/kernel"]
+    output_model["encoder.transformer.self_attn.linear_layers.1.bias"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/self/key/bias"]
+    output_model["encoder.transformer.self_attn.linear_layers.2.weight"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/self/value/kernel"]
+    output_model["encoder.transformer.self_attn.linear_layers.2.bias"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/self/value/bias"]
+    output_model["encoder.transformer.self_attn.final_linear.weight"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/output/dense/kernel"]
+    output_model["encoder.transformer.self_attn.final_linear.bias"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/attention_1/output/dense/bias"]
+    output_model["encoder.transformer.layer_norm_1.gamma"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/LayerNorm/gamma"]
+    output_model["encoder.transformer.layer_norm_1.beta"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/LayerNorm/beta"]
+    output_model["encoder.transformer.feed_forward.linear_1.weight"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/dense/kernel"]
+    output_model["encoder.transformer.feed_forward.linear_1.bias"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/dense/bias"]
+    output_model["encoder.transformer.feed_forward.linear_2.weight"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/output/dense/kernel"]
+    output_model["encoder.transformer.feed_forward.linear_2.bias"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/ffn_1/intermediate/output/dense/bias"]
+    output_model["encoder.transformer.layer_norm_2.gamma"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/LayerNorm_1/gamma"]
+    output_model["encoder.transformer.layer_norm_2.beta"] = \
+        input_model["bert/encoder/transformer/group_0/inner_group_0/LayerNorm_1/beta"]
+
+    output_model["target.sp.linear_1.weight"] = input_model["bert/pooler/dense/kernel"]
+    output_model["target.sp.linear_1.bias"] = input_model["bert/pooler/dense/bias"]
+    output_model["target.sp.linear_2.weight"] = input_model["cls/seq_relationship/output_weights"]
+    output_model["target.sp.linear_2.bias"] = input_model["cls/seq_relationship/output_bias"]
+    output_model["target.mlm.linear_1.weight"] = input_model["cls/predictions/transform/dense/kernel"]
+    output_model["target.mlm.linear_1.bias"] = input_model["cls/predictions/transform/dense/bias"]
+    output_model["target.mlm.layer_norm.gamma"] = input_model["cls/predictions/transform/LayerNorm/gamma"]
+    output_model["target.mlm.layer_norm.beta"] = input_model["cls/predictions/transform/LayerNorm/beta"]
+    output_model["target.mlm.linear_2.weight"] = input_model["bert/embeddings/word_embeddings"]
+    output_model["target.mlm.linear_2.bias"] = input_model["cls/predictions/output_bias"]
+
+    torch.save(output_model, args.output_model_path)
+
+
+if __name__ == "__main__":
+    main()