fix roberta

Author: zheyuye

scripts/conversion_toolkits/convert_fairseq_roberta.py

@@ -37,7 +37,7 @@ def convert_vocab(args, fairseq_model):
     fairseq_vocab = fairseq_model.task.dictionary
     # bos_word attr missing in fairseq_vocab
     fairseq_vocab.bos_word = fairseq_vocab[fairseq_vocab.bos_index]
-    
+
     assert os.path.exists(fairseq_dict_path), \
            '{} not found'.format(fairseq_dict_path)
     from mxnet.gluon.utils import download
@@ -63,7 +63,7 @@ def convert_vocab(args, fairseq_model):
     inter_vocab = list(inter_vocab.items())
     inter_vocab = sorted(inter_vocab, key=lambda x : x[1])
     tokens = [e[0] for e in inter_vocab]
-    
+
     tail = [fairseq_vocab[-4],
             fairseq_vocab[-3],
             fairseq_vocab[-2],
@@ -84,15 +84,15 @@ def convert_vocab(args, fairseq_model):
     gluon_vocab.save(vocab_save_path)
     os.remove(temp_vocab_file)
     os.remove(temp_merges_file)
-    
+
     gluon_tokenizer = HuggingFaceByteBPETokenizer(
         merges_save_path,
         vocab_save_path
     )
-    
+
     if args.test:
         test_vocab(fairseq_model, gluon_tokenizer)
-    
+
     vocab_size = len(fairseq_vocab)
     print('| converted dictionary: {} types'.format(vocab_size))
     return vocab_size
@@ -103,14 +103,14 @@ def test_vocab(fairseq_model, gluon_tokenizer, check_all_tokens=False):
     gluon_vocab = gluon_tokenizer.vocab
     assert len(fairseq_vocab) == \
            len(gluon_vocab)
-    
+
     # assert all_tokens
     # roberta with gpt2 bytebpe bpe does not provide all tokens directly
     if check_all_tokens:
         for i in range(len(fairseq_vocab)):
             assert fairseq_vocab[i] == gluon_vocab.all_tokens[i], \
                    '{}, {}, {}'.format(i, fairseq_vocab[i], gluon_vocab.all_tokens[i])
-    
+
     # assert special tokens
     for special_tokens in ['unk', 'pad', 'eos', 'bos']:
         assert getattr(fairseq_vocab, special_tokens + '_index') == \
@@ -121,7 +121,7 @@ def test_vocab(fairseq_model, gluon_tokenizer, check_all_tokens=False):
         assert fairseq_vocab[-1] == \
                gluon_vocab.all_tokens[-1] == \
                '<mask>'
-    
+
     sentence = "Hello, y'all! How are you Ⅷ 😁 😁 😁 ?" + \
                'GluonNLP is great！！！!!!' + \
                "GluonNLP-Amazon-Haibin-Leonard-Sheng-Shuai-Xingjian...../:!@# 'abc'"
@@ -131,7 +131,7 @@ def test_vocab(fairseq_model, gluon_tokenizer, check_all_tokens=False):
     # Notice: we may append bos and eos
     # manuually after tokenizing sentences
     assert fs_tokens.numpy().tolist()[1:-1] == gl_tokens
-    
+
     # assert decode
     fs_sentence = fairseq_model.decode(fs_tokens)
     gl_sentence = gluon_tokenizer.decode(gl_tokens)
@@ -170,7 +170,8 @@ def convert_params(fairseq_model,
     fairseq_prefix = 'model.decoder.'
     gluon_model = gluon_model_cls.from_cfg(
         gluon_cfg,
-        return_all_hiddens=True,
+        use_pooler=False,
+        output_all_encodings=True,
         prefix=gluon_prefix
     )
     gluon_model.initialize(ctx=ctx)
@@ -196,7 +197,7 @@ def convert_params(fairseq_model,
             np.concatenate([fs_q_weight, fs_k_weight, fs_v_weight], axis=0))
         gl_qkv_bias.set_data(
             np.concatenate([fs_q_bias, fs_k_bias, fs_v_bias], axis=0))
-        
+
         for k, v in [
             ('self_attn.out_proj.weight', 'proj_weight'),
             ('self_attn.out_proj.bias', 'proj_bias'),
@@ -230,20 +231,20 @@ def convert_params(fairseq_model,
         gl_name = gluon_prefix + v
         gluon_params[gl_name].set_data(
             fairseq_params[fs_name].cpu().numpy())
-    
+
     # position embed weight
     padding_idx = fairseq_model.task.dictionary.pad_index
     fs_pos_embed_name = fairseq_prefix + 'sentence_encoder.embed_positions.weight'
     gl_pos_embed_name = gluon_prefix + 'pos_embed_embed_weight'
     gluon_params[gl_pos_embed_name].set_data(
         fairseq_params[fs_pos_embed_name].cpu().numpy()[padding_idx + 1:,:])
-    
+
     # assert untie=False
     assert np.array_equal(
         fairseq_params[fairseq_prefix + 'sentence_encoder.embed_tokens.weight'].cpu().numpy(),
         fairseq_params[fairseq_prefix + 'lm_head.weight'].cpu().numpy()
     )
-    
+
     return gluon_model
 
 def test_model(fairseq_model, gluon_model, gpu):
@@ -272,16 +273,16 @@ def test_model(fairseq_model, gluon_model, gpu):
     fs_input_ids = torch.from_numpy(input_ids).cuda(gpu)
     if gpu is not None:
         fs_input_ids = fs_input_ids.cuda(gpu)
-    
+
     fairseq_model.model.eval()
-    
-    gl_x, gl_all_hiddens = \
+
+    gl_all_hiddens, gl_x = \
         gluon_model(gl_input_ids, gl_valid_length)
 
     fs_x, fs_extra = \
         fairseq_model.model.cuda(gpu)(fs_input_ids, return_all_hiddens=True)
     fs_all_hiddens = fs_extra['inner_states']
-    
+
     num_layers = fairseq_model.args.encoder_layers
     for i in range(num_layers + 1):
         gl_hidden = gl_all_hiddens[i].asnumpy()
@@ -317,7 +318,7 @@ def rename(save_dir):
         new_name = '{file_prefix}-{short_hash}.{file_sufix}'.format(
                 file_prefix=file_prefix,
                 short_hash=long_hash[:8],
-                file_sufix=file_sufix)   
+                file_sufix=file_sufix)
         new_path = os.path.join(save_dir, new_name)
         shutil.move(old_path, new_path)
         file_size = os.path.getsize(new_path)

src/gluonnlp/models/roberta.py

@@ -130,7 +130,7 @@ def __init__(self,
                  use_mlm=True,
                  untie_weight=False,
                  encoder_normalize_before=True,
-                 return_all_hiddens=False,
+                 output_all_encodings=False,
                  prefix=None,
                  params=None):
         """
@@ -160,7 +160,7 @@ def __init__(self,
         untie_weight
             Whether to untie weights between embeddings and classifiers
         encoder_normalize_before
-        return_all_hiddens
+        output_all_encodings
         prefix
         params
         """
@@ -182,7 +182,7 @@ def __init__(self,
         self.use_mlm = use_mlm
         self.untie_weight = untie_weight
         self.encoder_normalize_before = encoder_normalize_before
-        self.return_all_hiddens = return_all_hiddens
+        self.output_all_encodings = output_all_encodings
         with self.name_scope():
             self.tokens_embed = nn.Embedding(
                 input_dim=self.vocab_size,
@@ -220,7 +220,7 @@ def __init__(self,
                 bias_initializer=bias_initializer,
                 activation=self.activation,
                 dtype=self.dtype,
-                return_all_hiddens=self.return_all_hiddens
+                output_all_encodings=self.output_all_encodings
             )
             self.encoder.hybridize()
 
@@ -237,25 +237,63 @@ def __init__(self,
                     bias_initializer=bias_initializer
                 )
                 self.lm_head.hybridize()
-            # TODO support use_pooler
 
     def hybrid_forward(self, F, tokens, valid_length):
-        x = self.tokens_embed(tokens)
+        outputs = []
+        embedding = self.get_initial_embedding(F, tokens)
+
+        inner_states = self.encoder(x, valid_length)
+        if self.output_all_encodings:
+            contextual_embeddings = inner_states[-1]
+        else:
+            contextual_embeddings = inner_states
+        outputs.append(contextual_embeddings)
+
+        if self.use_pooler:
+            pooled_out = self.apply_pooling(contextual_embeddings)
+            outputs.append(pooled_out)
+            
+        if self.use_mlm:
+            mlm_output = self.lm_head(contextual_embeddings)
+            outputs.append(mlm_output)
+        return tuple(outputs) if len(outputs) > 1 else outputs[0]
+
+    def get_initial_embedding(self, F, inputs):
+        """Get the initial token embeddings that considers the token type and positional embeddings
+
+        Parameters
+        ----------
+        F
+        inputs
+            Shape (batch_size, seq_length)
+
+        Returns
+        -------
+        embedding
+            The initial embedding that will be fed into the encoder
+        """
+        embedding = self.tokens_embed(inputs)
         if self.pos_embed_type:
-            positional_embedding = self.pos_embed(F.npx.arange_like(x, axis=1))
+            positional_embedding = self.pos_embed(F.npx.arange_like(inputs, axis=1))
             positional_embedding = F.np.expand_dims(positional_embedding, axis=0)
-            x = x + positional_embedding
+            embedding = embedding + positional_embedding
         if self.embed_ln:
-            x = self.embed_ln(x)
-        x = self.embed_dropout(x)
-        inner_states = self.encoder(x, valid_length)
-        x = inner_states[-1]
-        if self.use_mlm:
-            x = self.lm_head(x)
-        if self.return_all_hiddens:
-            return x, inner_states
-        else:
-            return x
+            embedding = self.embed_ln(embedding)
+        embedding = self.embed_dropout(embedding)
+
+    def apply_pooling(self, sequence):
+        """Generate the representation given the inputs.
+
+        This is used for pre-training or fine-tuning a mobile bert model.
+        Get the first token of the whole sequence which is [CLS]
+
+        sequence:
+            Shape (batch_size, sequence_length, units)
+        return:
+            Shape (batch_size, units)
+        """
+        outputs = sequence[:, 0, :]
+        return outputs
 
     @staticmethod
     def get_cfg(key=None):
@@ -271,7 +309,7 @@ def from_cfg(cls,
                  use_mlm=True,
                  untie_weight=False,
                  encoder_normalize_before=True,
-                 return_all_hiddens=False,
+                 output_all_encodings=False,
                  prefix=None,
                  params=None):
         cfg = RobertaModel.get_cfg().clone_merge(cfg)
@@ -298,7 +336,7 @@ def from_cfg(cls,
                    use_mlm=use_mlm,
                    untie_weight=untie_weight,
                    encoder_normalize_before=encoder_normalize_before,
-                   return_all_hiddens=return_all_hiddens,
+                   output_all_encodings=output_all_encodings,
                    prefix=prefix,
                    params=params)
 
@@ -316,7 +354,7 @@ def __init__(self,
                  bias_initializer='zeros',
                  activation='gelu',
                  dtype='float32',
-                 return_all_hiddens=False,
+                 output_all_encodings=False,
                  prefix='encoder_',
                  params=None):
         super(RobertaEncoder, self).__init__(prefix=prefix, params=params)
@@ -329,7 +367,7 @@ def __init__(self,
         self.layer_norm_eps = layer_norm_eps
         self.activation = activation
         self.dtype = dtype
-        self.return_all_hiddens = return_all_hiddens
+        self.output_all_encodings = output_all_encodings
         with self.name_scope():
             self.all_layers = nn.HybridSequential(prefix='layers_')
             with self.all_layers.name_scope():
@@ -358,8 +396,8 @@ def hybrid_forward(self, F, x, valid_length):
             layer = self.all_layers[layer_idx]
             x, _ = layer(x, atten_mask)
             inner_states.append(x)
-        if not self.return_all_hiddens:
-            inner_states = [x]
+        if not self.output_all_encodings:
+            inner_states = x
         return inner_states
 
 @use_np
@@ -419,7 +457,8 @@ def list_pretrained_roberta():
 
 def get_pretrained_roberta(model_name: str = 'fairseq_roberta_base',
                            root: str = get_model_zoo_home_dir(),
-                           load_backbone: bool = True) \
+                           load_backbone: bool = True,
+                           load_mlm: bool = False) \
         -> Tuple[CN, HuggingFaceByteBPETokenizer, str]:
     """Get the pretrained RoBERTa weights