Add unit tests for UnifiedTransformer

paddlenlp/transformers/unified_transformer/modeling.py

@@ -159,7 +159,8 @@ def __init__(self,
                  hidden_dropout_prob=0.1,
                  max_position_embeddings=512,
                  type_vocab_size=2,
-                 role_type_size=None):
+                 role_type_size=None,
+                 pad_token_id=None):
         super(UnifiedTransformerEmbeddings, self).__init__()
         self.word_embeddings = nn.Embedding(vocab_size, hidden_size)
         self.position_embeddings = nn.Embedding(max_position_embeddings,
@@ -169,9 +170,39 @@ def __init__(self,
             role_type_size, hidden_size)
         self.dropout = nn.Dropout(hidden_dropout_prob)
 
-    def forward(self, input_ids, token_type_ids, position_ids, role_ids=None):
+        self.pad_token_id = pad_token_id
+
+    def forward(self,
+                input_ids,
+                token_type_ids=None,
+                position_ids=None,
+                role_ids=None):
+        if position_ids is None:
+            if self.pad_token_id is None:
+                position_ids = paddle.expand_as(
+                    paddle.arange(end=paddle.shape(input_ids)[1],
+                                  dtype="int64"), input_ids)
+            else:
+                # NOTE: If there is a unk_token_id in input_ids, the following logic is wrong.
+                # In that case, the position_ids must be provided.
+                # And this is for left padding input_ids.
+                num_pad = paddle.sum(
+                    (input_ids == self.pad_token_id).astype("float32"),
+                    axis=-1,
+                    keepdim=True)
+                position_ids = F.relu(
+                    paddle.expand_as(
+                        paddle.arange(end=paddle.shape(input_ids)[1],
+                                      dtype="float32"), input_ids) -
+                    num_pad).astype("int64")
+            position_ids.stop_gradient = True
+
         input_embedings = self.word_embeddings(input_ids)
         position_embeddings = self.position_embeddings(position_ids)
+
+        if token_type_ids is None:
+            token_type_ids = paddle.zeros_like(input_ids, dtype="int64")
+            token_type_ids.stop_gradient = True
         token_type_embeddings = self.token_type_embeddings(token_type_ids)
 
         embeddings = input_embedings + position_embeddings + token_type_embeddings
@@ -283,7 +314,8 @@ def __init__(self,
                                                        hidden_dropout_prob,
                                                        max_position_embeddings,
                                                        type_vocab_size,
-                                                       role_type_size)
+                                                       role_type_size,
+                                                       self.pad_token_id)
         encoder_layer = nn.TransformerEncoderLayer(
             hidden_size,
             num_attention_heads,
@@ -298,11 +330,17 @@ def __init__(self,
                                              encoder_norm)
         self.apply(self.init_weights)
 
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
     def forward(self,
                 input_ids,
-                token_type_ids,
-                position_ids,
-                attention_mask,
+                token_type_ids=None,
+                position_ids=None,
+                attention_mask=None,
                 use_cache=False,
                 cache=None,
                 role_ids=None):
@@ -382,6 +420,10 @@ def forward(self,
                     is_split_into_words=False)
                 outputs = model(**inputs)
         """
+        if attention_mask is None:
+            attention_mask = ((input_ids == self.pad_token_id).astype(
+                paddle.get_default_dtype()) * -1e4).unsqueeze([1, 2])
+            attention_mask.stop_gradient = True
 
         embedding_output = self.embeddings(input_ids,
                                            token_type_ids,
@@ -454,9 +496,9 @@ def __init__(self, unified_transformer):
 
     def forward(self,
                 input_ids,
-                token_type_ids,
-                position_ids,
-                attention_mask,
+                token_type_ids=None,
+                position_ids=None,
+                attention_mask=None,
                 masked_positions=None,
                 use_cache=False,
                 cache=None,
@@ -549,30 +591,62 @@ def prepare_faster_entry(self, kwargs):
 
     def adjust_logits_during_generation(self, logits):
         # pre-process distribution
-        logits[:, self.unified_transformer.unk_token_id] = -1e9
-        logits[:, self.unified_transformer.bos_token_id] = -1e9
-        logits[:, self.unified_transformer.mask_token_id] = -1e9
+        logits[:, self.unified_transformer.unk_token_id] = -1e4
+        logits[:, self.unified_transformer.bos_token_id] = -1e4
+        logits[:, self.unified_transformer.mask_token_id] = -1e4
         return logits
 
     def prepare_inputs_for_generation(self,
                                       input_ids,
-                                      token_type_ids,
-                                      position_ids,
-                                      attention_mask,
+                                      token_type_ids=None,
+                                      position_ids=None,
+                                      attention_mask=None,
                                       use_cache=False,
                                       cache=None,
                                       **kwargs):
 
         role_ids = kwargs.get("role_ids", None)
 
+        if position_ids is None:
+            if self.pad_token_id is None:
+                position_ids = paddle.expand_as(
+                    paddle.arange(end=paddle.shape(input_ids)[1],
+                                  dtype="int64"), input_ids)
+            else:
+                # NOTE: If there is a unk_token_id in input_ids, the following logic is wrong.
+                # In that case, the position_ids must be provided.
+                # And this is for left padding input_ids.
+                num_pad = paddle.sum(
+                    (input_ids == self.pad_token_id).astype("float32"),
+                    axis=-1,
+                    keepdim=True)
+                position_ids = F.relu(
+                    paddle.expand_as(
+                        paddle.arange(end=paddle.shape(input_ids)[1],
+                                      dtype="float32"), input_ids) -
+                    num_pad).astype("int64")
+            position_ids.stop_gradient = True
+
+        if token_type_ids is None:
+            token_type_ids = paddle.zeros_like(input_ids, dtype="int64")
+            token_type_ids.stop_gradient = True
+
+        if attention_mask is None:
+            attention_mask = ((input_ids == self.pad_token_id).astype(
+                paddle.get_default_dtype()) * -1e4).unsqueeze([1, 2])
+            attention_mask.stop_gradient = True
+
         # only last token for inputs_ids if cache is defined in kwargs
         if cache is not None:
             input_ids = input_ids[:, -1:]
-            token_type_ids = token_type_ids[:, -1:]
-            position_ids = position_ids[:, -1:]
-            attention_mask = attention_mask[:, :, -1:, :]
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1:]
+            if position_ids is not None:
+                position_ids = position_ids[:, -1:]
             if role_ids is not None:
                 role_ids = role_ids[:, -1:]
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, :, -1:, :]
 
         return {
             "input_ids": input_ids,

paddlenlp/transformers/unified_transformer/tokenizer.py

@@ -188,6 +188,13 @@ def vocab_size(self):
         """
         return len(self.vocab)
 
+    def get_vocab(self):
+        vocab = {
+            self.convert_ids_to_tokens(i): i
+            for i in range(self.vocab_size)
+        }
+        return vocab
+
     def preprocess_text(self,
                         inputs,
                         remove_space=True,

tests/transformers/unified_transformer/__init__.py

@@ -0,0 +1,13 @@
+# 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.