CoolFace
Modelpublic

k4tel/bert-multilingial-geolocation-prediction

sourceHugging Faceupdated 4y agoView on Hugging Face
0likes17downloads
modeling_bert.py1890 linesDownload Raw Back to root
1# coding=utf-82# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9#     http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16"""PyTorch BERT model."""17 18 19import math20import os21import warnings22from dataclasses import dataclass23from typing import List, Optional, Tuple, Union24 25import torch26import torch.utils.checkpoint27from torch import nn28from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss29 30from ...activations import ACT2FN31from ...modeling_outputs import (32    BaseModelOutputWithPastAndCrossAttentions,33    BaseModelOutputWithPoolingAndCrossAttentions,34    CausalLMOutputWithCrossAttentions,35    MaskedLMOutput,36    MultipleChoiceModelOutput,37    NextSentencePredictorOutput,38    QuestionAnsweringModelOutput,39    SequenceClassifierOutput,40    TokenClassifierOutput,41)42from ...modeling_utils import PreTrainedModel43from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer44from ...utils import (45    ModelOutput,46    add_code_sample_docstrings,47    add_start_docstrings,48    add_start_docstrings_to_model_forward,49    logging,50    replace_return_docstrings,51)52from .configuration_bert import BertConfig53 54 55logger = logging.get_logger(__name__)56 57_CHECKPOINT_FOR_DOC = "bert-base-uncased"58_CONFIG_FOR_DOC = "BertConfig"59_TOKENIZER_FOR_DOC = "BertTokenizer"60 61# TokenClassification docstring62_CHECKPOINT_FOR_TOKEN_CLASSIFICATION = "dbmdz/bert-large-cased-finetuned-conll03-english"63_TOKEN_CLASS_EXPECTED_OUTPUT = (64    "['O', 'I-ORG', 'I-ORG', 'I-ORG', 'O', 'O', 'O', 'O', 'O', 'I-LOC', 'O', 'I-LOC', 'I-LOC'] "65)66_TOKEN_CLASS_EXPECTED_LOSS = 0.0167 68# QuestionAnswering docstring69_CHECKPOINT_FOR_QA = "deepset/bert-base-cased-squad2"70_QA_EXPECTED_OUTPUT = "'a nice puppet'"71_QA_EXPECTED_LOSS = 7.4172_QA_TARGET_START_INDEX = 1473_QA_TARGET_END_INDEX = 1574 75# SequenceClassification docstring76_CHECKPOINT_FOR_SEQUENCE_CLASSIFICATION = "textattack/bert-base-uncased-yelp-polarity"77_SEQ_CLASS_EXPECTED_OUTPUT = "'LABEL_1'"78_SEQ_CLASS_EXPECTED_LOSS = 0.0179 80 81BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [82    "bert-base-uncased",83    "bert-large-uncased",84    "bert-base-cased",85    "bert-large-cased",86    "bert-base-multilingual-uncased",87    "bert-base-multilingual-cased",88    "bert-base-chinese",89    "bert-base-german-cased",90    "bert-large-uncased-whole-word-masking",91    "bert-large-cased-whole-word-masking",92    "bert-large-uncased-whole-word-masking-finetuned-squad",93    "bert-large-cased-whole-word-masking-finetuned-squad",94    "bert-base-cased-finetuned-mrpc",95    "bert-base-german-dbmdz-cased",96    "bert-base-german-dbmdz-uncased",97    "cl-tohoku/bert-base-japanese",98    "cl-tohoku/bert-base-japanese-whole-word-masking",99    "cl-tohoku/bert-base-japanese-char",100    "cl-tohoku/bert-base-japanese-char-whole-word-masking",101    "TurkuNLP/bert-base-finnish-cased-v1",102    "TurkuNLP/bert-base-finnish-uncased-v1",103    "wietsedv/bert-base-dutch-cased",104    # See all BERT models at https://huggingface.co/models?filter=bert105]106 107 108def load_tf_weights_in_bert(model, config, tf_checkpoint_path):109    """Load tf checkpoints in a pytorch model."""110    try:111        import re112 113        import numpy as np114        import tensorflow as tf115    except ImportError:116        logger.error(117            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "118            "https://www.tensorflow.org/install/ for installation instructions."119        )120        raise121    tf_path = os.path.abspath(tf_checkpoint_path)122    logger.info(f"Converting TensorFlow checkpoint from {tf_path}")123    # Load weights from TF model124    init_vars = tf.train.list_variables(tf_path)125    names = []126    arrays = []127    for name, shape in init_vars:128        logger.info(f"Loading TF weight {name} with shape {shape}")129        array = tf.train.load_variable(tf_path, name)130        names.append(name)131        arrays.append(array)132 133    for name, array in zip(names, arrays):134        name = name.split("/")135        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v136        # which are not required for using pretrained model137        if any(138            n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]139            for n in name140        ):141            logger.info(f"Skipping {'/'.join(name)}")142            continue143        pointer = model144        for m_name in name:145            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):146                scope_names = re.split(r"_(\d+)", m_name)147            else:148                scope_names = [m_name]149            if scope_names[0] == "kernel" or scope_names[0] == "gamma":150                pointer = getattr(pointer, "weight")151            elif scope_names[0] == "output_bias" or scope_names[0] == "beta":152                pointer = getattr(pointer, "bias")153            elif scope_names[0] == "output_weights":154                pointer = getattr(pointer, "weight")155            elif scope_names[0] == "squad":156                pointer = getattr(pointer, "classifier")157            else:158                try:159                    pointer = getattr(pointer, scope_names[0])160                except AttributeError:161                    logger.info(f"Skipping {'/'.join(name)}")162                    continue163            if len(scope_names) >= 2:164                num = int(scope_names[1])165                pointer = pointer[num]166        if m_name[-11:] == "_embeddings":167            pointer = getattr(pointer, "weight")168        elif m_name == "kernel":169            array = np.transpose(array)170        try:171            if pointer.shape != array.shape:172                raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")173        except AssertionError as e:174            e.args += (pointer.shape, array.shape)175            raise176        logger.info(f"Initialize PyTorch weight {name}")177        pointer.data = torch.from_numpy(array)178    return model179 180 181class BertEmbeddings(nn.Module):182    """Construct the embeddings from word, position and token_type embeddings."""183 184    def __init__(self, config):185        super().__init__()186        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)187        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)188        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)189 190        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load191        # any TensorFlow checkpoint file192        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)193        self.dropout = nn.Dropout(config.hidden_dropout_prob)194        # position_ids (1, len position emb) is contiguous in memory and exported when serialized195        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")196        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))197        self.register_buffer(198            "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False199        )200 201    def forward(202        self,203        input_ids: Optional[torch.LongTensor] = None,204        token_type_ids: Optional[torch.LongTensor] = None,205        position_ids: Optional[torch.LongTensor] = None,206        inputs_embeds: Optional[torch.FloatTensor] = None,207        past_key_values_length: int = 0,208    ) -> torch.Tensor:209        if input_ids is not None:210            input_shape = input_ids.size()211        else:212            input_shape = inputs_embeds.size()[:-1]213 214        seq_length = input_shape[1]215 216        if position_ids is None:217            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]218 219        # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs220        # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves221        # issue #5664222        if token_type_ids is None:223            if hasattr(self, "token_type_ids"):224                buffered_token_type_ids = self.token_type_ids[:, :seq_length]225                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)226                token_type_ids = buffered_token_type_ids_expanded227            else:228                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)229 230        if inputs_embeds is None:231            inputs_embeds = self.word_embeddings(input_ids)232        token_type_embeddings = self.token_type_embeddings(token_type_ids)233 234        embeddings = inputs_embeds + token_type_embeddings235        if self.position_embedding_type == "absolute":236            position_embeddings = self.position_embeddings(position_ids)237            embeddings += position_embeddings238        embeddings = self.LayerNorm(embeddings)239        embeddings = self.dropout(embeddings)240        return embeddings241 242 243class BertSelfAttention(nn.Module):244    def __init__(self, config, position_embedding_type=None):245        super().__init__()246        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):247            raise ValueError(248                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "249                f"heads ({config.num_attention_heads})"250            )251 252        self.num_attention_heads = config.num_attention_heads253        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)254        self.all_head_size = self.num_attention_heads * self.attention_head_size255 256        self.query = nn.Linear(config.hidden_size, self.all_head_size)257        self.key = nn.Linear(config.hidden_size, self.all_head_size)258        self.value = nn.Linear(config.hidden_size, self.all_head_size)259 260        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)261        self.position_embedding_type = position_embedding_type or getattr(262            config, "position_embedding_type", "absolute"263        )264        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":265            self.max_position_embeddings = config.max_position_embeddings266            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)267 268        self.is_decoder = config.is_decoder269 270    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:271        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)272        x = x.view(new_x_shape)273        return x.permute(0, 2, 1, 3)274 275    def forward(276        self,277        hidden_states: torch.Tensor,278        attention_mask: Optional[torch.FloatTensor] = None,279        head_mask: Optional[torch.FloatTensor] = None,280        encoder_hidden_states: Optional[torch.FloatTensor] = None,281        encoder_attention_mask: Optional[torch.FloatTensor] = None,282        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,283        output_attentions: Optional[bool] = False,284    ) -> Tuple[torch.Tensor]:285        mixed_query_layer = self.query(hidden_states)286 287        # If this is instantiated as a cross-attention module, the keys288        # and values come from an encoder; the attention mask needs to be289        # such that the encoder's padding tokens are not attended to.290        is_cross_attention = encoder_hidden_states is not None291 292        if is_cross_attention and past_key_value is not None:293            # reuse k,v, cross_attentions294            key_layer = past_key_value[0]295            value_layer = past_key_value[1]296            attention_mask = encoder_attention_mask297        elif is_cross_attention:298            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))299            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))300            attention_mask = encoder_attention_mask301        elif past_key_value is not None:302            key_layer = self.transpose_for_scores(self.key(hidden_states))303            value_layer = self.transpose_for_scores(self.value(hidden_states))304            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)305            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)306        else:307            key_layer = self.transpose_for_scores(self.key(hidden_states))308            value_layer = self.transpose_for_scores(self.value(hidden_states))309 310        query_layer = self.transpose_for_scores(mixed_query_layer)311 312        if self.is_decoder:313            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.314            # Further calls to cross_attention layer can then reuse all cross-attention315            # key/value_states (first "if" case)316            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of317            # all previous decoder key/value_states. Further calls to uni-directional self-attention318            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)319            # if encoder bi-directional self-attention `past_key_value` is always `None`320            past_key_value = (key_layer, value_layer)321 322        # Take the dot product between "query" and "key" to get the raw attention scores.323        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))324 325        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":326            seq_length = hidden_states.size()[1]327            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)328            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)329            distance = position_ids_l - position_ids_r330            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)331            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility332 333            if self.position_embedding_type == "relative_key":334                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)335                attention_scores = attention_scores + relative_position_scores336            elif self.position_embedding_type == "relative_key_query":337                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)338                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)339                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key340 341        attention_scores = attention_scores / math.sqrt(self.attention_head_size)342        if attention_mask is not None:343            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)344            attention_scores = attention_scores + attention_mask345 346        # Normalize the attention scores to probabilities.347        attention_probs = nn.functional.softmax(attention_scores, dim=-1)348 349        # This is actually dropping out entire tokens to attend to, which might350        # seem a bit unusual, but is taken from the original Transformer paper.351        attention_probs = self.dropout(attention_probs)352 353        # Mask heads if we want to354        if head_mask is not None:355            attention_probs = attention_probs * head_mask356 357        context_layer = torch.matmul(attention_probs, value_layer)358 359        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()360        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)361        context_layer = context_layer.view(new_context_layer_shape)362 363        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)364 365        if self.is_decoder:366            outputs = outputs + (past_key_value,)367        return outputs368 369 370class BertSelfOutput(nn.Module):371    def __init__(self, config):372        super().__init__()373        self.dense = nn.Linear(config.hidden_size, config.hidden_size)374        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)375        self.dropout = nn.Dropout(config.hidden_dropout_prob)376 377    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:378        hidden_states = self.dense(hidden_states)379        hidden_states = self.dropout(hidden_states)380        hidden_states = self.LayerNorm(hidden_states + input_tensor)381        return hidden_states382 383 384class BertAttention(nn.Module):385    def __init__(self, config, position_embedding_type=None):386        super().__init__()387        self.self = BertSelfAttention(config, position_embedding_type=position_embedding_type)388        self.output = BertSelfOutput(config)389        self.pruned_heads = set()390 391    def prune_heads(self, heads):392        if len(heads) == 0:393            return394        heads, index = find_pruneable_heads_and_indices(395            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads396        )397 398        # Prune linear layers399        self.self.query = prune_linear_layer(self.self.query, index)400        self.self.key = prune_linear_layer(self.self.key, index)401        self.self.value = prune_linear_layer(self.self.value, index)402        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)403 404        # Update hyper params and store pruned heads405        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)406        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads407        self.pruned_heads = self.pruned_heads.union(heads)408 409    def forward(410        self,411        hidden_states: torch.Tensor,412        attention_mask: Optional[torch.FloatTensor] = None,413        head_mask: Optional[torch.FloatTensor] = None,414        encoder_hidden_states: Optional[torch.FloatTensor] = None,415        encoder_attention_mask: Optional[torch.FloatTensor] = None,416        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,417        output_attentions: Optional[bool] = False,418    ) -> Tuple[torch.Tensor]:419        self_outputs = self.self(420            hidden_states,421            attention_mask,422            head_mask,423            encoder_hidden_states,424            encoder_attention_mask,425            past_key_value,426            output_attentions,427        )428        attention_output = self.output(self_outputs[0], hidden_states)429        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them430        return outputs431 432 433class BertIntermediate(nn.Module):434    def __init__(self, config):435        super().__init__()436        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)437        if isinstance(config.hidden_act, str):438            self.intermediate_act_fn = ACT2FN[config.hidden_act]439        else:440            self.intermediate_act_fn = config.hidden_act441 442    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:443        hidden_states = self.dense(hidden_states)444        hidden_states = self.intermediate_act_fn(hidden_states)445        return hidden_states446 447 448class BertOutput(nn.Module):449    def __init__(self, config):450        super().__init__()451        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)452        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)453        self.dropout = nn.Dropout(config.hidden_dropout_prob)454 455    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:456        hidden_states = self.dense(hidden_states)457        hidden_states = self.dropout(hidden_states)458        hidden_states = self.LayerNorm(hidden_states + input_tensor)459        return hidden_states460 461 462class BertLayer(nn.Module):463    def __init__(self, config):464        super().__init__()465        self.chunk_size_feed_forward = config.chunk_size_feed_forward466        self.seq_len_dim = 1467        self.attention = BertAttention(config)468        self.is_decoder = config.is_decoder469        self.add_cross_attention = config.add_cross_attention470        if self.add_cross_attention:471            if not self.is_decoder:472                raise ValueError(f"{self} should be used as a decoder model if cross attention is added")473            self.crossattention = BertAttention(config, position_embedding_type="absolute")474        self.intermediate = BertIntermediate(config)475        self.output = BertOutput(config)476 477    def forward(478        self,479        hidden_states: torch.Tensor,480        attention_mask: Optional[torch.FloatTensor] = None,481        head_mask: Optional[torch.FloatTensor] = None,482        encoder_hidden_states: Optional[torch.FloatTensor] = None,483        encoder_attention_mask: Optional[torch.FloatTensor] = None,484        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,485        output_attentions: Optional[bool] = False,486    ) -> Tuple[torch.Tensor]:487        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2488        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None489        self_attention_outputs = self.attention(490            hidden_states,491            attention_mask,492            head_mask,493            output_attentions=output_attentions,494            past_key_value=self_attn_past_key_value,495        )496        attention_output = self_attention_outputs[0]497 498        # if decoder, the last output is tuple of self-attn cache499        if self.is_decoder:500            outputs = self_attention_outputs[1:-1]501            present_key_value = self_attention_outputs[-1]502        else:503            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights504 505        cross_attn_present_key_value = None506        if self.is_decoder and encoder_hidden_states is not None:507            if not hasattr(self, "crossattention"):508                raise ValueError(509                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"510                    " by setting `config.add_cross_attention=True`"511                )512 513            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple514            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None515            cross_attention_outputs = self.crossattention(516                attention_output,517                attention_mask,518                head_mask,519                encoder_hidden_states,520                encoder_attention_mask,521                cross_attn_past_key_value,522                output_attentions,523            )524            attention_output = cross_attention_outputs[0]525            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights526 527            # add cross-attn cache to positions 3,4 of present_key_value tuple528            cross_attn_present_key_value = cross_attention_outputs[-1]529            present_key_value = present_key_value + cross_attn_present_key_value530 531        layer_output = apply_chunking_to_forward(532            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output533        )534        outputs = (layer_output,) + outputs535 536        # if decoder, return the attn key/values as the last output537        if self.is_decoder:538            outputs = outputs + (present_key_value,)539 540        return outputs541 542    def feed_forward_chunk(self, attention_output):543        intermediate_output = self.intermediate(attention_output)544        layer_output = self.output(intermediate_output, attention_output)545        return layer_output546 547 548class BertEncoder(nn.Module):549    def __init__(self, config):550        super().__init__()551        self.config = config552        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])553        self.gradient_checkpointing = False554 555    def forward(556        self,557        hidden_states: torch.Tensor,558        attention_mask: Optional[torch.FloatTensor] = None,559        head_mask: Optional[torch.FloatTensor] = None,560        encoder_hidden_states: Optional[torch.FloatTensor] = None,561        encoder_attention_mask: Optional[torch.FloatTensor] = None,562        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,563        use_cache: Optional[bool] = None,564        output_attentions: Optional[bool] = False,565        output_hidden_states: Optional[bool] = False,566        return_dict: Optional[bool] = True,567    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:568        all_hidden_states = () if output_hidden_states else None569        all_self_attentions = () if output_attentions else None570        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None571 572        next_decoder_cache = () if use_cache else None573        for i, layer_module in enumerate(self.layer):574            if output_hidden_states:575                all_hidden_states = all_hidden_states + (hidden_states,)576 577            layer_head_mask = head_mask[i] if head_mask is not None else None578            past_key_value = past_key_values[i] if past_key_values is not None else None579 580            if self.gradient_checkpointing and self.training:581 582                if use_cache:583                    logger.warning(584                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."585                    )586                    use_cache = False587 588                def create_custom_forward(module):589                    def custom_forward(*inputs):590                        return module(*inputs, past_key_value, output_attentions)591 592                    return custom_forward593 594                layer_outputs = torch.utils.checkpoint.checkpoint(595                    create_custom_forward(layer_module),596                    hidden_states,597                    attention_mask,598                    layer_head_mask,599                    encoder_hidden_states,600                    encoder_attention_mask,601                )602            else:603                layer_outputs = layer_module(604                    hidden_states,605                    attention_mask,606                    layer_head_mask,607                    encoder_hidden_states,608                    encoder_attention_mask,609                    past_key_value,610                    output_attentions,611                )612 613            hidden_states = layer_outputs[0]614            if use_cache:615                next_decoder_cache += (layer_outputs[-1],)616            if output_attentions:617                all_self_attentions = all_self_attentions + (layer_outputs[1],)618                if self.config.add_cross_attention:619                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)620 621        if output_hidden_states:622            all_hidden_states = all_hidden_states + (hidden_states,)623 624        if not return_dict:625            return tuple(626                v627                for v in [628                    hidden_states,629                    next_decoder_cache,630                    all_hidden_states,631                    all_self_attentions,632                    all_cross_attentions,633                ]634                if v is not None635            )636        return BaseModelOutputWithPastAndCrossAttentions(637            last_hidden_state=hidden_states,638            past_key_values=next_decoder_cache,639            hidden_states=all_hidden_states,640            attentions=all_self_attentions,641            cross_attentions=all_cross_attentions,642        )643 644 645class BertPooler(nn.Module):646    def __init__(self, config):647        super().__init__()648        self.dense = nn.Linear(config.hidden_size, config.hidden_size)649        self.activation = nn.Tanh()650 651    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:652        # We "pool" the model by simply taking the hidden state corresponding653        # to the first token.654        first_token_tensor = hidden_states[:, 0]655        pooled_output = self.dense(first_token_tensor)656        pooled_output = self.activation(pooled_output)657        return pooled_output658 659 660class BertPredictionHeadTransform(nn.Module):661    def __init__(self, config):662        super().__init__()663        self.dense = nn.Linear(config.hidden_size, config.hidden_size)664        if isinstance(config.hidden_act, str):665            self.transform_act_fn = ACT2FN[config.hidden_act]666        else:667            self.transform_act_fn = config.hidden_act668        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)669 670    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:671        hidden_states = self.dense(hidden_states)672        hidden_states = self.transform_act_fn(hidden_states)673        hidden_states = self.LayerNorm(hidden_states)674        return hidden_states675 676 677class BertLMPredictionHead(nn.Module):678    def __init__(self, config):679        super().__init__()680        self.transform = BertPredictionHeadTransform(config)681 682        # The output weights are the same as the input embeddings, but there is683        # an output-only bias for each token.684        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)685 686        self.bias = nn.Parameter(torch.zeros(config.vocab_size))687 688        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`689        self.decoder.bias = self.bias690 691    def forward(self, hidden_states):692        hidden_states = self.transform(hidden_states)693        hidden_states = self.decoder(hidden_states)694        return hidden_states695 696 697class BertOnlyMLMHead(nn.Module):698    def __init__(self, config):699        super().__init__()700        self.predictions = BertLMPredictionHead(config)701 702    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:703        prediction_scores = self.predictions(sequence_output)704        return prediction_scores705 706 707class BertOnlyNSPHead(nn.Module):708    def __init__(self, config):709        super().__init__()710        self.seq_relationship = nn.Linear(config.hidden_size, 2)711 712    def forward(self, pooled_output):713        seq_relationship_score = self.seq_relationship(pooled_output)714        return seq_relationship_score715 716 717class BertPreTrainingHeads(nn.Module):718    def __init__(self, config):719        super().__init__()720        self.predictions = BertLMPredictionHead(config)721        self.seq_relationship = nn.Linear(config.hidden_size, 2)722 723    def forward(self, sequence_output, pooled_output):724        prediction_scores = self.predictions(sequence_output)725        seq_relationship_score = self.seq_relationship(pooled_output)726        return prediction_scores, seq_relationship_score727 728 729class BertPreTrainedModel(PreTrainedModel):730    """731    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained732    models.733    """734 735    config_class = BertConfig736    load_tf_weights = load_tf_weights_in_bert737    base_model_prefix = "bert"738    supports_gradient_checkpointing = True739    _keys_to_ignore_on_load_missing = [r"position_ids"]740 741    def _init_weights(self, module):742        """Initialize the weights"""743        if isinstance(module, nn.Linear):744            # Slightly different from the TF version which uses truncated_normal for initialization745            # cf https://github.com/pytorch/pytorch/pull/5617746            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)747            if module.bias is not None:748                module.bias.data.zero_()749        elif isinstance(module, nn.Embedding):750            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)751            if module.padding_idx is not None:752                module.weight.data[module.padding_idx].zero_()753        elif isinstance(module, nn.LayerNorm):754            module.bias.data.zero_()755            module.weight.data.fill_(1.0)756 757    def _set_gradient_checkpointing(self, module, value=False):758        if isinstance(module, BertEncoder):759            module.gradient_checkpointing = value760 761 762@dataclass763class BertForPreTrainingOutput(ModelOutput):764    """765    Output type of [`BertForPreTraining`].766 767    Args:768        loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):769            Total loss as the sum of the masked language modeling loss and the next sequence prediction770            (classification) loss.771        prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):772            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).773        seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):774            Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation775            before SoftMax).776        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):777            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of778            shape `(batch_size, sequence_length, hidden_size)`.779 780            Hidden-states of the model at the output of each layer plus the initial embedding outputs.781        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):782            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,783            sequence_length)`.784 785            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention786            heads.787    """788 789    loss: Optional[torch.FloatTensor] = None790    prediction_logits: torch.FloatTensor = None791    seq_relationship_logits: torch.FloatTensor = None792    hidden_states: Optional[Tuple[torch.FloatTensor]] = None793    attentions: Optional[Tuple[torch.FloatTensor]] = None794 795 796BERT_START_DOCSTRING = r"""797 798    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the799    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads800    etc.)801 802    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.803    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage804    and behavior.805 806    Parameters:807        config ([`BertConfig`]): Model configuration class with all the parameters of the model.808            Initializing with a config file does not load the weights associated with the model, only the809            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.810"""811 812BERT_INPUTS_DOCSTRING = r"""813    Args:814        input_ids (`torch.LongTensor` of shape `({0})`):815            Indices of input sequence tokens in the vocabulary.816 817            Indices can be obtained using [`BertTokenizer`]. See [`PreTrainedTokenizer.encode`] and818            [`PreTrainedTokenizer.__call__`] for details.819 820            [What are input IDs?](../glossary#input-ids)821        attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):822            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:823 824            - 1 for tokens that are **not masked**,825            - 0 for tokens that are **masked**.826 827            [What are attention masks?](../glossary#attention-mask)828        token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):829            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,830            1]`:831 832            - 0 corresponds to a *sentence A* token,833            - 1 corresponds to a *sentence B* token.834 835            [What are token type IDs?](../glossary#token-type-ids)836        position_ids (`torch.LongTensor` of shape `({0})`, *optional*):837            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,838            config.max_position_embeddings - 1]`.839 840            [What are position IDs?](../glossary#position-ids)841        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):842            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:843 844            - 1 indicates the head is **not masked**,845            - 0 indicates the head is **masked**.846 847        inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):848            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This849            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the850            model's internal embedding lookup matrix.851        output_attentions (`bool`, *optional*):852            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned853            tensors for more detail.854        output_hidden_states (`bool`, *optional*):855            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for856            more detail.857        return_dict (`bool`, *optional*):858            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.859"""860 861 862@add_start_docstrings(863    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",864    BERT_START_DOCSTRING,865)866class BertModel(BertPreTrainedModel):867    """868 869    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of870    cross-attention is added between the self-attention layers, following the architecture described in [Attention is871    all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,872    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.873 874    To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set875    to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and876    `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.877    """878 879    def __init__(self, config, add_pooling_layer=True):880        super().__init__(config)881        self.config = config882 883        self.embeddings = BertEmbeddings(config)884        self.encoder = BertEncoder(config)885 886        self.pooler = BertPooler(config) if add_pooling_layer else None887 888        # Initialize weights and apply final processing889        self.post_init()890 891    def get_input_embeddings(self):892        return self.embeddings.word_embeddings893 894    def set_input_embeddings(self, value):895        self.embeddings.word_embeddings = value896 897    def _prune_heads(self, heads_to_prune):898        """899        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base900        class PreTrainedModel901        """902        for layer, heads in heads_to_prune.items():903            self.encoder.layer[layer].attention.prune_heads(heads)904 905    @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))906    @add_code_sample_docstrings(907        processor_class=_TOKENIZER_FOR_DOC,908        checkpoint=_CHECKPOINT_FOR_DOC,909        output_type=BaseModelOutputWithPoolingAndCrossAttentions,910        config_class=_CONFIG_FOR_DOC,911    )912    def forward(913        self,914        input_ids: Optional[torch.Tensor] = None,915        attention_mask: Optional[torch.Tensor] = None,916        token_type_ids: Optional[torch.Tensor] = None,917        position_ids: Optional[torch.Tensor] = None,918        head_mask: Optional[torch.Tensor] = None,919        inputs_embeds: Optional[torch.Tensor] = None,920        encoder_hidden_states: Optional[torch.Tensor] = None,921        encoder_attention_mask: Optional[torch.Tensor] = None,922        past_key_values: Optional[List[torch.FloatTensor]] = None,923        use_cache: Optional[bool] = None,924        output_attentions: Optional[bool] = None,925        output_hidden_states: Optional[bool] = None,926        return_dict: Optional[bool] = None,927    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:928        r"""929        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):930            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if931            the model is configured as a decoder.932        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):933            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in934            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:935 936            - 1 for tokens that are **not masked**,937            - 0 for tokens that are **masked**.938        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):939            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.940 941            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that942            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all943            `decoder_input_ids` of shape `(batch_size, sequence_length)`.944        use_cache (`bool`, *optional*):945            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see946            `past_key_values`).947        """948        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions949        output_hidden_states = (950            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states951        )952        return_dict = return_dict if return_dict is not None else self.config.use_return_dict953 954        if self.config.is_decoder:955            use_cache = use_cache if use_cache is not None else self.config.use_cache956        else:957            use_cache = False958 959        if input_ids is not None and inputs_embeds is not None:960            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")961        elif input_ids is not None:962            input_shape = input_ids.size()963        elif inputs_embeds is not None:964            input_shape = inputs_embeds.size()[:-1]965        else:966            raise ValueError("You have to specify either input_ids or inputs_embeds")967 968        batch_size, seq_length = input_shape969        device = input_ids.device if input_ids is not None else inputs_embeds.device970 971        # past_key_values_length972        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0973 974        if attention_mask is None:975            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)976 977        if token_type_ids is None:978            if hasattr(self.embeddings, "token_type_ids"):979                buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]980                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)981                token_type_ids = buffered_token_type_ids_expanded982            else:983                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)984 985        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]986        # ourselves in which case we just need to make it broadcastable to all heads.987        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)988 989        # If a 2D or 3D attention mask is provided for the cross-attention990        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]991        if self.config.is_decoder and encoder_hidden_states is not None:992            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()993            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)994            if encoder_attention_mask is None:995                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)996            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)997        else:998            encoder_extended_attention_mask = None999 1000        # Prepare head mask if needed1001        # 1.0 in head_mask indicate we keep the head1002        # attention_probs has shape bsz x n_heads x N x N1003        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]1004        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]1005        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)1006 1007        embedding_output = self.embeddings(1008            input_ids=input_ids,1009            position_ids=position_ids,1010            token_type_ids=token_type_ids,1011            inputs_embeds=inputs_embeds,1012            past_key_values_length=past_key_values_length,1013        )1014        encoder_outputs = self.encoder(1015            embedding_output,1016            attention_mask=extended_attention_mask,1017            head_mask=head_mask,1018            encoder_hidden_states=encoder_hidden_states,1019            encoder_attention_mask=encoder_extended_attention_mask,1020            past_key_values=past_key_values,1021            use_cache=use_cache,1022            output_attentions=output_attentions,1023            output_hidden_states=output_hidden_states,1024            return_dict=return_dict,1025        )1026        sequence_output = encoder_outputs[0]1027        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None1028 1029        if not return_dict:1030            return (sequence_output, pooled_output) + encoder_outputs[1:]1031 1032        return BaseModelOutputWithPoolingAndCrossAttentions(1033            last_hidden_state=sequence_output,1034            pooler_output=pooled_output,1035            past_key_values=encoder_outputs.past_key_values,1036            hidden_states=encoder_outputs.hidden_states,1037            attentions=encoder_outputs.attentions,1038            cross_attentions=encoder_outputs.cross_attentions,1039        )1040 1041 1042@add_start_docstrings(1043    """1044    Bert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next1045    sentence prediction (classification)` head.1046    """,1047    BERT_START_DOCSTRING,1048)1049class BertForPreTraining(BertPreTrainedModel):1050    def __init__(self, config):1051        super().__init__(config)1052 1053        self.bert = BertModel(config)1054        self.cls = BertPreTrainingHeads(config)1055 1056        # Initialize weights and apply final processing1057        self.post_init()1058 1059    def get_output_embeddings(self):1060        return self.cls.predictions.decoder1061 1062    def set_output_embeddings(self, new_embeddings):1063        self.cls.predictions.decoder = new_embeddings1064 1065    @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))1066    @replace_return_docstrings(output_type=BertForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)1067    def forward(1068        self,1069        input_ids: Optional[torch.Tensor] = None,1070        attention_mask: Optional[torch.Tensor] = None,1071        token_type_ids: Optional[torch.Tensor] = None,1072        position_ids: Optional[torch.Tensor] = None,1073        head_mask: Optional[torch.Tensor] = None,1074        inputs_embeds: Optional[torch.Tensor] = None,1075        labels: Optional[torch.Tensor] = None,1076        next_sentence_label: Optional[torch.Tensor] = None,1077        output_attentions: Optional[bool] = None,1078        output_hidden_states: Optional[bool] = None,1079        return_dict: Optional[bool] = None,1080    ) -> Union[Tuple[torch.Tensor], BertForPreTrainingOutput]:1081        r"""1082            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1083                Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,1084                config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked),1085                the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`1086            next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1087                Labels for computing the next sequence prediction (classification) loss. Input should be a sequence1088                pair (see `input_ids` docstring) Indices should be in `[0, 1]`:1089 1090                - 0 indicates sequence B is a continuation of sequence A,1091                - 1 indicates sequence B is a random sequence.1092            kwargs (`Dict[str, any]`, optional, defaults to *{}*):1093                Used to hide legacy arguments that have been deprecated.1094 1095        Returns:1096 1097        Example:1098 1099        ```python1100        >>> from transformers import BertTokenizer, BertForPreTraining1101        >>> import torch1102 1103        >>> tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")1104        >>> model = BertForPreTraining.from_pretrained("bert-base-uncased")1105 1106        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")1107        >>> outputs = model(**inputs)1108 1109        >>> prediction_logits = outputs.prediction_logits1110        >>> seq_relationship_logits = outputs.seq_relationship_logits1111        ```1112        """1113        return_dict = return_dict if return_dict is not None else self.config.use_return_dict1114 1115        outputs = self.bert(1116            input_ids,1117            attention_mask=attention_mask,1118            token_type_ids=token_type_ids,1119            position_ids=position_ids,1120            head_mask=head_mask,1121            inputs_embeds=inputs_embeds,1122            output_attentions=output_attentions,1123            output_hidden_states=output_hidden_states,1124            return_dict=return_dict,1125        )1126 1127        sequence_output, pooled_output = outputs[:2]1128        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)1129 1130        total_loss = None1131        if labels is not None and next_sentence_label is not None:1132            loss_fct = CrossEntropyLoss()1133            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))1134            next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))1135            total_loss = masked_lm_loss + next_sentence_loss1136 1137        if not return_dict:1138            output = (prediction_scores, seq_relationship_score) + outputs[2:]1139            return ((total_loss,) + output) if total_loss is not None else output1140 1141        return BertForPreTrainingOutput(1142            loss=total_loss,1143            prediction_logits=prediction_scores,1144            seq_relationship_logits=seq_relationship_score,1145            hidden_states=outputs.hidden_states,1146            attentions=outputs.attentions,1147        )1148 1149 1150@add_start_docstrings(1151    """Bert Model with a `language modeling` head on top for CLM fine-tuning.""", BERT_START_DOCSTRING1152)1153class BertLMHeadModel(BertPreTrainedModel):1154 1155    _keys_to_ignore_on_load_unexpected = [r"pooler"]1156    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]1157 1158    def __init__(self, config):1159        super().__init__(config)1160 1161        if not config.is_decoder:1162            logger.warning("If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`")1163 1164        self.bert = BertModel(config, add_pooling_layer=False)1165        self.cls = BertOnlyMLMHead(config)1166 1167        # Initialize weights and apply final processing1168        self.post_init()1169 1170    def get_output_embeddings(self):1171        return self.cls.predictions.decoder1172 1173    def set_output_embeddings(self, new_embeddings):1174        self.cls.predictions.decoder = new_embeddings1175 1176    @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))1177    @add_code_sample_docstrings(1178        processor_class=_TOKENIZER_FOR_DOC,1179        checkpoint=_CHECKPOINT_FOR_DOC,1180        output_type=CausalLMOutputWithCrossAttentions,1181        config_class=_CONFIG_FOR_DOC,1182    )1183    def forward(1184        self,1185        input_ids: Optional[torch.Tensor] = None,1186        attention_mask: Optional[torch.Tensor] = None,1187        token_type_ids: Optional[torch.Tensor] = None,1188        position_ids: Optional[torch.Tensor] = None,1189        head_mask: Optional[torch.Tensor] = None,1190        inputs_embeds: Optional[torch.Tensor] = None,1191        encoder_hidden_states: Optional[torch.Tensor] = None,1192        encoder_attention_mask: Optional[torch.Tensor] = None,1193        labels: Optional[torch.Tensor] = None,1194        past_key_values: Optional[List[torch.Tensor]] = None,1195        use_cache: Optional[bool] = None,1196        output_attentions: Optional[bool] = None,1197        output_hidden_states: Optional[bool] = None,1198        return_dict: Optional[bool] = None,1199    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:1200        r"""

Showing the first 1,200 of 1890 lines. Download the file for the rest.