Source code for gluonnlp.model.bert
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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.
"""BERT models."""
# pylint: disable=too-many-lines
__all__ = ['BERTModel', 'RoBERTaModel', 'BERTEncoder', 'BERTClassifier',
'RoBERTaClassifier', 'bert_12_768_12', 'bert_24_1024_16',
'ernie_12_768_12', 'roberta_12_768_12', 'roberta_24_1024_16',
'DistilBERTModel', 'distilbert_6_768_12']
import os
import mxnet as mx
from mxnet.gluon import HybridBlock, nn
from mxnet.gluon.model_zoo import model_store
from ..base import get_home_dir
from .block import GELU
from .seq2seq_encoder_decoder import Seq2SeqEncoder
from .transformer import PositionwiseFFN
from .utils import _load_pretrained_params, _load_vocab
###############################################################################
# COMPONENTS #
###############################################################################
class DotProductSelfAttentionCell(HybridBlock):
r"""Multi-head Dot Product Self Attention Cell.
In the DotProductSelfAttentionCell, the input query/key/value will be linearly projected
for `num_heads` times with different projection matrices. Each projected key, value, query
will be used to calculate the attention weights and values. The output of each head will be
concatenated to form the final output.
This is a more efficient implementation of MultiHeadAttentionCell with
DotProductAttentionCell as the base_cell:
score = <W_q h_q, W_k h_k> / sqrt(dim_q)
Parameters
----------
units : int
Total number of projected units for query. Must be divided exactly by num_heads.
num_heads : int
Number of parallel attention heads
use_bias : bool, default True
Whether to use bias when projecting the query/key/values
weight_initializer : str or `Initializer` or None, default None
Initializer of the weights.
bias_initializer : str or `Initializer`, default 'zeros'
Initializer of the bias.
prefix : str or None, default None
See document of `Block`.
params : str or None, default None
See document of `Block`.
Inputs:
- **qkv** : Symbol or NDArray
Query / Key / Value vector. Shape (query_length, batch_size, C_in)
- **valid_len** : Symbol or NDArray or None, default None
Valid length of the query/key/value slots. Shape (batch_size, query_length)
Outputs:
- **context_vec** : Symbol or NDArray
Shape (query_length, batch_size, context_vec_dim)
- **att_weights** : Symbol or NDArray
Attention weights of multiple heads.
Shape (batch_size, num_heads, query_length, memory_length)
"""
def __init__(self, units, num_heads, dropout=0.0, use_bias=True,
weight_initializer=None, bias_initializer='zeros',
prefix=None, params=None):
super().__init__(prefix=prefix, params=params)
self._num_heads = num_heads
self._use_bias = use_bias
self._dropout = dropout
self.units = units
with self.name_scope():
if self._use_bias:
self.query_bias = self.params.get('query_bias', shape=(self.units,),
init=bias_initializer)
self.key_bias = self.params.get('key_bias', shape=(self.units,),
init=bias_initializer)
self.value_bias = self.params.get('value_bias', shape=(self.units,),
init=bias_initializer)
weight_shape = (self.units, self.units)
self.query_weight = self.params.get('query_weight', shape=weight_shape,
init=weight_initializer,
allow_deferred_init=True)
self.key_weight = self.params.get('key_weight', shape=weight_shape,
init=weight_initializer,
allow_deferred_init=True)
self.value_weight = self.params.get('value_weight', shape=weight_shape,
init=weight_initializer,
allow_deferred_init=True)
self.dropout_layer = nn.Dropout(self._dropout)
def _collect_params_with_prefix(self, prefix=''):
# the registered parameter names in v0.8 are the following:
# prefix_proj_query.weight, prefix_proj_query.bias
# prefix_proj_value.weight, prefix_proj_value.bias
# prefix_proj_key.weight, prefix_proj_key.bias
# this is a temporary fix to keep backward compatibility, due to an issue in MXNet:
# https://github.com/apache/incubator-mxnet/issues/17220
if prefix:
prefix += '.'
ret = {prefix + 'proj_' + k.replace('_', '.') : v for k, v in self._reg_params.items()}
for name, child in self._children.items():
ret.update(child._collect_params_with_prefix(prefix + name))
return ret
# pylint: disable=arguments-differ
def hybrid_forward(self, F, qkv, valid_len, query_bias, key_bias, value_bias,
query_weight, key_weight, value_weight):
# interleaved_matmul_selfatt ops assume the projection is done with interleaving
# weights for query/key/value. The concatenated weight should have shape
# (num_heads, C_out/num_heads * 3, C_in).
query_weight = query_weight.reshape(shape=(self._num_heads, -1, 0), reverse=True)
key_weight = key_weight.reshape(shape=(self._num_heads, -1, 0), reverse=True)
value_weight = value_weight.reshape(shape=(self._num_heads, -1, 0), reverse=True)
in_weight = F.concat(query_weight, key_weight, value_weight, dim=-2)
in_weight = in_weight.reshape(shape=(-1, 0), reverse=True)
# concat bias
query_bias = query_bias.reshape(shape=(self._num_heads, -1), reverse=True)
key_bias = key_bias.reshape(shape=(self._num_heads, -1), reverse=True)
value_bias = value_bias.reshape(shape=(self._num_heads, -1), reverse=True)
in_bias = F.stack(query_bias, key_bias, value_bias, axis=1).reshape(-1)
# qkv_proj shape = (seq_length, batch_size, num_heads * head_dim * 3)
qkv_proj = F.FullyConnected(data=qkv, weight=in_weight, bias=in_bias,
num_hidden=self.units*3, no_bias=False, flatten=False)
att_score = F.contrib.interleaved_matmul_selfatt_qk(qkv_proj, heads=self._num_heads)
if valid_len is not None:
valid_len = F.broadcast_axis(F.expand_dims(valid_len, axis=1),
axis=1, size=self._num_heads)
valid_len = valid_len.reshape(shape=(-1, 0), reverse=True)
att_weights = F.softmax(att_score, length=valid_len, use_length=True, axis=-1)
else:
att_weights = F.softmax(att_score, axis=-1)
# att_weights shape = (batch_size, seq_length, seq_length)
att_weights = self.dropout_layer(att_weights)
context_vec = F.contrib.interleaved_matmul_selfatt_valatt(qkv_proj, att_weights,
heads=self._num_heads)
att_weights = att_weights.reshape(shape=(-1, self._num_heads, 0, 0), reverse=True)
return context_vec, att_weights
class BERTEncoderCell(HybridBlock):
"""Structure of the BERT Encoder Cell.
Parameters
----------
units : int
Number of units for the output
hidden_size : int
number of units in the hidden layer of position-wise feed-forward networks
num_heads : int
Number of heads in multi-head attention
dropout : float
output_attention: bool
Whether to output the attention weights
attention_use_bias : float, default True
Whether to use bias term in the attention cell
weight_initializer : str or Initializer
Initializer for the input weights matrix, used for the linear
transformation of the inputs.
bias_initializer : str or Initializer
Initializer for the bias vector.
prefix : str, default None
Prefix for name of `Block`s. (and name of weight if params is `None`).
params : Parameter or None
Container for weight sharing between cells. Created if `None`.
activation : str, default 'gelu'
Activation methods in PositionwiseFFN
layer_norm_eps : float, default 1e-5
Epsilon for layer_norm
Inputs:
- **inputs** : input sequence. Shape (length, batch_size, C_in)
- **valid_length** : valid length of inputs for attention. Shape (batch_size, length)
Outputs:
- **outputs**: output tensor of the transformer encoder cell.
Shape (length, batch_size, C_out)
- **additional_outputs**: the additional output of all the BERT encoder cell.
"""
def __init__(self, units=128, hidden_size=512, num_heads=4,
dropout=0.0, output_attention=False,
attention_use_bias=True,
weight_initializer=None, bias_initializer='zeros',
prefix=None, params=None, activation='gelu',
layer_norm_eps=1e-5):
super().__init__(prefix=prefix, params=params)
self._dropout = dropout
self._output_attention = output_attention
with self.name_scope():
if dropout:
self.dropout_layer = nn.Dropout(rate=dropout)
self.attention_cell = DotProductSelfAttentionCell(units, num_heads,
use_bias=attention_use_bias,
dropout=dropout)
self.proj = nn.Dense(units=units, flatten=False, use_bias=True,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer, prefix='proj_')
self.ffn = PositionwiseFFN(units=units, hidden_size=hidden_size, dropout=dropout,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer, activation=activation,
layer_norm_eps=layer_norm_eps)
self.layer_norm = nn.LayerNorm(in_channels=units, epsilon=layer_norm_eps)
def hybrid_forward(self, F, inputs, valid_len=None): # pylint: disable=arguments-differ
"""Transformer Encoder Attention Cell.
Parameters
----------
inputs : Symbol or NDArray
Input sequence. Shape (length, batch_size, C_in)
valid_len : Symbol or NDArray or None
Valid length for inputs. Shape (batch_size, length)
Returns
-------
encoder_cell_outputs: list
Outputs of the encoder cell. Contains:
- outputs of the transformer encoder cell. Shape (length, batch_size, C_out)
- additional_outputs of all the transformer encoder cell
"""
outputs, attention_weights = self.attention_cell(inputs, valid_len)
outputs = self.proj(outputs)
if self._dropout:
outputs = self.dropout_layer(outputs)
# use residual
outputs = outputs + inputs
outputs = self.layer_norm(outputs)
outputs = self.ffn(outputs)
additional_outputs = []
if self._output_attention:
additional_outputs.append(attention_weights)
return outputs, additional_outputs
[docs]class BERTEncoder(HybridBlock, Seq2SeqEncoder):
"""Structure of the BERT Encoder.
Different from the original encoder for transformer, `BERTEncoder` uses
learnable positional embedding, a 'gelu' activation functions and a
separate epsilon value for LayerNorm.
Parameters
----------
num_layers : int
Number of attention layers.
units : int
Number of units for the output.
hidden_size : int
number of units in the hidden layer of position-wise feed-forward networks
max_length : int
Maximum length of the input sequence
num_heads : int
Number of heads in multi-head attention
dropout : float
Dropout probability of the attention probabilities and embedding.
output_attention: bool, default False
Whether to output the attention weights
output_all_encodings: bool, default False
Whether to output encodings of all encoder cells
weight_initializer : str or Initializer
Initializer for the input weights matrix, used for the linear
transformation of the inputs.
bias_initializer : str or Initializer
Initializer for the bias vector.
prefix : str, default None.
Prefix for name of `Block`s. (and name of weight if params is `None`).
params : Parameter or None
Container for weight sharing between cells. Created if `None`.
activation : str, default 'gelu'
Activation methods in PositionwiseFFN
layer_norm_eps : float, default 1e-12
Epsilon for layer_norm
Inputs:
- **inputs** : input sequence of shape (length, batch_size, C_in)
- **states** : list of tensors for initial states and valid length for self attention.
- **valid_length** : valid lengths of each sequence. Usually used when part of sequence
has been padded. Shape is (batch_size, )
Outputs:
- **outputs** : the output of the encoder. Shape is (length, batch_size, C_out)
- **additional_outputs** : list of tensors.
Either be an empty list or contains the attention weights in this step.
The attention weights will have shape (batch_size, num_heads, length, mem_length)
"""
def __init__(self, *, num_layers=2, units=512, hidden_size=2048,
max_length=50, num_heads=4, dropout=0.0,
output_attention=False, output_all_encodings=False, weight_initializer=None,
bias_initializer='zeros', prefix=None, params=None, activation='gelu',
layer_norm_eps=1e-12):
super().__init__(prefix=prefix, params=params)
assert units % num_heads == 0,\
'In BERTEncoder, The units should be divided exactly ' \
'by the number of heads. Received units={}, num_heads={}' \
.format(units, num_heads)
self._max_length = max_length
self._units = units
self._output_attention = output_attention
self._output_all_encodings = output_all_encodings
self._dropout = dropout
self._layer_norm_eps = layer_norm_eps
with self.name_scope():
if dropout:
self.dropout_layer = nn.Dropout(rate=dropout)
self.layer_norm = nn.LayerNorm(in_channels=units, epsilon=self._layer_norm_eps)
self.position_weight = self.params.get('position_weight', shape=(max_length, units),
init=weight_initializer)
self.transformer_cells = nn.HybridSequential()
for i in range(num_layers):
cell = BERTEncoderCell(
units=units, hidden_size=hidden_size, num_heads=num_heads,
weight_initializer=weight_initializer,
bias_initializer=bias_initializer, dropout=dropout,
output_attention=output_attention, prefix='transformer%d_' % i,
activation=activation, layer_norm_eps=layer_norm_eps)
self.transformer_cells.add(cell)
def __call__(self, inputs, states=None, valid_length=None): # pylint: disable=arguments-differ
"""Encode the inputs given the states and valid sequence length.
Parameters
----------
inputs : NDArray or Symbol
Input sequence. Shape (batch_size, length, C_in)
states : list of NDArrays or Symbols
Initial states. The list of initial states and valid length for self attention
valid_length : NDArray or Symbol
Valid lengths of each sequence. This is usually used when part of sequence has
been padded. Shape (batch_size,)
Returns
-------
encoder_outputs: list
Outputs of the encoder. Contains:
- outputs of the transformer encoder. Shape (batch_size, length, C_out)
- additional_outputs of all the transformer encoder
"""
return super().__call__(inputs, states, valid_length)
[docs] def hybrid_forward(self, F, inputs, states=None, valid_length=None, position_weight=None):
# pylint: disable=arguments-differ
"""Encode the inputs given the states and valid sequence length.
Parameters
----------
inputs : NDArray or Symbol
Input sequence. Shape (length, batch_size, C_in)
states : list of NDArrays or Symbols
Initial states. The list of initial states and valid length for self attention
valid_length : NDArray or Symbol
Valid lengths of each sequence. This is usually used when part of sequence has
been padded. Shape (batch_size,)
Returns
-------
outputs : NDArray or Symbol, or List[NDArray] or List[Symbol]
If output_all_encodings flag is False, then the output of the last encoder.
If output_all_encodings flag is True, then the list of all outputs of all encoders.
In both cases, shape of the tensor(s) is/are (length, batch_size, C_out)
additional_outputs : list
Either be an empty list or contains the attention weights in this step.
The attention weights will have shape (batch_size, length) or
(batch_size, num_heads, length, length)
"""
# axis 0 is for length
steps = F.contrib.arange_like(inputs, axis=0)
if valid_length is not None:
zeros = F.zeros_like(steps)
# valid_length for attention, shape = (batch_size, seq_length)
attn_valid_len = F.broadcast_add(F.reshape(valid_length, shape=(-1, 1)),
F.reshape(zeros, shape=(1, -1)))
attn_valid_len = F.cast(attn_valid_len, dtype='int32')
if states is None:
states = [attn_valid_len]
else:
states.append(attn_valid_len)
else:
attn_valid_len = None
if states is None:
states = [steps]
else:
states.append(steps)
# positional encoding
positional_embed = F.Embedding(steps, position_weight, self._max_length, self._units)
inputs = F.broadcast_add(inputs, F.expand_dims(positional_embed, axis=1))
if self._dropout:
inputs = self.dropout_layer(inputs)
inputs = self.layer_norm(inputs)
outputs = inputs
all_encodings_outputs = []
additional_outputs = []
for cell in self.transformer_cells:
outputs, attention_weights = cell(inputs, attn_valid_len)
inputs = outputs
if self._output_all_encodings:
if valid_length is not None:
outputs = F.SequenceMask(outputs, sequence_length=valid_length,
use_sequence_length=True, axis=0)
all_encodings_outputs.append(outputs)
if self._output_attention:
additional_outputs.append(attention_weights)
if valid_length is not None and not self._output_all_encodings:
# if self._output_all_encodings, SequenceMask is already applied above
outputs = F.SequenceMask(outputs, sequence_length=valid_length,
use_sequence_length=True, axis=0)
if self._output_all_encodings:
return all_encodings_outputs, additional_outputs
return outputs, additional_outputs
###############################################################################
# FULL MODEL #
###############################################################################
[docs]class BERTModel(HybridBlock):
"""Generic Model for BERT (Bidirectional Encoder Representations from Transformers).
Parameters
----------
encoder : BERTEncoder
Bidirectional encoder that encodes the input sentence.
vocab_size : int or None, default None
The size of the vocabulary.
token_type_vocab_size : int or None, default None
The vocabulary size of token types (number of segments).
units : int or None, default None
Number of units for the final pooler layer.
embed_size : int or None, default None
Size of the embedding vectors. It is used to generate the word and token type
embeddings if word_embed and token_type_embed are None.
embed_initializer : Initializer, default None
Initializer of the embedding weights. It is used to generate the source and target
embeddings if word_embed and token_type_embed are None.
word_embed : Block or None, default None
The word embedding. If set to None, word_embed will be constructed using embed_size.
token_type_embed : Block or None, default None
The token type embedding (segment embedding). If set to None and the token_type_embed will
be constructed using embed_size.
use_pooler : bool, default True
Whether to include the pooler which converts the encoded sequence tensor of shape
(batch_size, seq_length, units) to a tensor of shape (batch_size, units)
for segment level classification task.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
use_classifier : bool, default True
Whether to include the classifier for next sentence classification.
use_token_type_embed : bool, default True
Whether to include token type embedding (segment embedding).
prefix : str or None
See document of `mx.gluon.Block`.
params : ParameterDict or None
See document of `mx.gluon.Block`.
Inputs:
- **inputs**: input sequence tensor, shape (batch_size, seq_length)
- **token_types**: optional input token type tensor, shape (batch_size, seq_length).
If the inputs contain two sequences, then the token type of the first
sequence differs from that of the second one.
- **valid_length**: optional tensor of input sequence valid lengths, shape (batch_size,)
- **masked_positions**: optional tensor of position of tokens for masked LM decoding,
shape (batch_size, num_masked_positions).
Outputs:
- **sequence_outputs**: Encoded sequence, which can be either a tensor of the last
layer of the Encoder, or a list of all sequence encodings of all layers.
In both cases shape of the tensor(s) is/are (batch_size, seq_length, units).
- **attention_outputs**: output list of all intermediate encodings per layer
Returned only if BERTEncoder.output_attention is True.
List of num_layers length of tensors of shape
(batch_size, num_attention_heads, seq_length, seq_length)
- **pooled_output**: output tensor of pooled representation of the first tokens.
Returned only if use_pooler is True. Shape (batch_size, units)
- **next_sentence_classifier_output**: output tensor of next sentence classification.
Returned only if use_classifier is True. Shape (batch_size, 2)
- **masked_lm_outputs**: output tensor of sequence decoding for masked language model
prediction. Returned only if use_decoder True.
Shape (batch_size, num_masked_positions, vocab_size)
"""
def __init__(self, encoder, vocab_size=None, token_type_vocab_size=None, units=None,
embed_size=None, embed_initializer=None,
word_embed=None, token_type_embed=None, use_pooler=True, use_decoder=True,
use_classifier=True, use_token_type_embed=True, prefix=None, params=None):
super().__init__(prefix=prefix, params=params)
self._use_decoder = use_decoder
self._use_classifier = use_classifier
self._use_pooler = use_pooler
self._use_token_type_embed = use_token_type_embed
self._units = units
self.encoder = encoder
# Construct word embedding
self.word_embed = self._get_embed(word_embed, vocab_size, embed_size,
embed_initializer, 'word_embed_')
# Construct token type embedding
if use_token_type_embed:
self.token_type_embed = self._get_embed(token_type_embed, token_type_vocab_size,
embed_size, embed_initializer,
'token_type_embed_')
if self._use_pooler:
# Construct pooler
self.pooler = self._get_pooler(units, 'pooler_')
if self._use_classifier:
# Construct classifier for next sentence predicition
self.classifier = self._get_classifier('cls_')
else:
assert not use_classifier, 'Cannot use classifier if use_pooler is False'
if self._use_decoder:
# Construct decoder for masked language model
self.decoder = self._get_decoder(units, vocab_size, self.word_embed[0], 'decoder_')
def _get_classifier(self, prefix):
""" Construct a decoder for the next sentence prediction task """
with self.name_scope():
classifier = nn.Dense(2, prefix=prefix)
return classifier
def _get_decoder(self, units, vocab_size, embed, prefix):
""" Construct a decoder for the masked language model task """
with self.name_scope():
decoder = nn.HybridSequential(prefix=prefix)
decoder.add(nn.Dense(units, flatten=False))
decoder.add(GELU())
decoder.add(nn.LayerNorm(in_channels=units, epsilon=self.encoder._layer_norm_eps))
decoder.add(nn.Dense(vocab_size, flatten=False, params=embed.collect_params()))
assert decoder[3].weight == list(embed.collect_params().values())[0], \
'The weights of word embedding are not tied with those of decoder'
return decoder
def _get_embed(self, embed, vocab_size, embed_size, initializer, prefix):
""" Construct an embedding block. """
if embed is None:
assert embed_size is not None, '"embed_size" cannot be None if "word_embed" or ' \
'token_type_embed is not given.'
with self.name_scope():
embed = nn.HybridSequential(prefix=prefix)
with embed.name_scope():
embed.add(nn.Embedding(input_dim=vocab_size, output_dim=embed_size,
weight_initializer=initializer))
assert isinstance(embed, HybridBlock)
return embed
def _get_pooler(self, units, prefix):
""" Construct pooler.
The pooler slices and projects the hidden output of first token
in the sequence for segment level classification.
"""
with self.name_scope():
pooler = nn.Dense(units=units, flatten=False, activation='tanh',
prefix=prefix)
return pooler
def __call__(self, inputs, token_types, valid_length=None, masked_positions=None):
# pylint: disable=dangerous-default-value, arguments-differ
"""Generate the representation given the inputs.
This is used in training or fine-tuning a BERT model.
"""
return super().__call__(inputs, token_types, valid_length, masked_positions)
[docs] def hybrid_forward(self, F, inputs, token_types, valid_length=None, masked_positions=None):
# pylint: disable=arguments-differ
"""Generate the representation given the inputs.
This is used in training or fine-tuning a BERT model.
"""
outputs = []
seq_out, attention_out = self._encode_sequence(inputs, token_types, valid_length)
outputs.append(seq_out)
if self.encoder._output_all_encodings:
assert isinstance(seq_out, list)
output = seq_out[-1]
else:
output = seq_out
if attention_out:
outputs.append(attention_out)
if self._use_pooler:
pooled_out = self._apply_pooling(output)
outputs.append(pooled_out)
if self._use_classifier:
next_sentence_classifier_out = self.classifier(pooled_out)
outputs.append(next_sentence_classifier_out)
if self._use_decoder:
assert masked_positions is not None, \
'masked_positions tensor is required for decoding masked language model'
decoder_out = self._decode(F, output, masked_positions)
outputs.append(decoder_out)
return tuple(outputs) if len(outputs) > 1 else outputs[0]
def _encode_sequence(self, inputs, token_types, valid_length=None):
"""Generate the representation given the input sequences.
This is used for pre-training or fine-tuning a BERT model.
"""
# embedding
embedding = self.word_embed(inputs)
if self._use_token_type_embed:
type_embedding = self.token_type_embed(token_types)
embedding = embedding + type_embedding
# (batch, seq_len, C) -> (seq_len, batch, C)
embedding = embedding.transpose((1, 0, 2))
# encoding
outputs, additional_outputs = self.encoder(embedding, valid_length=valid_length)
# (seq_len, batch, C) -> (batch, seq_len, C)
if isinstance(outputs, (list, tuple)):
outputs = [o.transpose((1, 0, 2)) for o in outputs]
else:
outputs = outputs.transpose((1, 0, 2))
return outputs, additional_outputs
def _apply_pooling(self, sequence):
"""Generate the representation given the inputs.
This is used for pre-training or fine-tuning a BERT model.
"""
outputs = sequence.slice(begin=(0, 0, 0), end=(None, 1, None))
outputs = outputs.reshape(shape=(-1, self._units))
return self.pooler(outputs)
def _decode(self, F, sequence, masked_positions):
"""Generate unnormalized prediction for the masked language model task.
This is only used for pre-training the BERT model.
Inputs:
- **sequence**: input tensor of sequence encodings.
Shape (batch_size, seq_length, units).
- **masked_positions**: input tensor of position of tokens for masked LM decoding.
Shape (batch_size, num_masked_positions). For each sample in the batch, the values
in this tensor must not be out of bound considering the length of the sequence.
Outputs:
- **masked_lm_outputs**: output tensor of token predictions for target masked_positions.
Shape (batch_size, num_masked_positions, vocab_size).
"""
masked_positions = masked_positions.astype('int32')
mask_shape = masked_positions.shape_array()
num_masked_positions = mask_shape.slice(begin=(1,), end=(2,)).astype('int32')
idx_arange = F.contrib.arange_like(masked_positions.reshape((-1, )), axis=0)
batch_idx = F.broadcast_div(idx_arange, num_masked_positions)
# batch_idx_1d = [0,0,0,1,1,1,2,2,2...]
# masked_positions_1d = [1,2,4,0,3,4,2,3,5...]
batch_idx_1d = batch_idx.reshape((1, -1))
masked_positions_1d = masked_positions.reshape((1, -1))
position_idx = F.concat(batch_idx_1d, masked_positions_1d, dim=0)
encoded = F.gather_nd(sequence, position_idx)
encoded = encoded.reshape_like(masked_positions, lhs_begin=-2, lhs_end=-1, rhs_begin=0)
decoded = self.decoder(encoded)
return decoded
[docs]class RoBERTaModel(BERTModel):
"""Generic Model for BERT (Bidirectional Encoder Representations from Transformers).
Parameters
----------
encoder : BERTEncoder
Bidirectional encoder that encodes the input sentence.
vocab_size : int or None, default None
The size of the vocabulary.
units : int or None, default None
Number of units for the final pooler layer.
embed_size : int or None, default None
Size of the embedding vectors. It is used to generate the word and token type
embeddings if word_embed and token_type_embed are None.
embed_initializer : Initializer, default None
Initializer of the embedding weights. It is used to generate the source and target
embeddings if word_embed and token_type_embed are None.
word_embed : Block or None, default None
The word embedding. If set to None, word_embed will be constructed using embed_size.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
prefix : str or None
See document of `mx.gluon.Block`.
params : ParameterDict or None
See document of `mx.gluon.Block`.
Inputs:
- **inputs**: input sequence tensor, shape (batch_size, seq_length)
- **valid_length**: optional tensor of input sequence valid lengths, shape (batch_size,)
- **masked_positions**: optional tensor of position of tokens for masked LM decoding,
shape (batch_size, num_masked_positions).
Outputs:
- **sequence_outputs**: Encoded sequence, which can be either a tensor of the last
layer of the Encoder, or a list of all sequence encodings of all layers.
In both cases shape of the tensor(s) is/are (batch_size, seq_length, units).
- **attention_outputs**: output list of all intermediate encodings per layer
Returned only if BERTEncoder.output_attention is True.
List of num_layers length of tensors of shape
(num_masks, num_attention_heads, seq_length, seq_length)
- **masked_lm_outputs**: output tensor of sequence decoding for masked language model
prediction. Returned only if use_decoder True.
Shape (batch_size, num_masked_positions, vocab_size)
"""
def __init__(self, encoder, vocab_size=None, units=None,
embed_size=None, embed_initializer=None,
word_embed=None, use_decoder=True,
prefix=None, params=None):
super(RoBERTaModel, self).__init__(encoder, vocab_size=vocab_size,
token_type_vocab_size=None, units=units,
embed_size=embed_size,
embed_initializer=embed_initializer,
word_embed=word_embed, token_type_embed=None,
use_pooler=False, use_decoder=use_decoder,
use_classifier=False, use_token_type_embed=False,
prefix=prefix, params=params)
def __call__(self, inputs, valid_length=None, masked_positions=None):
# pylint: disable=dangerous-default-value
"""Generate the representation given the inputs.
This is used in training or fine-tuning a BERT model.
"""
return super(RoBERTaModel, self).__call__(inputs, [], valid_length=valid_length,
masked_positions=masked_positions)
[docs]class DistilBERTModel(BERTModel):
"""DistilBERT Model.
Parameters
----------
encoder : BERTEncoder
Bidirectional encoder that encodes the input sentence.
vocab_size : int or None, default None
The size of the vocabulary.
units : int or None, default None
Number of units for the final pooler layer.
embed_size : int or None, default None
Size of the embedding vectors. It is used to generate the word and token type
embeddings if word_embed and token_type_embed are None.
embed_initializer : Initializer, default None
Initializer of the embedding weights. It is used to generate the source and target
embeddings if word_embed and token_type_embed are None.
word_embed : Block or None, default None
The word embedding. If set to None, word_embed will be constructed using embed_size.
prefix : str or None
See document of `mx.gluon.Block`.
params : ParameterDict or None
See document of `mx.gluon.Block`.
Inputs:
- **inputs**: input sequence tensor, shape (batch_size, seq_length)
- **valid_length**: optional tensor of input sequence valid lengths, shape (batch_size,)
Outputs:
- **sequence_outputs**: Encoded sequence, which can be either a tensor of the last
layer of the Encoder, or a list of all sequence encodings of all layers.
In both cases shape of the tensor(s) is/are (batch_size, seq_length, units).
- **attention_outputs**: output list of all intermediate encodings per layer
Returned only if BERTEncoder.output_attention is True.
List of num_layers length of tensors of shape
(num_masks, num_attention_heads, seq_length, seq_length)
"""
def __init__(self, encoder, vocab_size=None, units=None,
embed_size=None, embed_initializer=None,
word_embed=None, prefix=None, params=None):
super(DistilBERTModel, self).__init__(encoder, vocab_size=vocab_size,
token_type_vocab_size=None, units=units,
embed_size=embed_size,
embed_initializer=embed_initializer,
word_embed=word_embed, token_type_embed=None,
use_pooler=False, use_decoder=False,
use_classifier=False, use_token_type_embed=False,
prefix=prefix, params=params)
def __call__(self, inputs, valid_length=None):
# pylint: disable=dangerous-default-value, signature-differs
"""Generate the representation given the inputs.
This is used in fine-tuning a DistilBERT model.
"""
return super(DistilBERTModel, self).__call__(inputs, [], valid_length=valid_length)
[docs]class BERTClassifier(HybridBlock):
"""Model for sentence (pair) classification task with BERT.
The model feeds token ids and token type ids into BERT to get the
pooled BERT sequence representation, then apply a Dense layer for
classification.
Parameters
----------
bert: BERTModel
Bidirectional encoder with transformer.
num_classes : int, default is 2
The number of target classes.
dropout : float or None, default 0.0.
Dropout probability for the bert output.
prefix : str or None
See document of `mx.gluon.Block`.
params : ParameterDict or None
See document of `mx.gluon.Block`.
"""
def __init__(self, bert, num_classes=2, dropout=0.0,
prefix=None, params=None):
super(BERTClassifier, self).__init__(prefix=prefix, params=params)
self.bert = bert
with self.name_scope():
self.classifier = nn.HybridSequential(prefix=prefix)
if dropout:
self.classifier.add(nn.Dropout(rate=dropout))
self.classifier.add(nn.Dense(units=num_classes))
def __call__(self, inputs, token_types, valid_length=None):
# pylint: disable=dangerous-default-value, arguments-differ
"""Generate the unnormalized score for the given the input sequences.
Parameters
----------
inputs : NDArray or Symbol, shape (batch_size, seq_length)
Input words for the sequences.
token_types : NDArray or Symbol, shape (batch_size, seq_length)
Token types for the sequences, used to indicate whether the word belongs to the
first sentence or the second one.
valid_length : NDArray or Symbol, or None, shape (batch_size)
Valid length of the sequence. This is used to mask the padded tokens.
Returns
-------
outputs : NDArray or Symbol
Shape (batch_size, num_classes)
"""
return super(BERTClassifier, self).__call__(inputs, token_types, valid_length)
[docs] def hybrid_forward(self, F, inputs, token_types, valid_length=None):
# pylint: disable=arguments-differ
"""Generate the unnormalized score for the given the input sequences.
Parameters
----------
inputs : NDArray or Symbol, shape (batch_size, seq_length)
Input words for the sequences.
token_types : NDArray or Symbol, shape (batch_size, seq_length)
Token types for the sequences, used to indicate whether the word belongs to the
first sentence or the second one.
valid_length : NDArray or None, shape (batch_size)
Valid length of the sequence. This is used to mask the padded tokens.
Returns
-------
outputs : NDArray
Shape (batch_size, num_classes)
"""
_, pooler_out = self.bert(inputs, token_types, valid_length)
return self.classifier(pooler_out)
[docs]class RoBERTaClassifier(HybridBlock):
"""Model for sentence (pair) classification task with BERT.
The model feeds token ids and token type ids into BERT to get the
pooled BERT sequence representation, then apply a Dense layer for
classification.
Parameters
----------
bert: RoBERTaModel
The RoBERTa model.
num_classes : int, default is 2
The number of target classes.
dropout : float or None, default 0.0.
Dropout probability for the RoBERTa output.
prefix : str or None
See document of `mx.gluon.Block`.
params : ParameterDict or None
See document of `mx.gluon.Block`.
Inputs:
- **inputs**: input sequence tensor, shape (batch_size, seq_length)
- **valid_length**: optional tensor of input sequence valid lengths.
Shape (batch_size, num_classes).
Outputs:
- **output**: Regression output, shape (batch_size, num_classes)
"""
def __init__(self, roberta, num_classes=2, dropout=0.0,
prefix=None, params=None):
super(RoBERTaClassifier, self).__init__(prefix=prefix, params=params)
self.roberta = roberta
self._units = roberta._units
with self.name_scope():
self.classifier = nn.HybridSequential(prefix=prefix)
if dropout:
self.classifier.add(nn.Dropout(rate=dropout))
self.classifier.add(nn.Dense(units=self._units, activation='tanh'))
if dropout:
self.classifier.add(nn.Dropout(rate=dropout))
self.classifier.add(nn.Dense(units=num_classes))
def __call__(self, inputs, valid_length=None):
# pylint: disable=dangerous-default-value, arguments-differ
"""Generate the unnormalized score for the given the input sequences.
Parameters
----------
inputs : NDArray or Symbol, shape (batch_size, seq_length)
Input words for the sequences.
valid_length : NDArray or Symbol, or None, shape (batch_size)
Valid length of the sequence. This is used to mask the padded tokens.
Returns
-------
outputs : NDArray or Symbol
Shape (batch_size, num_classes)
"""
return super(RoBERTaClassifier, self).__call__(inputs, valid_length)
[docs] def hybrid_forward(self, F, inputs, valid_length=None):
# pylint: disable=arguments-differ
"""Generate the unnormalized score for the given the input sequences.
Parameters
----------
inputs : NDArray or Symbol, shape (batch_size, seq_length)
Input words for the sequences.
valid_length : NDArray or Symbol, or None, shape (batch_size)
Valid length of the sequence. This is used to mask the padded tokens.
Returns
-------
outputs : NDArray or Symbol
Shape (batch_size, num_classes)
"""
seq_out = self.roberta(inputs, valid_length)
assert not isinstance(seq_out, (tuple, list)), 'Expected one output from RoBERTaModel'
outputs = seq_out.slice(begin=(0, 0, 0), end=(None, 1, None))
outputs = outputs.reshape(shape=(-1, self._units))
return self.classifier(outputs)
###############################################################################
# GET MODEL #
###############################################################################
model_store._model_sha1.update(
{name: checksum for checksum, name in [
('5656dac6965b5054147b0375337d5a6a7a2ff832', 'bert_12_768_12_book_corpus_wiki_en_cased'),
('75cc780f085e8007b3bf6769c6348bb1ff9a3074', 'bert_12_768_12_book_corpus_wiki_en_uncased'),
('e0864cc40b3d00fcfb1a878a728650d9148c9a1d',
'distilbert_6_768_12_distilbert_book_corpus_wiki_en_uncased'),
('a56e24015a777329c795eed4ed21c698af03c9ff',
'bert_12_768_12_openwebtext_book_corpus_wiki_en_uncased'),
('5cf21fcddb5ae1a4c21c61201643460c9d65d3b0',
'roberta_12_768_12_openwebtext_ccnews_stories_books_cased'),
('d1b7163e9628e2fd51c9a9f3a0dc519d4fc24add',
'roberta_24_1024_16_openwebtext_ccnews_stories_books_cased'),
('237f39851b24f0b56d70aa20efd50095e3926e26', 'bert_12_768_12_wiki_multilingual_uncased'),
('b0f57a207f85a7d361bb79de80756a8c9a4276f7', 'bert_12_768_12_wiki_multilingual_cased'),
('885ebb9adc249a170c5576e90e88cfd1bbd98da6', 'bert_12_768_12_wiki_cn_cased'),
('4e685a966f8bf07d533bd6b0e06c04136f23f620', 'bert_24_1024_16_book_corpus_wiki_en_cased'),
('24551e1446180e045019a87fc4ffbf714d99c0b5', 'bert_24_1024_16_book_corpus_wiki_en_uncased'),
('6c82d963fc8fa79c35dd6cb3e1725d1e5b6aa7d7', 'bert_12_768_12_scibert_scivocab_uncased'),
('adf9c81e72ac286a37b9002da8df9e50a753d98b', 'bert_12_768_12_scibert_scivocab_cased'),
('75acea8e8386890120533d6c0032b0b3fcb2d536', 'bert_12_768_12_scibert_basevocab_uncased'),
('8e86e5de55d6dae99123312cd8cdd8183a75e057', 'bert_12_768_12_scibert_basevocab_cased'),
('a07780385add682f609772e81ec64aca77c9fb05', 'bert_12_768_12_biobert_v1.0_pmc_cased'),
('280ad1cc487db90489f86189e045e915b35e7489', 'bert_12_768_12_biobert_v1.0_pubmed_cased'),
('8a8c75441f028a6b928b11466f3d30f4360dfff5',
'bert_12_768_12_biobert_v1.0_pubmed_pmc_cased'),
('55f15c5d23829f6ee87622b68711b15fef50e55b', 'bert_12_768_12_biobert_v1.1_pubmed_cased'),
('60281c98ba3572dfdaac75131fa96e2136d70d5c', 'bert_12_768_12_clinicalbert_uncased'),
('f869f3f89e4237a769f1b7edcbdfe8298b480052', 'ernie_12_768_12_baidu_ernie_uncased'),
('ccf0593e03b91b73be90c191d885446df935eb64', 'bert_12_768_12_kobert_news_wiki_ko_cased')
]})
roberta_12_768_12_hparams = {
'num_layers': 12,
'units': 768,
'hidden_size': 3072,
'max_length': 512,
'num_heads': 12,
'dropout': 0.1,
'embed_size': 768,
'word_embed': None,
'layer_norm_eps': 1e-5
}
roberta_24_1024_16_hparams = {
'num_layers': 24,
'units': 1024,
'hidden_size': 4096,
'max_length': 512,
'num_heads': 16,
'dropout': 0.1,
'embed_size': 1024,
'word_embed': None,
'layer_norm_eps': 1e-5
}
distilbert_6_768_12_hparams = {
'attention_cell': 'multi_head',
'num_layers': 6,
'units': 768,
'hidden_size': 3072,
'max_length': 512,
'num_heads': 12,
'scaled': True,
'dropout': 0.1,
'use_residual': True,
'embed_size': 768,
'word_embed': None,
}
bert_12_768_12_hparams = {
'num_layers': 12,
'units': 768,
'hidden_size': 3072,
'max_length': 512,
'num_heads': 12,
'dropout': 0.1,
'embed_size': 768,
'token_type_vocab_size': 2,
'word_embed': None,
}
bert_24_1024_16_hparams = {
'num_layers': 24,
'units': 1024,
'hidden_size': 4096,
'max_length': 512,
'num_heads': 16,
'dropout': 0.1,
'embed_size': 1024,
'token_type_vocab_size': 2,
'word_embed': None,
}
ernie_12_768_12_hparams = {
'num_layers': 12,
'units': 768,
'hidden_size': 3072,
'max_length': 513,
'num_heads': 12,
'dropout': 0.1,
'embed_size': 768,
'token_type_vocab_size': 2,
'word_embed': None,
'activation': 'relu',
'layer_norm_eps': 1e-5
}
bert_hparams = {
'distilbert_6_768_12': distilbert_6_768_12_hparams,
'bert_12_768_12': bert_12_768_12_hparams,
'bert_24_1024_16': bert_24_1024_16_hparams,
'roberta_12_768_12': roberta_12_768_12_hparams,
'roberta_24_1024_16': roberta_24_1024_16_hparams,
'ernie_12_768_12': ernie_12_768_12_hparams
}
[docs]def bert_12_768_12(dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(),
root=os.path.join(get_home_dir(), 'models'), use_pooler=True, use_decoder=True,
use_classifier=True, pretrained_allow_missing=False,
hparam_allow_override=False, **kwargs):
"""Generic BERT BASE model.
The number of layers (L) is 12, number of units (H) is 768, and the
number of self-attention heads (A) is 12.
Parameters
----------
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
The supported datasets are 'book_corpus_wiki_en_cased',
'book_corpus_wiki_en_uncased', 'wiki_cn_cased',
'openwebtext_book_corpus_wiki_en_uncased',
'wiki_multilingual_uncased', 'wiki_multilingual_cased',
'scibert_scivocab_uncased', 'scibert_scivocab_cased',
'scibert_basevocab_uncased', 'scibert_basevocab_cased',
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed',
'clinicalbert',
'kobert_news_wiki_ko_cased'
vocab : gluonnlp.vocab.BERTVocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_pooler : bool, default True
Whether to include the pooler which converts the encoded sequence tensor of shape
(batch_size, seq_length, units) to a tensor of shape (batch_size, units)
for for segment level classification task.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
Note that
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed',
'clinicalbert'
do not include these parameters.
use_classifier : bool, default True
Whether to include the classifier for next sentence classification.
Note that
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed'
do not include these parameters.
pretrained_allow_missing : bool, default False
Whether to ignore if any parameters for the BERTModel are missing in
the pretrained weights for model.
Some BERTModels for example do not provide decoder or classifier
weights. In that case it is still possible to construct a BERTModel
with use_decoder=True and/or use_classifier=True, but the respective
parameters will be missing from the pretrained file.
If pretrained_allow_missing=True, this will be ignored and the
parameters will be left uninitialized. Otherwise AssertionError is
raised.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
The pretrained parameters for dataset_name
'openwebtext_book_corpus_wiki_en_uncased' were obtained by running the
GluonNLP BERT pre-training script on OpenWebText.
The pretrained parameters for dataset_name 'scibert_scivocab_uncased',
'scibert_scivocab_cased', 'scibert_basevocab_uncased',
'scibert_basevocab_cased' were obtained by converting the parameters
published by "Beltagy, I., Cohan, A., & Lo, K. (2019). Scibert: Pretrained
contextualized embeddings for scientific text. arXiv preprint
arXiv:1903.10676."
The pretrained parameters for dataset_name 'biobert_v1.0_pmc',
'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc', 'biobert_v1.1_pubmed'
were obtained by converting the parameters published by "Lee, J., Yoon, W.,
Kim, S., Kim, D., Kim, S., So, C. H., & Kang, J. (2019). Biobert:
pre-trained biomedical language representation model for biomedical text
mining. arXiv preprint arXiv:1901.08746."
The pretrained parameters for dataset_name 'clinicalbert' were obtained by
converting the parameters published by "Huang, K., Altosaar, J., &
Ranganath, R. (2019). ClinicalBERT: Modeling Clinical Notes and Predicting
Hospital Readmission. arXiv preprint arXiv:1904.05342."
Returns
-------
BERTModel, gluonnlp.vocab.BERTVocab
"""
return get_bert_model(model_name='bert_12_768_12', vocab=vocab, dataset_name=dataset_name,
pretrained=pretrained, ctx=ctx, use_pooler=use_pooler,
use_decoder=use_decoder, use_classifier=use_classifier, root=root,
pretrained_allow_missing=pretrained_allow_missing,
hparam_allow_override=hparam_allow_override, **kwargs)
[docs]def bert_24_1024_16(dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(), use_pooler=True,
use_decoder=True, use_classifier=True,
root=os.path.join(get_home_dir(), 'models'),
pretrained_allow_missing=False,
hparam_allow_override=False, **kwargs):
"""Generic BERT LARGE model.
The number of layers (L) is 24, number of units (H) is 1024, and the
number of self-attention heads (A) is 16.
Parameters
----------
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
Options include 'book_corpus_wiki_en_uncased' and 'book_corpus_wiki_en_cased'.
vocab : gluonnlp.vocab.BERTVocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_pooler : bool, default True
Whether to include the pooler which converts the encoded sequence tensor of shape
(batch_size, seq_length, units) to a tensor of shape (batch_size, units)
for for segment level classification task.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
use_classifier : bool, default True
Whether to include the classifier for next sentence classification.
pretrained_allow_missing : bool, default False
Whether to ignore if any parameters for the BERTModel are missing in
the pretrained weights for model.
Some BERTModels for example do not provide decoder or classifier
weights. In that case it is still possible to construct a BERTModel
with use_decoder=True and/or use_classifier=True, but the respective
parameters will be missing from the pretrained file.
If pretrained_allow_missing=True, this will be ignored and the
parameters will be left uninitialized. Otherwise AssertionError is
raised.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
BERTModel, gluonnlp.vocab.BERTVocab
"""
return get_bert_model(model_name='bert_24_1024_16', vocab=vocab, dataset_name=dataset_name,
pretrained=pretrained, ctx=ctx, use_pooler=use_pooler,
use_decoder=use_decoder, use_classifier=use_classifier, root=root,
pretrained_allow_missing=pretrained_allow_missing,
hparam_allow_override=hparam_allow_override, **kwargs)
[docs]def distilbert_6_768_12(dataset_name='distil_book_corpus_wiki_en_uncased', vocab=None,
pretrained=True, ctx=mx.cpu(),
output_attention=False,
output_all_encodings=False,
root=os.path.join(get_home_dir(), 'models'),
hparam_allow_override=False, **kwargs):
"""DistilBERT model: https://arxiv.org/abs/1910.01108
The number of layers (L) is 6, number of units (H) is 768, and the
number of self-attention heads (A) is 12.
Parameters
----------
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
Options include 'book_corpus_wiki_en_uncased' and 'book_corpus_wiki_en_cased'.
vocab : gluonnlp.vocab.BERTVocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
DistilBERTModel, gluonnlp.vocab.Vocab
"""
model_name = 'distilbert_6_768_12'
predefined_args = bert_hparams[model_name].copy()
if not hparam_allow_override:
mutable_args = ['use_residual', 'dropout', 'word_embed']
mutable_args = frozenset(mutable_args)
assert all((k not in kwargs or k in mutable_args) for k in predefined_args), \
'Cannot override predefined model settings.'
predefined_args.update(kwargs)
# encoder
encoder = BERTEncoder(num_layers=predefined_args['num_layers'],
units=predefined_args['units'],
hidden_size=predefined_args['hidden_size'],
max_length=predefined_args['max_length'],
num_heads=predefined_args['num_heads'],
dropout=predefined_args['dropout'],
output_attention=output_attention,
output_all_encodings=output_all_encodings,
activation=predefined_args.get('activation', 'gelu'),
layer_norm_eps=predefined_args.get('layer_norm_eps', 1e-5))
from ..vocab import Vocab # pylint: disable=import-outside-toplevel
bert_vocab = _load_vocab(dataset_name, vocab, root, cls=Vocab)
# DistilBERT
net = DistilBERTModel(encoder, len(bert_vocab),
units=predefined_args['units'],
embed_size=predefined_args['embed_size'],
word_embed=predefined_args['word_embed'])
if pretrained:
_load_pretrained_params(net, model_name, dataset_name, root, ctx,
allow_missing=False)
return net, bert_vocab
[docs]def roberta_12_768_12(dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(),
use_decoder=True,
root=os.path.join(get_home_dir(), 'models'),
hparam_allow_override=False, **kwargs):
"""Generic RoBERTa BASE model.
The number of layers (L) is 12, number of units (H) is 768, and the
number of self-attention heads (A) is 12.
Parameters
----------
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
Options include 'book_corpus_wiki_en_uncased' and 'book_corpus_wiki_en_cased'.
vocab : gluonnlp.vocab.Vocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
RoBERTaModel, gluonnlp.vocab.Vocab
"""
return get_roberta_model(model_name='roberta_12_768_12', vocab=vocab, dataset_name=dataset_name,
pretrained=pretrained, ctx=ctx,
use_decoder=use_decoder, root=root,
hparam_allow_override=hparam_allow_override, **kwargs)
[docs]def roberta_24_1024_16(dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(),
use_decoder=True,
root=os.path.join(get_home_dir(), 'models'),
hparam_allow_override=False, **kwargs):
"""Generic RoBERTa LARGE model.
The number of layers (L) is 24, number of units (H) is 1024, and the
number of self-attention heads (A) is 16.
Parameters
----------
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
Options include 'book_corpus_wiki_en_uncased' and 'book_corpus_wiki_en_cased'.
vocab : gluonnlp.vocab.Vocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
RoBERTaModel, gluonnlp.vocab.Vocab
"""
return get_roberta_model(model_name='roberta_24_1024_16', vocab=vocab,
dataset_name=dataset_name, pretrained=pretrained, ctx=ctx,
use_decoder=use_decoder, root=root,
hparam_allow_override=hparam_allow_override, **kwargs)
[docs]def ernie_12_768_12(dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(),
root=os.path.join(get_home_dir(), 'models'), use_pooler=True, use_decoder=True,
use_classifier=True, hparam_allow_override=False, **kwargs):
"""Baidu ERNIE model.
Reference:
https://arxiv.org/pdf/1904.09223.pdf
The number of layers (L) is 12, number of units (H) is 768, and the
number of self-attention heads (A) is 12.
Parameters
----------
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
The supported datasets are 'baidu_ernie'
vocab : gluonnlp.vocab.BERTVocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_pooler : bool, default True
Whether to include the pooler which converts the encoded sequence tensor of shape
(batch_size, seq_length, units) to a tensor of shape (batch_size, units)
for for segment level classification task.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
use_classifier : bool, default True
Whether to include the classifier for next sentence classification.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
(BERTModel, gluonnlp.vocab.BERTVocab)
"""
return get_bert_model(model_name='ernie_12_768_12', vocab=vocab, dataset_name=dataset_name,
pretrained=pretrained, ctx=ctx, use_pooler=use_pooler,
use_decoder=use_decoder, use_classifier=use_classifier, root=root,
pretrained_allow_missing=False,
hparam_allow_override=hparam_allow_override, **kwargs)
def get_roberta_model(model_name=None, dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(),
use_decoder=True, output_attention=False, output_all_encodings=False,
root=os.path.join(get_home_dir(), 'models'), ignore_extra=False,
hparam_allow_override=False, **kwargs):
"""Any RoBERTa pretrained model.
Parameters
----------
model_name : str or None, default None
Options include 'bert_24_1024_16' and 'bert_12_768_12'.
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
The supported datasets for model_name of either roberta_24_1024_16 and
roberta_12_768_12 include 'openwebtext_ccnews_stories_books'.
vocab : gluonnlp.vocab.Vocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
Note that
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed',
'clinicalbert'
do not include these parameters.
output_attention : bool, default False
Whether to include attention weights of each encoding cell to the output.
output_all_encodings : bool, default False
Whether to output encodings of all encoder cells.
ignore_extra : bool, default False
Whether to silently ignore parameters from the file that are not
present in this Block.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
RoBERTaModel, gluonnlp.vocab.Vocab
"""
predefined_args = bert_hparams[model_name].copy()
if not hparam_allow_override:
mutable_args = ['use_residual', 'dropout', 'word_embed']
mutable_args = frozenset(mutable_args)
assert all((k not in kwargs or k in mutable_args) for k in predefined_args), \
'Cannot override predefined model settings.'
predefined_args.update(kwargs)
# encoder
encoder = BERTEncoder(num_layers=predefined_args['num_layers'],
units=predefined_args['units'],
hidden_size=predefined_args['hidden_size'],
max_length=predefined_args['max_length'],
num_heads=predefined_args['num_heads'],
dropout=predefined_args['dropout'],
output_attention=output_attention,
output_all_encodings=output_all_encodings,
activation=predefined_args.get('activation', 'gelu'),
layer_norm_eps=predefined_args.get('layer_norm_eps', 1e-5))
from ..vocab import Vocab # pylint: disable=import-outside-toplevel
bert_vocab = _load_vocab(dataset_name, vocab, root, cls=Vocab)
# BERT
net = RoBERTaModel(encoder, len(bert_vocab),
units=predefined_args['units'],
embed_size=predefined_args['embed_size'],
word_embed=predefined_args['word_embed'],
use_decoder=use_decoder)
if pretrained:
ignore_extra = ignore_extra or not use_decoder
_load_pretrained_params(net, model_name, dataset_name, root, ctx, ignore_extra=ignore_extra,
allow_missing=False)
return net, bert_vocab
def get_bert_model(model_name=None, dataset_name=None, vocab=None, pretrained=True, ctx=mx.cpu(),
use_pooler=True, use_decoder=True, use_classifier=True, output_attention=False,
output_all_encodings=False, use_token_type_embed=True,
root=os.path.join(get_home_dir(), 'models'),
pretrained_allow_missing=False, ignore_extra=False,
hparam_allow_override=False, **kwargs):
"""Any BERT pretrained model.
Parameters
----------
model_name : str or None, default None
Options include 'bert_24_1024_16' and 'bert_12_768_12'.
dataset_name : str or None, default None
If not None, the dataset name is used to load a vocabulary for the
dataset. If the `pretrained` argument is set to True, the dataset name
is further used to select the pretrained parameters to load.
The supported datasets for model_name of either bert_24_1024_16 and
bert_12_768_12 are 'book_corpus_wiki_en_cased',
'book_corpus_wiki_en_uncased'.
For model_name bert_12_768_12 'wiki_cn_cased',
'wiki_multilingual_uncased', 'wiki_multilingual_cased',
'scibert_scivocab_uncased', 'scibert_scivocab_cased',
'scibert_basevocab_uncased','scibert_basevocab_cased',
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed',
'clinicalbert',
'kobert_news_wiki_ko_cased'
are additionally supported.
vocab : gluonnlp.vocab.BERTVocab or None, default None
Vocabulary for the dataset. Must be provided if dataset_name is not
specified. Ignored if dataset_name is specified.
pretrained : bool, default True
Whether to load the pretrained weights for model.
ctx : Context, default CPU
The context in which to load the pretrained weights.
root : str, default '$MXNET_HOME/models'
Location for keeping the model parameters.
MXNET_HOME defaults to '~/.mxnet'.
use_pooler : bool, default True
Whether to include the pooler which converts the encoded sequence tensor of shape
(batch_size, seq_length, units) to a tensor of shape (batch_size, units)
for for segment level classification task.
use_decoder : bool, default True
Whether to include the decoder for masked language model prediction.
Note that
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed',
'clinicalbert'
do not include these parameters.
use_classifier : bool, default True
Whether to include the classifier for next sentence classification.
Note that
'biobert_v1.0_pmc', 'biobert_v1.0_pubmed', 'biobert_v1.0_pubmed_pmc',
'biobert_v1.1_pubmed'
do not include these parameters.
output_attention : bool, default False
Whether to include attention weights of each encoding cell to the output.
output_all_encodings : bool, default False
Whether to output encodings of all encoder cells.
pretrained_allow_missing : bool, default False
Whether to ignore if any parameters for the BERTModel are missing in
the pretrained weights for model.
Some BERTModels for example do not provide decoder or classifier
weights. In that case it is still possible to construct a BERTModel
with use_decoder=True and/or use_classifier=True, but the respective
parameters will be missing from the pretrained file.
If pretrained_allow_missing=True, this will be ignored and the
parameters will be left uninitialized. Otherwise AssertionError is
raised.
ignore_extra : bool, default False
Whether to silently ignore parameters from the file that are not
present in this Block.
hparam_allow_override : bool, default False
If set to True, pre-defined hyper-parameters of the model
(e.g. the number of layers, hidden units) can be overriden.
Returns
-------
(BERTModel, gluonnlp.vocab.BERTVocab)
"""
predefined_args = bert_hparams[model_name].copy()
if not hparam_allow_override:
mutable_args = ['use_residual', 'dropout', 'word_embed']
mutable_args = frozenset(mutable_args)
assert all((k not in kwargs or k in mutable_args) for k in predefined_args), \
'Cannot override predefined model settings.'
predefined_args.update(kwargs)
# encoder
encoder = BERTEncoder(num_layers=predefined_args['num_layers'],
units=predefined_args['units'],
hidden_size=predefined_args['hidden_size'],
max_length=predefined_args['max_length'],
num_heads=predefined_args['num_heads'],
dropout=predefined_args['dropout'],
output_attention=output_attention,
output_all_encodings=output_all_encodings,
activation=predefined_args.get('activation', 'gelu'),
layer_norm_eps=predefined_args.get('layer_norm_eps', 1e-12))
from ..vocab import BERTVocab # pylint: disable=import-outside-toplevel
bert_vocab = _load_vocab(dataset_name, vocab, root, cls=BERTVocab)
# BERT
net = BERTModel(encoder, len(bert_vocab),
token_type_vocab_size=predefined_args['token_type_vocab_size'],
units=predefined_args['units'],
embed_size=predefined_args['embed_size'],
word_embed=predefined_args['word_embed'],
use_pooler=use_pooler, use_decoder=use_decoder,
use_classifier=use_classifier,
use_token_type_embed=use_token_type_embed)
if pretrained:
ignore_extra = ignore_extra or not (use_pooler and use_decoder and use_classifier)
_load_pretrained_params(net, model_name, dataset_name, root, ctx, ignore_extra=ignore_extra,
allow_missing=pretrained_allow_missing)
return net, bert_vocab