import hetu as ht import numpy as np ''' Bert Module Architecture & Input/Output Tensor Size BertModel Inputs: input_ids: [batch_size, seq_len], word token indices in the vocabulary BertModel Outputs: sequence_output: [batch_size, seq_len, hidden_size] (from BertEncoder) pooled_output: [batch_size, hidden_size] (from BertPooler) BertModel: --[batch_size, seq_len]-- BertEmbeddings: Embedding(word/position/token_type) LayerNorm Dropout --[batch_size, seq_len, hidden_size]-- --[batch_size, seq_len, hidden_size]-- BertEncoder: BertLayer(num_hidden_layers): BertAttention: BertSelfAttention --[batch_size, seq_len, hidden_size]-- BertSelfOutput: Linear Dropout Add & LayerNorm --[batch_size, seq_len, hidden_size]-- BertIntermediate: Linear + Act(gule) --[batch_size, seq_len, intermediate_size]-- BertOutput: Linear Dropout Add & LayerNorm --[batch_size, seq_len, hidden_size]-- --[batch_size, seq_len, hidden_size]-- BertPooler: (Slice, select [cls]) --[batch_size, hidden_size]-- Linear + Act(Tanh) --[batch_size, hidden_size]-- Bert ''' ''' BertEmbeddings: --------------------------------------------------------------------------------------------------''' class BertEmbeddings(object): """Construct the embeddings from word, position and token_type embeddings. """ def __init__(self, config): self.seq_len = config.max_position_embeddings self.batch_size = config.batch_size self.word_embeddings = Embedding(config.vocab_size, config.hidden_size, "word_embeddings") self.position_embeddings = Embedding(config.max_position_embeddings, config.hidden_size, 'position_embeddings') self.token_type_embeddings = Embedding(config.type_vocab_size, config.hidden_size, 'token_type_embeddings') self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12) self.dropout = Dropout(config.hidden_dropout_prob) def __call__(self, input_ids, token_type_ids): ''' inputs: input_ids: [batch_size, seq_len] token_type_ids: [batch_size, seq_len] outputs: embeddings: [batch_size, seq_len, hidden_size] ''' seq_length= self.seq_len batch_size = self.batch_size position_ids = ht.Variable('position_ids', value=np.arange(seq_length).reshape((1,-1)).repeat(batch_size,axis=0), dtype=np.long, trainable=False, ctx=input_ids.ctx) '''Embedding Size inputs_id:[batch_size, seq_len], embedding_table:[vocab_size, hidden_size] position_ids:[batch_size, seq_len], embedding_table:[seq_len, hidden_size] token_type_ids:[batch_size, seq_len], embedding_tabel:[type_vocab_size, hidden_size] --> embeddings: [batch_size, seq_len, hidden_size] ''' words_embeddings = self.word_embeddings(input_ids) position_embeddings = self.position_embeddings(position_ids) token_type_embeddings = self.token_type_embeddings(token_type_ids) embeddings = words_embeddings + position_embeddings + token_type_embeddings embeddings = self.LayerNorm(embeddings) embeddings = self.dropout(embeddings) return embeddings '''-----------------------------------------------------------------------------------------------''' ''' BertEncoder & BertLayer: --------------------------------------------------------------------------------------------------''' class BertEncoder(object): def __init__(self, config): self.output_hidden_states = config.output_hidden_states self.layer = [BertLayer(config) for _ in range(config.num_hidden_layers)] def __call__(self, hidden_states, attention_mask=None): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] attention_mask: [batch_size, num_heads, seq_len, seq_len] outputs: hidden_states: [batch_size, seq_len, hidden_size] all_hidden_states: optional, num_hidden_layers * [batch_size, seq_len, hidden_size] ''' for i, layer_module in enumerate(self.layer): hidden_states = layer_module(hidden_states, attention_mask) return hidden_states # last-layer hidden state class BertLayer(object): def __init__(self, config): self.attention = BertAttention(config) self.intermediate = BertIntermediate(config) self.output = BertOutput(config) def __call__(self, hidden_states, attention_mask): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] attention_mask: [batch_size, num_heads, seq_len, seq_len] outputs: layer_output: [batch_size, seq_len, hidden_size] ''' attention_output = self.attention(hidden_states, attention_mask) intermediate_output = self.intermediate(attention_output) layer_output = self.output(intermediate_output, attention_output) return layer_output '''-----------------------------------------------------------------------------------------------''' ''' BertAttention & BertSelfAttention & BertSelfOutput --------------------------------------------------------------------------------------------------''' class BertAttention(object): def __init__(self, config): self.self = BertSelfAttention(config) self.output = BertSelfOutput(config) def __call__(self, input_tensor, attention_mask): ''' inputs: input_tensor: [batch_size, seq_len, hidden_size] attention_mask: [batch_size, num_heads, seq_len, seq_len] outputs: attention_output: [batch_size, seq_len, hidden_size] ''' self_output = self.self(input_tensor, attention_mask) attention_output = self.output(self_output, input_tensor) return attention_output class BertSelfAttention(object): def __init__(self, config): if config.hidden_size % config.num_attention_heads != 0: raise ValueError( "The hidden size (%d) is not a multiple of the number of attention " "heads (%d)" % (config.hidden_size, config.num_attention_heads)) self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.all_head_size = self.num_attention_heads * self.attention_head_size #all_head_size == hidden_size self.hidden_size = config.hidden_size self.seq_len = config.max_position_embeddings self.batch_size = config.batch_size linear_input_shape = [self.batch_size, self.seq_len, self.hidden_size] self.query = Linear(config.hidden_size, self.all_head_size, input_shape=linear_input_shape) self.key = Linear(config.hidden_size, self.all_head_size, input_shape=linear_input_shape) self.value = Linear(config.hidden_size, self.all_head_size, input_shape=linear_input_shape) self.dropout = Dropout(config.attention_probs_dropout_prob) def transpose_for_scores(self, input_tensor): output_tensor = ht.array_reshape_op( input_tensor, [self.batch_size, self.seq_len, self.num_attention_heads, self.attention_head_size]) output_tensor = ht.transpose_op(output_tensor, [0, 2, 1, 3]) return output_tensor def __call__(self, hidden_states, attention_mask): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] attention_mask: [batch_size, 1, 1, seq_len] outputs: context_layer: [batch_size, seq_len, hidden_size] ''' # linear transformation mixed_query_layer = self.query(hidden_states) # [batch_size, seq_len, hidden_size] mixed_key_layer = self.key(hidden_states) # [batch_size, seq_len, hidden_size] mixed_value_layer = self.value(hidden_states) # [batch_size, seq_len, hidden_size] # transpose query_layer = self.transpose_for_scores(mixed_query_layer) # [batch_size, num_heads, seq_len, head_size] key_layer = self.transpose_for_scores(mixed_key_layer) # [batch_size, num_heads, seq_len, head_size] value_layer = self.transpose_for_scores(mixed_value_layer) # [batch_size, num_heads, seq_len, head_size] # score key_layer_scaled = key_layer * (1.0 / np.sqrt(float(self.attention_head_size))) attention_scores = ht.batch_matmul_op(query_layer, key_layer_scaled, trans_B=True) # [batch_size, num_heads, seq_len, seq_len] # Apply the attention mask is (precomputed for all layers in BertModel forward() function) attention_scores = attention_scores + ht.broadcastto_op(attention_mask, attention_scores) # [batch_size, num_heads, seq_len, seq_len] # Normalize the attention scores to probabilities. attention_probs = ht.softmax_op(attention_scores) # This is actually dropping out entire tokens to attend to, which might # seem a bit unusual, but is taken from the original Transformer paper. attention_probs = self.dropout(attention_probs) context_layer = ht.batch_matmul_op(attention_probs, value_layer) # [batch_size, num_heads, seq_len, head_size] context_layer = ht.transpose_op(context_layer, [0, 2, 1, 3]) # [batch_size, seq_len, num_heads, head_size] context_layer = ht.array_reshape_op(context_layer, [-1, self.seq_len, self.all_head_size]) # [batch_size, seq_len, hidden_size] return context_layer class BertSelfOutput(object): def __init__(self, config): linear_input_shape = [config.batch_size, config.max_position_embeddings, config.hidden_size] self.dense = Linear(config.hidden_size, config.hidden_size, input_shape=linear_input_shape) self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12) self.dropout = Dropout(config.hidden_dropout_prob) def __call__(self, hidden_states, input_tensor): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] input_tensor: [batch_size, seq_len, hidden_size] outputs: hidden_states: [batch_size, seq_len, hidden_size] ''' hidden_states = self.dense(hidden_states) hidden_states = self.dropout(hidden_states) hidden_states = self.LayerNorm(hidden_states + input_tensor) return hidden_states '''-----------------------------------------------------------------------------------------------''' ''' BertIntermediate & BertOutput （2-layer FeedForward) --------------------------------------------------------------------------------------------------''' class BertIntermediate(object): def __init__(self, config): if config.hidden_act == "relu": self.intermediate_act_fn = ht.relu_op elif config.hidden_act == "gelu": self.intermediate_act_fn = ht.gelu_op print("Gelu activation is not implemented yet.") assert(False) linear_input_shape = [config.batch_size, config.max_position_embeddings, config.hidden_size] self.dense = Linear(config.hidden_size, config.intermediate_size, activation = self.intermediate_act_fn, input_shape=linear_input_shape) def __call__(self, hidden_states): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] outputs: hidden_states: [batch_size, seq_len, intermediate_size] ''' hidden_states = self.dense(hidden_states) return hidden_states class BertOutput(object): def __init__(self, config): linear_input_shape = [config.batch_size, config.max_position_embeddings, config.intermediate_size] self.dense = Linear(config.intermediate_size, config.hidden_size, input_shape=linear_input_shape) self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12) self.dropout = Dropout(config.hidden_dropout_prob) def __call__(self, hidden_states, input_tensor): ''' inputs: hidden_states: [batch_size, seq_len, intermediate_size] outputs: hidden_states: [batch_size, seq_len, hidden_size] ''' hidden_states = self.dense(hidden_states) hidden_states = self.dropout(hidden_states) hidden_states = self.LayerNorm(hidden_states + input_tensor) return hidden_states '''-----------------------------------------------------------------------------------------------''' ''' BertPooler --------------------------------------------------------------------------------------------------''' class BertPooler(object): def __init__(self, config): self.dense = Linear(config.hidden_size, config.hidden_size, activation = ht.tanh_op) self.batch_size = config.batch_size self.hidden_size = config.hidden_size def __call__(self, hidden_states): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] outputs: pooled_output: [batch_size, hidden_size] ''' first_token_tensor = ht.slice_op(hidden_states,(0,0,0),(self.batch_size,1,self.hidden_size)) first_token_tensor = ht.array_reshape_op(first_token_tensor, [self.batch_size, self.hidden_size]) pooled_output = self.dense(first_token_tensor) return pooled_output '''-----------------------------------------------------------------------------------------------''' ''' Bert Downstream Heads --------------------------------------------------------------------------------------------------''' class BertPredictionHeadTransform(object): def __init__(self, config): if config.hidden_act == "relu": self.hidden_act = ht.relu_op elif config.hidden_act == "gelu": self.hidden_act = ht.gelu_op print("Gelu activation is not implemented yet.") assert(False) linear_input_shape = [config.batch_size, config.max_position_embeddings, config.hidden_size] self.dense_act = Linear(config.hidden_size, config.hidden_size, activation=self.hidden_act, input_shape=linear_input_shape) self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12) def __call__(self, hidden_states): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] outputs: hidden_states: [batch_size, seq_len, hidden_size] ''' hidden_states = self.dense_act(hidden_states) hidden_states = self.LayerNorm(hidden_states) return hidden_states class BertLMPredictionHead(object): def __init__(self, config, bert_model_embedding_weights): ''' bert_model_embedding_weights: [vocab_size, hidden_size] ''' self.transform = BertPredictionHeadTransform(config) linear_input_shape = [config.batch_size, config.max_position_embeddings, config.hidden_size] self.decoder = Linear(config.hidden_size, config.vocab_size, bias_initializer=ht.init.zeros,input_shape=linear_input_shape) self.decoder.weights = ht.transpose_op(bert_model_embedding_weights) def __call__(self, hidden_states): ''' inputs: hidden_states: [batch_size, seq_len, hidden_size] outputs: hidden_states: [batch_size, seq_len, vocab_size] ''' hidden_states = self.transform(hidden_states) hidden_states = self.decoder(hidden_states) return hidden_states class BertOnlyMLMHead(object): def __init__(self, config, bert_model_embedding_weights): self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights) def __call__(self, sequence_output): ''' inputs: sequence_output: [batch_size, seq_len, hidden_size] outputs: prediction_scores: [batch_size, seq_len, vocab_size] ''' prediction_scores = self.predictions(sequence_output) return prediction_scores class BertOnlyNSPHead(object): def __init__(self, config): self.seq_relationship = Linear(config.hidden_size, 2) def __call__(self, pooled_output): ''' inputs: pooled_output: [batch_size, hidden_size] outputs: seq_relationship_score: [batch_size, 2] ''' seq_relationship_score = self.seq_relationship(pooled_output) return seq_relationship_score class BertPreTrainingHeads(object): def __init__(self, config, bert_model_embedding_weights): self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights) self.seq_relationship = Linear(config.hidden_size, 2) def __call__(self, sequence_output, pooled_output): ''' inputs: sequence_output: [batch_size, seq_len, hidden_size] pooled_output: [batch_size, hidden_size] outputs: prediction_scores: [batch_size, seq_len, vocab_size] seq_relationship_score: [batch_size, 2] ''' prediction_scores = self.predictions(sequence_output) seq_relationship_score = self.seq_relationship(pooled_output) return prediction_scores, seq_relationship_score '''-----------------------------------------------------------------------------------------------''' ''' BertModel: --------------------------------------------------------------------------------------------------''' class BertModel(object): """BERT model ("Bidirectional Embedding Representations from a Transformer"). Params: config: a BertConfig class instance with the configuration to build a new model Inputs: `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. Outputs: Tuple of (encoded_layers, pooled_output) `encoded_layers`: controled by `output_all_encoded_layers` argument: - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size], - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding to the last attention block of shape [batch_size, sequence_length, hidden_size], `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a classifier pretrained on top of the hidden state associated to the first character of the input (`CLS`) to train on the Next-Sentence task (see BERT's paper). Example usage: ```python # Already been converted into WordPiece token ids input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]]) input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]]) token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]]) config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072) model = modeling.BertModel(config=config) all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask) ``` """ def __init__(self, config): self.embeddings = BertEmbeddings(config) self.encoder = BertEncoder(config) self.pooler = BertPooler(config) self.batch_size=config.batch_size self.seq_len=config.max_position_embeddings def __call__(self, input_ids, token_type_ids, attention_mask): extended_attention_mask = ht.array_reshape_op(attention_mask, [self.batch_size, 1, 1, self.seq_len]) extended_attention_mask = (extended_attention_mask+(-1.0)) * 10000.0 embedding_output = self.embeddings(input_ids, token_type_ids) sequence_output = self.encoder(embedding_output, extended_attention_mask) pooled_output = self.pooler(sequence_output) return sequence_output, pooled_output '''-----------------------------------------------------------------------------------------------''' ''' BertForPreTraining: --------------------------------------------------------------------------------------------------''' class BertForPreTraining(object): """BERT model with pre-training heads. This module comprises the BERT model followed by the two pre-training heads: - the masked language modeling head, and - the next sentence classification head. Params: config: a BertConfig class instance with the configuration to build a new model. Inputs: `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. `masked_lm_labels`: optional masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss is only computed for the labels set in [0, ..., vocab_size] `next_sentence_label`: optional next sentence classification loss: torch.LongTensor of shape [batch_size] with indices selected in [0, 1]. 0 => next sentence is the continuation, 1 => next sentence is a random sentence. Outputs: if `masked_lm_labels` and `next_sentence_label` are not `None`: Outputs the total_loss which is the sum of the masked language modeling loss and the next sentence classification loss. if `masked_lm_labels` or `next_sentence_label` is `None`: Outputs a tuple comprising - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and - the next sentence classification logits of shape [batch_size, 2]. Example usage: ```python # Already been converted into WordPiece token ids input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]]) input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]]) token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]]) config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072) model = BertForPreTraining(config) masked_lm_logits_scores, seq_relationship_logits = model(input_ids, token_type_ids, input_mask) ``` """ def __init__(self, config): self.bert = BertModel(config) self.cls = BertPreTrainingHeads(config, self.bert.embeddings.word_embeddings.weight) self.vocab_size=config.vocab_size def __call__(self, input_ids, token_type_ids, attention_mask, masked_lm_labels=None, next_sentence_label=None): sequence_output, pooled_output = self.bert(input_ids, token_type_ids, attention_mask) prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) return_op = [prediction_scores, seq_relationship_score] if masked_lm_labels is not None and next_sentence_label is not None: ''' masked_lm_labels: [batch_size, seq_len, vocab_size], one hot form, masked places are filled with 0 prediction_scores: [batch_size, seq_len, vocab_size] next_sentence_label: [batch_size, 2], one hot form, masked places are filled with 0 seq_relationship_score: [batch_size, 2] masked_lm_loss: [batch_size*seq_len] next_sentence_loss: [batch_size] ''' masked_lm_loss = ht.softmaxcrossentropy_sparse_op(prediction_scores, masked_lm_labels, ignored_index=-1) next_sentence_loss = ht.softmaxcrossentropy_sparse_op(seq_relationship_score, next_sentence_label, ignored_index=-1) return_op += [masked_lm_loss, next_sentence_loss] return return_op class BertForMaskedLM(object): """BERT model with the masked language modeling head. This module comprises the BERT model followed by the masked language modeling head. Params: config: a BertConfig class instance with the configuration to build a new model. Inputs: `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. `masked_lm_labels`: masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss is only computed for the labels set in [0, ..., vocab_size] Outputs: if `masked_lm_labels` is not `None`: Outputs the masked language modeling loss. if `masked_lm_labels` is `None`: Outputs the masked language modeling logits of shape [batch_size, sequence_length, vocab_size]. Example usage: ```python # Already been converted into WordPiece token ids input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]]) input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]]) token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]]) config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072) model = BertForMaskedLM(config) masked_lm_logits_scores = model(input_ids, token_type_ids, input_mask) ``` """ def __init__(self, config): self.bert = BertModel(config) self.cls = BertOnlyMLMHead(config, self.bert.embeddings.word_embeddings.weight) self.vocab_size=config.vocab_size def __call__(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None): sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask) prediction_scores = self.cls(sequence_output) return_op = [prediction_scores] if masked_lm_labels is not None: ''' masked_lm_labels: [batch_size, seq_len, vocab_size], one hot form, masked places are filled with 0 prediction_scores: [batch_size, seq_len, vocab_size] masked_lm_loss: [batch_size*seq_len] ''' masked_lm_loss = ht.softmaxcrossentropy_sparse_op(prediction_scores, masked_lm_labels, ignored_index=-1) return_op += [masked_lm_loss] return return_op class BertForNextSentencePrediction(object): """BERT model with next sentence prediction head. This module comprises the BERT model followed by the next sentence classification head. Params: config: a BertConfig class instance with the configuration to build a new model. Inputs: `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. `next_sentence_label`: next sentence classification loss: torch.LongTensor of shape [batch_size] with indices selected in [0, 1]. 0 => next sentence is the continuation, 1 => next sentence is a random sentence. Outputs: if `next_sentence_label` is not `None`: Outputs the total_loss which is the sum of the masked language modeling loss and the next sentence classification loss. if `next_sentence_label` is `None`: Outputs the next sentence classification logits of shape [batch_size, 2]. Example usage: ```python # Already been converted into WordPiece token ids input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]]) input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]]) token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]]) config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072) model = BertForNextSentencePrediction(config) seq_relationship_logits = model(input_ids, token_type_ids, input_mask) ``` """ def __init__(self, config): self.bert = BertModel(config) self.cls = BertOnlyNSPHead(config) def __call__(self, input_ids, token_type_ids=None, attention_mask=None, next_sentence_label=None): _, pooled_output = self.bert(input_ids, token_type_ids, attention_mask) seq_relationship_score = self.cls(pooled_output) return_op = [seq_relationship_score] if next_sentence_label is not None: ''' next_sentence_label: [batch_size, 2], one hot form, masked places are filled with 0 seq_relationship_score: [batch_size, 2] next_sentence_loss: [batch_size] ''' next_sentence_loss = ht.softmaxcrossentropy_sparse_op(seq_relationship_score, next_sentence_label, ignored_index=-1) return_op += [next_sentence_loss] return return_op '''-----------------------------------------------------------------------------------------------''' ''' Bert Layer utils (Embedding & BerLayerNorm & Dropout & Linear) --------------------------------------------------------------------------------------------------''' class Embedding(object): def __init__(self, num_embeddings, embedding_dim, embedding_name=None, initializer=ht.init.xavier_normal): self.weight = initializer(name=embedding_name, shape=(num_embeddings, embedding_dim)) def __call__(self, input_tensor): return ht.embedding_lookup_op(self.weight, input_tensor) class BertLayerNorm(object): def __init__(self, hidden_size, eps=1e-12): self.eps=eps self.scale = ht.init.ones(name='layer_norm_scale', shape=(hidden_size, )) self.bias = ht.init.zeros(name='layer_norm_bias', shape=(hidden_size, )) def __call__(self, input_tensor): return ht.layer_normalization_op(input_tensor, self.scale, self.bias, eps=self.eps) class Dropout(object): def __init__(self, dropout_prob=None): self.dropout_prob = dropout_prob def __call__(self, input_tensor): if self.dropout_prob is None or self.dropout_prob == 0.0: return input_tensor output = ht.dropout_op(input_tensor, 1.0 - self.dropout_prob) return output class Linear(object): def __init__(self, in_features, out_features, bias=True, activation=None, kernel_initializer=ht.init.xavier_normal, bias_initializer=ht.init.zeros, input_shape=None): self.bias_flag = bias self.activation = activation self.weights = kernel_initializer(name='dense_weights', shape=(in_features, out_features)) if self.bias_flag: self.bias = bias_initializer(name='dense_bias', shape=(out_features,)) self.input_shape=input_shape self.in_features = in_features self.out_features = out_features if self.input_shape is not None and self.input_shape[-1]!=in_features: print("Specified in_features is not equal to input_shape[-1].") assert(False) def __call__(self, input_tensor): if self.input_shape is not None and len(self.input_shape)!=2: input_tensor = ht.array_reshape_op(input_tensor, [-1, self.in_features]) outputs = ht.matmul_op(input_tensor, self.weights) if self.bias_flag: outputs = outputs + ht.broadcastto_op(self.bias, outputs) if self.activation is not None: outputs = self.activation(outputs) if self.input_shape is not None and len(self.input_shape)!=2: outputs = ht.array_reshape_op(outputs, self.input_shape[:-1]+[self.out_features]) return outputs '''-----------------------------------------------------------------------------------------------'''