Styjun

Google's BERT is pretrained on next sentence prediction tasks, but I'm wondering if it's possible to call the next sentence prediction function on new data.

The idea is: given sentence A and given sentence B, I want a probabilistic label for whether or not sentence B follows sentence A. BERT is pretrained on a huge set of data, so I was hoping to use this next sentence prediction on new sentence data. I can't seem to figure out if this next sentence prediction function can be called and if so, how. Thanks for your help!

Hugging face did it for you: https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/pytorch_pretrained_bert/modeling.py#L854

class BertForNextSentencePrediction(BertPreTrainedModel):
 """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):
 super(BertForNextSentencePrediction, self).__init__(config)
 self.bert = BertModel(config)
 self.cls = BertOnlyNSPHead(config)
 self.apply(self.init_bert_weights)

 def forward(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,
 output_all_encoded_layers=False)
 seq_relationship_score = self.cls( pooled_output)

 if next_sentence_label is not None:
 loss_fct = CrossEntropyLoss(ignore_index=-1)
 next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
 return next_sentence_loss
 else:
 return seq_relationship_score