|
19 | 19 | from .. import PretrainedModel, register_base_model
|
20 | 20 |
|
21 | 21 | __all__ = [
|
22 |
| - 'ErnieModel', 'ErniePretrainedModel', 'ErnieForSequenceClassification', |
23 |
| - 'ErnieForTokenClassification', 'ErnieForQuestionAnswering', |
24 |
| - 'ErnieForPretraining', 'ErniePretrainingCriterion' |
| 22 | + 'ErnieModel', |
| 23 | + 'ErniePretrainedModel', |
| 24 | + 'ErnieForSequenceClassification', |
| 25 | + 'ErnieForTokenClassification', |
| 26 | + 'ErnieForQuestionAnswering', |
| 27 | + 'ErnieForPretraining', |
| 28 | + 'ErniePretrainingCriterion', |
| 29 | + 'ErnieForMaskedLM', |
25 | 30 | ]
|
26 | 31 |
|
27 | 32 |
|
@@ -770,3 +775,87 @@ def forward(self, prediction_scores, seq_relationship_score,
|
770 | 775 | next_sentence_loss = F.cross_entropy(
|
771 | 776 | seq_relationship_score, next_sentence_labels, reduction='none')
|
772 | 777 | return paddle.mean(masked_lm_loss), paddle.mean(next_sentence_loss)
|
| 778 | + |
| 779 | + |
| 780 | +class ErnieOnlyMLMHead(nn.Layer): |
| 781 | + def __init__(self, hidden_size, vocab_size, activation, embedding_weights): |
| 782 | + super().__init__() |
| 783 | + self.predictions = ErnieLMPredictionHead( |
| 784 | + hidden_size=hidden_size, |
| 785 | + vocab_size=vocab_size, |
| 786 | + activation=activation, |
| 787 | + embedding_weights=embedding_weights) |
| 788 | + |
| 789 | + def forward(self, sequence_output, masked_positions=None): |
| 790 | + prediction_scores = self.predictions(sequence_output, masked_positions) |
| 791 | + return prediction_scores |
| 792 | + |
| 793 | + |
| 794 | +class ErnieForMaskedLM(ErniePretrainedModel): |
| 795 | + """ |
| 796 | + Ernie Model with a `masked language modeling` head on top. |
| 797 | +
|
| 798 | + Args: |
| 799 | + ernie (:class:`ErnieModel`): |
| 800 | + An instance of :class:`ErnieModel`. |
| 801 | +
|
| 802 | + """ |
| 803 | + |
| 804 | + def __init__(self, ernie): |
| 805 | + super(ErnieForMaskedLM, self).__init__() |
| 806 | + self.ernie = ernie |
| 807 | + self.cls = ErnieOnlyMLMHead( |
| 808 | + self.ernie.config["hidden_size"], |
| 809 | + self.ernie.config["vocab_size"], |
| 810 | + self.ernie.config["hidden_act"], |
| 811 | + embedding_weights=self.ernie.embeddings.word_embeddings.weight) |
| 812 | + |
| 813 | + self.apply(self.init_weights) |
| 814 | + |
| 815 | + def forward(self, |
| 816 | + input_ids, |
| 817 | + token_type_ids=None, |
| 818 | + position_ids=None, |
| 819 | + attention_mask=None): |
| 820 | + r""" |
| 821 | +
|
| 822 | + Args: |
| 823 | + input_ids (Tensor): |
| 824 | + See :class:`ErnieModel`. |
| 825 | + token_type_ids (Tensor, optional): |
| 826 | + See :class:`ErnieModel`. |
| 827 | + position_ids (Tensor, optional): |
| 828 | + See :class:`ErnieModel`. |
| 829 | + attention_mask (Tensor, optional): |
| 830 | + See :class:`ErnieModel`. |
| 831 | +
|
| 832 | + Returns: |
| 833 | + Tensor: Returns tensor `prediction_scores`, The scores of masked token prediction. |
| 834 | + Its data type should be float32 and shape is [batch_size, sequence_length, vocab_size]. |
| 835 | +
|
| 836 | + Example: |
| 837 | + .. code-block:: |
| 838 | +
|
| 839 | + import paddle |
| 840 | + from paddlenlp.transformers import ErnieForMaskedLM, ErnieTokenizer |
| 841 | +
|
| 842 | + tokenizer = ErnieTokenizer.from_pretrained('ernie-1.0') |
| 843 | + model = ErnieForMaskedLM.from_pretrained('ernie-1.0') |
| 844 | + |
| 845 | + inputs = tokenizer("Welcome to use PaddlePaddle and PaddleNLP!") |
| 846 | + inputs = {k:paddle.to_tensor([v]) for (k, v) in inputs.items()} |
| 847 | +
|
| 848 | + logits = model(**inputs) |
| 849 | + print(logits.shape) |
| 850 | + # [1, 17, 18000] |
| 851 | +
|
| 852 | + """ |
| 853 | + |
| 854 | + outputs = self.ernie( |
| 855 | + input_ids, |
| 856 | + token_type_ids=token_type_ids, |
| 857 | + position_ids=position_ids, |
| 858 | + attention_mask=attention_mask) |
| 859 | + sequence_output = outputs[0] |
| 860 | + prediction_scores = self.cls(sequence_output, masked_positions=None) |
| 861 | + return prediction_scores |
0 commit comments