The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of cross-attention is added between the self-attention layers, following the architecture described in `Attention is all you need`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar
| 625 | BERT_START_DOCSTRING, |
| 626 | ) |
| 627 | class BertModel(BertPreTrainedModel): |
| 628 | """ |
| 629 | |
| 630 | The model can behave as an encoder (with only self-attention) as well |
| 631 | as a decoder, in which case a layer of cross-attention is added between |
| 632 | the self-attention layers, following the architecture described in `Attention is all you need`_ by Ashish Vaswani, |
| 633 | Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. |
| 634 | |
| 635 | To behave as an decoder the model needs to be initialized with the |
| 636 | :obj:`is_decoder` argument of the configuration set to :obj:`True`; an |
| 637 | :obj:`encoder_hidden_states` is expected as an input to the forward pass. |
| 638 | |
| 639 | .. _`Attention is all you need`: |
| 640 | https://arxiv.org/abs/1706.03762 |
| 641 | |
| 642 | """ |
| 643 | |
| 644 | def __init__(self, config): |
| 645 | super().__init__(config) |
| 646 | self.config = config |
| 647 | |
| 648 | self.embeddings = BertEmbeddings(config) |
| 649 | self.encoder = BertEncoder(config) |
| 650 | self.pooler = BertPooler(config) |
| 651 | |
| 652 | self.init_weights() |
| 653 | |
| 654 | def get_input_embeddings(self): |
| 655 | return self.embeddings.word_embeddings |
| 656 | |
| 657 | def set_input_embeddings(self, value): |
| 658 | self.embeddings.word_embeddings = value |
| 659 | |
| 660 | def _prune_heads(self, heads_to_prune): |
| 661 | """ Prunes heads of the model. |
| 662 | heads_to_prune: dict of {layer_num: list of heads to prune in this layer} |
| 663 | See base class PreTrainedModel |
| 664 | """ |
| 665 | for layer, heads in heads_to_prune.items(): |
| 666 | self.encoder.layer[layer].attention.prune_heads(heads) |
| 667 | |
| 668 | @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)")) |
| 669 | @add_code_sample_docstrings(tokenizer_class=_TOKENIZER_FOR_DOC, checkpoint="bert-base-uncased") |
| 670 | def forward( |
| 671 | self, |
| 672 | input_ids=None, |
| 673 | attention_mask=None, |
| 674 | token_type_ids=None, |
| 675 | position_ids=None, |
| 676 | head_mask=None, |
| 677 | inputs_embeds=None, |
| 678 | encoder_hidden_states=None, |
| 679 | encoder_attention_mask=None, |
| 680 | output_attentions=None, |
| 681 | output_hidden_states=None, |
| 682 | ): |
| 683 | r""" |
| 684 | Return: |