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 Vas
| 644 | |
| 645 | |
| 646 | class BertModel(BertPreTrainedModel): |
| 647 | """ |
| 648 | The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of |
| 649 | cross-attention is added between the self-attention layers, following the architecture described in `Attention is |
| 650 | all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, |
| 651 | Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. |
| 652 | argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an |
| 653 | input to the forward pass. |
| 654 | """ |
| 655 | |
| 656 | def __init__(self, config, add_pooling_layer=True): |
| 657 | super().__init__(config) |
| 658 | self.config = config |
| 659 | |
| 660 | self.embeddings = BertEmbeddings(config) |
| 661 | |
| 662 | self.encoder = BertEncoder(config) |
| 663 | |
| 664 | self.pooler = BertPooler(config) if add_pooling_layer else None |
| 665 | |
| 666 | self.init_weights() |
| 667 | |
| 668 | |
| 669 | def get_input_embeddings(self): |
| 670 | return self.embeddings.word_embeddings |
| 671 | |
| 672 | def set_input_embeddings(self, value): |
| 673 | self.embeddings.word_embeddings = value |
| 674 | |
| 675 | def _prune_heads(self, heads_to_prune): |
| 676 | """ |
| 677 | Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base |
| 678 | class PreTrainedModel |
| 679 | """ |
| 680 | for layer, heads in heads_to_prune.items(): |
| 681 | self.encoder.layer[layer].attention.prune_heads(heads) |
| 682 | |
| 683 | |
| 684 | def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor: |
| 685 | """ |
| 686 | Makes broadcastable attention and causal masks so that future and masked tokens are ignored. |
| 687 | |
| 688 | Arguments: |
| 689 | attention_mask (:obj:`torch.Tensor`): |
| 690 | Mask with ones indicating tokens to attend to, zeros for tokens to ignore. |
| 691 | input_shape (:obj:`Tuple[int]`): |
| 692 | The shape of the input to the model. |
| 693 | device: (:obj:`torch.device`): |
| 694 | The device of the input to the model. |
| 695 | |
| 696 | Returns: |
| 697 | :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`. |
| 698 | """ |
| 699 | # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] |
| 700 | # ourselves in which case we just need to make it broadcastable to all heads. |
| 701 | if attention_mask.dim() == 3: |
| 702 | extended_attention_mask = attention_mask[:, None, :, :] |
| 703 | elif attention_mask.dim() == 2: |