Add PLATO-XL example (#1708)

FrostML · web-flow · commit dbd26feb8031 · 2022-03-03T15:48:36.000+08:00
diff --git a/examples/dialogue/plato-xl/README.md b/examples/dialogue/plato-xl/README.md
@@ -0,0 +1,45 @@
+# PLATO-XL
+
+## 模型简介
+
+构建高质量的开放领域（Open-Domain）的对话机器人，使得它能用自然语言与人自由地交流，这一直是自然语言处理领域终极目标之一。
+
+为了能够简易地构建一个高质量的开放域聊天机器人，本项目在 Paddle 上实现了 PLATO-XL 的预测模型，并实现了高性能的预测加速，而整套 float16 的方案可以确保在 32G V100 单卡上就能 load 并执行 11B 的 PLATO-XL 模型，无需再涉及 float32 相关计算。
+
+此外，PLATO-XL 72-layers, 32-heads, 3072-hidden，网络参数量较大，即使是在使用 float16 的情况下，72 层网络至少需要显存约 24G，并且需要保证当前使用的 GPU 支持 float16 的计算。
+
+其中：
+* 支持 float16 的 GPU 信息可以在 NVIDIA [官网](https://docs.nvidia.com/deeplearning/tensorrt/support-matrix/index.html#hardware-precision-matrix) 上查询；
+* 您当前使用的 GPU 的 compute capability 同样可以在 NVIDIA [官网](https://developer.nvidia.com/zh-cn/cuda-gpus#compute) 上找到，与上面链接中表格对应。
+
+PLATO-XL 的训练过程及其他细节详见 [Knover](https://github.com/PaddlePaddle/Knover/tree/develop/projects/PLATO-XL)
+
+## 快速开始
+
+### 环境依赖
+
+- python 3.7+
+- sentencepiece
+
+安装方式：
+``` python
+pip install sentencepiece
+```
+
+### 高性能生成
+
+使用 `infer.py` 脚本进行测试，无需单独下载预训练，脚本将自行下载。运行如下命令即可进行高性能预测，forward 将自动循环 200 次前向以供性能测试需要。
+
+```shell
+export CUDA_VISIBLE_DEVICES=0
+python infer.py --use_role --position_style relative --max_out_len 64 --min_out_len 1 --topk 4
+```
+
+该脚本各个参数含义如下：
+
+* `--use_role`: 是否使用 role embedding。
+* `--position_style`: 位置编码方式，这里可以选择是 "relative" 或是 "continuous"。
+* `--max_out_len`: 最长的输出的长度。
+* `--min_out_len`: 最短的输出长度。
+* `--topk`: 用于 top_k sampling 的 k 值的设定。
+* `--topp`: 用于 top_p sampling 的 p 值的设定。
diff --git a/examples/dialogue/plato-xl/infer.py b/examples/dialogue/plato-xl/infer.py
@@ -0,0 +1,153 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+import time
+import argparse
+from pprint import pprint
+
+import paddle
+
+from paddlenlp.transformers import UnifiedTransformerModel, UnifiedTransformerLMHeadModel, UnifiedTransformerTokenizer
+
+
+def setup_args():
+    """Setup arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--use_role",
+        action="store_true",
+        help="Whether to use role embeddings. ")
+    parser.add_argument(
+        "--position_style",
+        default="relative",
+        choices=["continuous", "relative"],
+        type=str,
+        help="The type for positional embedding. Default is continuous. ")
+    parser.add_argument(
+        "--max_out_len",
+        default=64,
+        type=int,
+        help="Maximum output sequence length. ")
+    parser.add_argument(
+        "--min_out_len",
+        default=1,
+        type=int,
+        help="Minimum output sequence length. ")
+    parser.add_argument(
+        "--topk",
+        default=4,
+        type=int,
+        help="The k value for topk_sampling. Default is 4. ")
+    parser.add_argument(
+        "--topp",
+        default=1.0,
+        type=float,
+        help="The p value for topp_sampling. Default is 0.0f. ")
+    parser.add_argument(
+        "--use_fp16_decoding",
+        action="store_true",
+        help="Whether to use fp16 decoding to predict. ")
+    parser.add_argument(
+        "--decoding_strategy",
+        default="sampling",
+        choices=["sampling", "beam_search"],
+        type=str,
+        help="The main strategy to decode. ")
+    parser.add_argument(
+        "--num_beams",
+        default=4,
+        type=int,
+        help="The number of candidate to procedure beam search. ")
+
+    args = parser.parse_args()
+
+    return args
+
+
+def postprocess_response(token_ids, tokenizer):
+    """Post-process the decoded sequence. Truncate from the first <eos>."""
+    eos_pos = len(token_ids)
+    for i, tok_id in enumerate(token_ids):
+        if tok_id == tokenizer.sep_token_id:
+            eos_pos = i
+            break
+    token_ids = token_ids[:eos_pos]
+    tokens = tokenizer.convert_ids_to_tokens(token_ids)
+    tokens = tokenizer.merge_subword(tokens)
+    return tokens
+
+
+def infer(args):
+    model_name = 'plato-xl'
+    model = UnifiedTransformerLMHeadModel.from_pretrained(model_name)
+    tokenizer = UnifiedTransformerTokenizer.from_pretrained(model_name)
+
+    context = [
+        "Hi , Becky , what's up ?",
+        "Not much , except that my mother-in-law is driving me up the wall .",
+        "What's the problem ?"
+    ]
+
+    data = tokenizer.dialogue_encode(
+        history=context,
+        add_start_token_as_response=True,
+        return_length=True,
+        return_role_ids=args.use_role,
+        position_style=args.position_style)
+
+    for name in data:
+        if name == "attention_mask":
+            data[name] = paddle.to_tensor(
+                data[name], dtype="float32").reshape([1, 1, 41, 41])
+        else:
+            data[name] = paddle.to_tensor(
+                data[name], dtype="int32").reshape([1, -1])
+
+    for i in range(200):
+        if 100 == i:
+            paddle.device.cuda.synchronize()
+            start = time.time()
+
+        outputs, _ = model.generate(
+            input_ids=data['input_ids'],
+            token_type_ids=data['token_type_ids'],
+            position_ids=data['position_ids'],
+            attention_mask=data['attention_mask'],
+            role_ids=data.get('role_ids', None),
+            seq_len=data['seq_len'],
+            max_length=args.max_out_len,
+            min_length=args.min_out_len,
+            decode_strategy=args.decoding_strategy,
+            top_k=args.topk,
+            top_p=args.topp,
+            num_beams=args.num_beams,
+            use_fp16_decoding=args.use_fp16_decoding,
+            use_faster=True)
+
+    paddle.device.cuda.synchronize()
+    print("Average time of FasterGeneration of PLATO-XL model is {}ms. ".format(
+        (time.time() - start) / 100 * 1000))
+
+    result = postprocess_response(outputs[0].numpy(), tokenizer)
+    result = " ".join(result)
+
+    print("Model input:", context)
+    print("Result:", result)
+
+
+if __name__ == "__main__":
+    args = setup_args()
+    pprint(args)
+
+    infer(args)
diff --git a/paddlenlp/transformers/unified_transformer/modeling.py b/paddlenlp/transformers/unified_transformer/modeling.py
@@ -96,6 +96,26 @@ class UnifiedTransformerPretrainedModel(PretrainedModel):
             "eos_token_id": 2,
             "mask_token_id": 30000,
         },
+        "plato-xl": {
+            "vocab_size": 8001,
+            "hidden_size": 3072,
+            "num_hidden_layers": 72,
+            "num_attention_heads": 32,
+            "intermediate_size": 18432,
+            "hidden_act": "gelu",
+            "hidden_dropout_prob": 0.1,
+            "attention_probs_dropout_prob": 0.1,
+            "normalize_before": True,
+            "max_position_embeddings": 1024,
+            "type_vocab_size": 2,
+            "role_type_size": 128,
+            "initializer_range": 0.02,
+            "unk_token_id": 0,
+            "pad_token_id": 0,
+            "bos_token_id": 1,
+            "eos_token_id": 2,
+            "mask_token_id": 8000,
+        }
     }
     resource_files_names = {"model_state": "model_state.pdparams"}
     pretrained_resource_files_map = {
@@ -106,6 +126,8 @@ class UnifiedTransformerPretrainedModel(PretrainedModel):
             "https://bj.bcebos.com/paddlenlp/models/transformers/unified_transformer/unified_transformer-12L-cn-luge.pdparams",
             "plato-mini":
             "https://bj.bcebos.com/paddlenlp/models/transformers/unified_transformer/plato-mini.pdparams",
+            "plato-xl":
+            "https://bj.bcebos.com/paddlenlp/models/transformers/unified_transformer/plato-xl.pdparams",
         }
     }
     base_model_prefix = "unified_transformer"
@@ -115,7 +137,9 @@ def init_weights(self, layer):
         if isinstance(layer, (nn.Linear, nn.Embedding)):
             # In the dygraph mode, use the `set_value` to reset the parameter directly,
             # and reset the `state_dict` to update parameter in static mode.
-            if isinstance(layer.weight, paddle.Tensor):
+            if isinstance(
+                    layer.weight,
+                    paddle.Tensor) and paddle.get_default_dtype() == "float32":
                 layer.weight.set_value(
                     paddle.tensor.normal(
                         mean=0.0,
@@ -133,20 +157,27 @@ def __init__(self,
                  hidden_size=768,
                  hidden_dropout_prob=0.1,
                  max_position_embeddings=512,
-                 type_vocab_size=2):
+                 type_vocab_size=2,
+                 role_type_size=None):
         super(UnifiedTransformerEmbeddings, self).__init__()
         self.word_embeddings = nn.Embedding(vocab_size, hidden_size)
         self.position_embeddings = nn.Embedding(max_position_embeddings,
                                                 hidden_size)
         self.token_type_embeddings = nn.Embedding(type_vocab_size, hidden_size)
+        self.role_embeddings = None if role_type_size is None else nn.Embedding(
+            role_type_size, hidden_size)
         self.dropout = nn.Dropout(hidden_dropout_prob)
 
-    def forward(self, input_ids, token_type_ids, position_ids):
+    def forward(self, input_ids, token_type_ids, position_ids, role_ids=None):
         input_embedings = self.word_embeddings(input_ids)
         position_embeddings = self.position_embeddings(position_ids)
         token_type_embeddings = self.token_type_embeddings(token_type_ids)
 
         embeddings = input_embedings + position_embeddings + token_type_embeddings
+
+        if self.role_embeddings is not None:
+            embeddings += self.role_embeddings(role_ids)
+
         embeddings = self.dropout(embeddings)
         return embeddings
 
@@ -221,25 +252,25 @@ class UnifiedTransformerModel(UnifiedTransformerPretrainedModel):
             The id of special token `mask_token`. Defaults to 30000.
     """
 
-    def __init__(
-            self,
-            vocab_size,
-            hidden_size=768,
-            num_hidden_layers=12,
-            num_attention_heads=12,
-            intermediate_size=3072,
-            hidden_act="gelu",
-            hidden_dropout_prob=0.1,
-            attention_probs_dropout_prob=0.1,
-            normalize_before=True,
-            max_position_embeddings=512,
-            type_vocab_size=2,
-            initializer_range=0.02,
-            unk_token_id=0,
-            pad_token_id=0,
-            bos_token_id=1,
-            eos_token_id=2,
-            mask_token_id=30000, ):
+    def __init__(self,
+                 vocab_size,
+                 hidden_size=768,
+                 num_hidden_layers=12,
+                 num_attention_heads=12,
+                 intermediate_size=3072,
+                 hidden_act="gelu",
+                 hidden_dropout_prob=0.1,
+                 attention_probs_dropout_prob=0.1,
+                 normalize_before=True,
+                 max_position_embeddings=512,
+                 type_vocab_size=2,
+                 initializer_range=0.02,
+                 unk_token_id=0,
+                 pad_token_id=0,
+                 bos_token_id=1,
+                 eos_token_id=2,
+                 mask_token_id=30000,
+                 role_type_size=None):
         super(UnifiedTransformerModel, self).__init__()
         self.unk_token_id = unk_token_id
         self.pad_token_id = pad_token_id
@@ -250,7 +281,7 @@ def __init__(
 
         self.embeddings = UnifiedTransformerEmbeddings(
             vocab_size, hidden_size, hidden_dropout_prob,
-            max_position_embeddings, type_vocab_size)
+            max_position_embeddings, type_vocab_size, role_type_size)
         encoder_layer = nn.TransformerEncoderLayer(
             hidden_size,
             num_attention_heads,
diff --git a/paddlenlp/transformers/unified_transformer/tokenizer.py b/paddlenlp/transformers/unified_transformer/tokenizer.py
