devbox
IT小白
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

Bert的位置编码：position_embeddings（绝对位置可学习参数式编码）_bert中的positonal embedding

作者：IT小白 | 2024-03-25 20:47:13

赞

踩

bert中的positonal embedding

position_embeddings（绝对位置可学习参数式编码）

  1. class BertEmbeddings(nn.Module):
  2.     """Construct the embeddings from word, position and token_type embeddings.
  3.     """
  4.     def __init__(self, config):
  5.         super(BertEmbeddings, self).__init__()
  6.         self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
  7.         self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
  8.         self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
  9.  
  10.         # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
  11.         # any TensorFlow checkpoint file
  12.         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
  13.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  14.  
  15.     def forward(self, input_ids, token_type_ids=None):
  16.         seq_length = input_ids.size(1)  # 32
  17.         position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)  # torch.Size([32])
  18.         position_ids = position_ids.unsqueeze(0).expand_as(input_ids)   # torch.Size([32])---->torch.Size([128, 32])
  19.         if token_type_ids is None:
  20.             token_type_ids = torch.zeros_like(input_ids)
  21.  
  22.         words_embeddings = self.word_embeddings(input_ids)  # torch.Size([128, 32])---->torch.Size([128, 32, 768])
  23.         position_embeddings = self.position_embeddings(position_ids)    # torch.Size([128, 32])---->torch.Size([128, 32, 768])
  24.         token_type_embeddings = self.token_type_embeddings(token_type_ids)  # torch.Size([128, 32])---->torch.Size([128, 32, 768])
  25.  
  26.         embeddings = words_embeddings + position_embeddings + token_type_embeddings # torch.Size([128, 32, 768])
  27.         embeddings = self.LayerNorm(embeddings) # torch.Size([128, 32, 768])---->torch.Size([128, 32, 768])
  28.         embeddings = self.dropout(embeddings)   # torch.Size([128, 32, 768])
  29.         return embeddings

完整代码：

  1. # coding=utf-8
  2. # 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 at
  8. #
  9. #     http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # 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 and
  15. # limitations under the License.
  16. """PyTorch BERT model."""
  17.  
  18. from __future__ import absolute_import, division, print_function, unicode_literals
  19.  
  20. import copy
  21. import json
  22. import logging
  23. import math
  24. import os
  25. import shutil
  26. import tarfile
  27. import tempfile
  28. import sys
  29. from io import open
  30.  
  31. import torch
  32. from torch import nn
  33. from torch.nn import CrossEntropyLoss
  34.  
  35. from .file_utils import cached_path, WEIGHTS_NAME, CONFIG_NAME
  36.  
  37. logger = logging.getLogger(__name__)
  38.  
  39. PRETRAINED_MODEL_ARCHIVE_MAP = {
  40.     'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased.tar.gz",
  41.     'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased.tar.gz",
  42.     'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased.tar.gz",
  43.     'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased.tar.gz",
  44.     'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased.tar.gz",
  45.     'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased.tar.gz",
  46.     'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese.tar.gz",
  47. }
  48. BERT_CONFIG_NAME = 'bert_config.json'
  49. TF_WEIGHTS_NAME = 'model.ckpt'
  50.  
  51. def load_tf_weights_in_bert(model, tf_checkpoint_path):
  52.     """ Load tf checkpoints in a pytorch model
  53.     """
  54.     try:
  55.         import re
  56.         import numpy as np
  57.         import tensorflow as tf
  58.     except ImportError:
  59.         print("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
  60.             "https://www.tensorflow.org/install/ for installation instructions.")
  61.         raise
  62.     tf_path = os.path.abspath(tf_checkpoint_path)
  63.     print("Converting TensorFlow checkpoint from {}".format(tf_path))
  64.     # Load weights from TF model
  65.     init_vars = tf.train.list_variables(tf_path)
  66.     names = []
  67.     arrays = []
  68.     for name, shape in init_vars:
  69.         print("Loading TF weight {} with shape {}".format(name, shape))
  70.         array = tf.train.load_variable(tf_path, name)
  71.         names.append(name)
  72.         arrays.append(array)
  73.  
  74.     for name, array in zip(names, arrays):
  75.         name = name.split('/')
  76.         # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
  77.         # which are not required for using pretrained model
  78.         if any(n in ["adam_v", "adam_m", "global_step"] for n in name):
  79.             print("Skipping {}".format("/".join(name)))
  80.             continue
  81.         pointer = model
  82.         for m_name in name:
  83.             if re.fullmatch(r'[A-Za-z]+_\d+', m_name):
  84.                 l = re.split(r'_(\d+)', m_name)
  85.             else:
  86.                 l = [m_name]
  87.             if l[0] == 'kernel' or l[0] == 'gamma':
  88.                 pointer = getattr(pointer, 'weight')
  89.             elif l[0] == 'output_bias' or l[0] == 'beta':
  90.                 pointer = getattr(pointer, 'bias')
  91.             elif l[0] == 'output_weights':
  92.                 pointer = getattr(pointer, 'weight')
  93.             elif l[0] == 'squad':
  94.                 pointer = getattr(pointer, 'classifier')
  95.             else:
  96.                 try:
  97.                     pointer = getattr(pointer, l[0])
  98.                 except AttributeError:
  99.                     print("Skipping {}".format("/".join(name)))
  100.                     continue
  101.             if len(l) >= 2:
  102.                 num = int(l[1])
  103.                 pointer = pointer[num]
  104.         if m_name[-11:] == '_embeddings':
  105.             pointer = getattr(pointer, 'weight')
  106.         elif m_name == 'kernel':
  107.             array = np.transpose(array)
  108.         try:
  109.             assert pointer.shape == array.shape
  110.         except AssertionError as e:
  111.             e.args += (pointer.shape, array.shape)
  112.             raise
  113.         print("Initialize PyTorch weight {}".format(name))
  114.         pointer.data = torch.from_numpy(array)
  115.     return model
  116.  
  117.  
  118. def gelu(x):
  119.     """Implementation of the gelu activation function.
  120.         For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
  121.         0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
  122.         Also see https://arxiv.org/abs/1606.08415
  123.     """
  124.     return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
  125.  
  126.  
  127. def swish(x):
  128.     return x * torch.sigmoid(x)
  129.  
  130.  
  131. ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish}
  132.  
  133.  
  134. class BertConfig(object):
  135.     """Configuration class to store the configuration of a `BertModel`.
  136.     """
  137.     def __init__(self,
  138.                  vocab_size_or_config_json_file,
  139.                  hidden_size=768,
  140.                  num_hidden_layers=12,
  141.                  num_attention_heads=12,
  142.                  intermediate_size=3072,
  143.                  hidden_act="gelu",
  144.                  hidden_dropout_prob=0.1,
  145.                  attention_probs_dropout_prob=0.1,
  146.                  max_position_embeddings=512,
  147.                  type_vocab_size=2,
  148.                  initializer_range=0.02):
  149.         """Constructs BertConfig.
  151.         Args:
  152.             vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
  153.             hidden_size: Size of the encoder layers and the pooler layer.
  154.             num_hidden_layers: Number of hidden layers in the Transformer encoder.
  155.             num_attention_heads: Number of attention heads for each attention layer in
  156.                 the Transformer encoder.
  157.             intermediate_size: The size of the "intermediate" (i.e., feed-forward)
  158.                 layer in the Transformer encoder.
  159.             hidden_act: The non-linear activation function (function or string) in the
  160.                 encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
  161.             hidden_dropout_prob: The dropout probabilitiy for all fully connected
  162.                 layers in the embeddings, encoder, and pooler.
  163.             attention_probs_dropout_prob: The dropout ratio for the attention
  164.                 probabilities.
  165.             max_position_embeddings: The maximum sequence length that this model might
  166.                 ever be used with. Typically set this to something large just in case
  167.                 (e.g., 512 or 1024 or 2048).
  168.             type_vocab_size: The vocabulary size of the `token_type_ids` passed into
  169.                 `BertModel`.
  170.             initializer_range: The sttdev of the truncated_normal_initializer for
  171.                 initializing all weight matrices.
  172.         """
  173.         if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
  174.                         and isinstance(vocab_size_or_config_json_file, unicode)):
  175.             with open(vocab_size_or_config_json_file, "r", encoding='utf-8') as reader:
  176.                 json_config = json.loads(reader.read())
  177.             for key, value in json_config.items():
  178.                 self.__dict__[key] = value
  179.         elif isinstance(vocab_size_or_config_json_file, int):
  180.             self.vocab_size = vocab_size_or_config_json_file
  181.             self.hidden_size = hidden_size
  182.             self.num_hidden_layers = num_hidden_layers
  183.             self.num_attention_heads = num_attention_heads
  184.             self.hidden_act = hidden_act
  185.             self.intermediate_size = intermediate_size
  186.             self.hidden_dropout_prob = hidden_dropout_prob
  187.             self.attention_probs_dropout_prob = attention_probs_dropout_prob
  188.             self.max_position_embeddings = max_position_embeddings
  189.             self.type_vocab_size = type_vocab_size
  190.             self.initializer_range = initializer_range
  191.         else:
  192.             raise ValueError("First argument must be either a vocabulary size (int)"
  193.                              "or the path to a pretrained model config file (str)")
  194.  
  195.     @classmethod
  196.     def from_dict(cls, json_object):
  197.         """Constructs a `BertConfig` from a Python dictionary of parameters."""
  198.         config = BertConfig(vocab_size_or_config_json_file=-1)
  199.         for key, value in json_object.items():
  200.             config.__dict__[key] = value
  201.         return config
  202.  
  203.     @classmethod
  204.     def from_json_file(cls, json_file):
  205.         """Constructs a `BertConfig` from a json file of parameters."""
  206.         with open(json_file, "r", encoding='utf-8') as reader:
  207.             text = reader.read()
  208.         return cls.from_dict(json.loads(text))
  209.  
  210.     def __repr__(self):
  211.         return str(self.to_json_string())
  212.  
  213.     def to_dict(self):
  214.         """Serializes this instance to a Python dictionary."""
  215.         output = copy.deepcopy(self.__dict__)
  216.         return output
  217.  
  218.     def to_json_string(self):
  219.         """Serializes this instance to a JSON string."""
  220.         return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
  221.  
  222.     def to_json_file(self, json_file_path):
  223.         """ Save this instance to a json file."""
  224.         with open(json_file_path, "w", encoding='utf-8') as writer:
  225.             writer.write(self.to_json_string())
  226.  
  227. try:
  228.     from apex.normalization.fused_layer_norm import FusedLayerNorm as BertLayerNorm
  229. except ImportError:
  230.     logger.info("Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex .")
  231.     class BertLayerNorm(nn.Module):
  232.         def __init__(self, hidden_size, eps=1e-12):
  233.             """Construct a layernorm module in the TF style (epsilon inside the square root).
  234.             """
  235.             super(BertLayerNorm, self).__init__()
  236.             self.weight = nn.Parameter(torch.ones(hidden_size))
  237.             self.bias = nn.Parameter(torch.zeros(hidden_size))
  238.             self.variance_epsilon = eps
  239.  
  240.         def forward(self, x):
  241.             u = x.mean(-1, keepdim=True)
  242.             s = (x - u).pow(2).mean(-1, keepdim=True)
  243.             x = (x - u) / torch.sqrt(s + self.variance_epsilon)
  244.             return self.weight * x + self.bias
  245.  
  246. class BertEmbeddings(nn.Module):
  247.     """Construct the embeddings from word, position and token_type embeddings.
  248.     """
  249.     def __init__(self, config):
  250.         super(BertEmbeddings, self).__init__()
  251.         self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
  252.         self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
  253.         self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
  254.  
  255.         # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
  256.         # any TensorFlow checkpoint file
  257.         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
  258.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  259.  
  260.     def forward(self, input_ids, token_type_ids=None):
  261.         seq_length = input_ids.size(1)
  262.         position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
  263.         position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
  264.         if token_type_ids is None:
  265.             token_type_ids = torch.zeros_like(input_ids)
  266.  
  267.         words_embeddings = self.word_embeddings(input_ids)
  268.         position_embeddings = self.position_embeddings(position_ids)
  269.         token_type_embeddings = self.token_type_embeddings(token_type_ids)
  270.  
  271.         embeddings = words_embeddings + position_embeddings + token_type_embeddings
  272.         embeddings = self.LayerNorm(embeddings)
  273.         embeddings = self.dropout(embeddings)
  274.         return embeddings
  275.  
  276.  
  277. class BertSelfAttention(nn.Module):
  278.     def __init__(self, config):
  279.         super(BertSelfAttention, self).__init__()
  280.         if config.hidden_size % config.num_attention_heads != 0:
  281.             raise ValueError(
  282.                 "The hidden size (%d) is not a multiple of the number of attention "
  283.                 "heads (%d)" % (config.hidden_size, config.num_attention_heads))
  284.         self.num_attention_heads = config.num_attention_heads
  285.         self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
  286.         self.all_head_size = self.num_attention_heads * self.attention_head_size
  287.  
  288.         self.query = nn.Linear(config.hidden_size, self.all_head_size)
  289.         self.key = nn.Linear(config.hidden_size, self.all_head_size)
  290.         self.value = nn.Linear(config.hidden_size, self.all_head_size)
  291.  
  292.         self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
  293.  
  294.     def transpose_for_scores(self, x):
  295.         new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
  296.         x = x.view(*new_x_shape)
  297.         return x.permute(0, 2, 1, 3)
  298.  
  299.     def forward(self, hidden_states, attention_mask):
  300.         mixed_query_layer = self.query(hidden_states)
  301.         mixed_key_layer = self.key(hidden_states)
  302.         mixed_value_layer = self.value(hidden_states)
  303.  
  304.         query_layer = self.transpose_for_scores(mixed_query_layer)
  305.         key_layer = self.transpose_for_scores(mixed_key_layer)
  306.         value_layer = self.transpose_for_scores(mixed_value_layer)
  307.  
  308.         # Take the dot product between "query" and "key" to get the raw attention scores.
  309.         attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
  310.         attention_scores = attention_scores / math.sqrt(self.attention_head_size)
  311.         # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
  312.         attention_scores = attention_scores + attention_mask
  313.  
  314.         # Normalize the attention scores to probabilities.
  315.         attention_probs = nn.Softmax(dim=-1)(attention_scores)
  316.  
  317.         # This is actually dropping out entire tokens to attend to, which might
  318.         # seem a bit unusual, but is taken from the original Transformer paper.
  319.         attention_probs = self.dropout(attention_probs)
  320.  
  321.         context_layer = torch.matmul(attention_probs, value_layer)
  322.         context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
  323.         new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
  324.         context_layer = context_layer.view(*new_context_layer_shape)
  325.         return context_layer
  326.  
  327.  
  328. class BertSelfOutput(nn.Module):
  329.     def __init__(self, config):
  330.         super(BertSelfOutput, self).__init__()
  331.         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
  332.         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
  333.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  334.  
  335.     def forward(self, hidden_states, input_tensor):
  336.         hidden_states = self.dense(hidden_states)
  337.         hidden_states = self.dropout(hidden_states)
  338.         hidden_states = self.LayerNorm(hidden_states + input_tensor)
  339.         return hidden_states
  340.  
  341.  
  342. class BertAttention(nn.Module):
  343.     def __init__(self, config):
  344.         super(BertAttention, self).__init__()
  345.         self.self = BertSelfAttention(config)
  346.         self.output = BertSelfOutput(config)
  347.  
  348.     def forward(self, input_tensor, attention_mask):
  349.         self_output = self.self(input_tensor, attention_mask)
  350.         attention_output = self.output(self_output, input_tensor)
  351.         return attention_output
  352.  
  353.  
  354. class BertIntermediate(nn.Module):
  355.     def __init__(self, config):
  356.         super(BertIntermediate, self).__init__()
  357.         self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
  358.         if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
  359.             self.intermediate_act_fn = ACT2FN[config.hidden_act]
  360.         else:
  361.             self.intermediate_act_fn = config.hidden_act
  362.  
  363.     def forward(self, hidden_states):
  364.         hidden_states = self.dense(hidden_states)
  365.         hidden_states = self.intermediate_act_fn(hidden_states)
  366.         return hidden_states
  367.  
  368.  
  369. class BertOutput(nn.Module):
  370.     def __init__(self, config):
  371.         super(BertOutput, self).__init__()
  372.         self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
  373.         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
  374.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  375.  
  376.     def forward(self, hidden_states, input_tensor):
  377.         hidden_states = self.dense(hidden_states)
  378.         hidden_states = self.dropout(hidden_states)
  379.         hidden_states = self.LayerNorm(hidden_states + input_tensor)
  380.         return hidden_states
  381.  
  382.  
  383. class BertLayer(nn.Module):
  384.     def __init__(self, config):
  385.         super(BertLayer, self).__init__()
  386.         self.attention = BertAttention(config)
  387.         self.intermediate = BertIntermediate(config)
  388.         self.output = BertOutput(config)
  389.  
  390.     def forward(self, hidden_states, attention_mask):
  391.         attention_output = self.attention(hidden_states, attention_mask)
  392.         intermediate_output = self.intermediate(attention_output)
  393.         layer_output = self.output(intermediate_output, attention_output)
  394.         return layer_output
  395.  
  396.  
  397. class BertEncoder(nn.Module):
  398.     def __init__(self, config):
  399.         super(BertEncoder, self).__init__()
  400.         layer = BertLayer(config)
  401.         self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])
  402.  
  403.     def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True):
  404.         all_encoder_layers = []
  405.         for layer_module in self.layer:
  406.             hidden_states = layer_module(hidden_states, attention_mask)
  407.             if output_all_encoded_layers:
  408.                 all_encoder_layers.append(hidden_states)
  409.         if not output_all_encoded_layers:
  410.             all_encoder_layers.append(hidden_states)
  411.         return all_encoder_layers
  412.  
  413.  
  414. class BertPooler(nn.Module):
  415.     def __init__(self, config):
  416.         super(BertPooler, self).__init__()
  417.         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
  418.         self.activation = nn.Tanh()
  419.  
  420.     def forward(self, hidden_states):
  421.         # We "pool" the model by simply taking the hidden state corresponding
  422.         # to the first token.
  423.         first_token_tensor = hidden_states[:, 0]
  424.         pooled_output = self.dense(first_token_tensor)
  425.         pooled_output = self.activation(pooled_output)
  426.         return pooled_output
  427.  
  428.  
  429. class BertPredictionHeadTransform(nn.Module):
  430.     def __init__(self, config):
  431.         super(BertPredictionHeadTransform, self).__init__()
  432.         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
  433.         if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
  434.             self.transform_act_fn = ACT2FN[config.hidden_act]
  435.         else:
  436.             self.transform_act_fn = config.hidden_act
  437.         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
  438.  
  439.     def forward(self, hidden_states):
  440.         hidden_states = self.dense(hidden_states)
  441.         hidden_states = self.transform_act_fn(hidden_states)
  442.         hidden_states = self.LayerNorm(hidden_states)
  443.         return hidden_states
  444.  
  445.  
  446. class BertLMPredictionHead(nn.Module):
  447.     def __init__(self, config, bert_model_embedding_weights):
  448.         super(BertLMPredictionHead, self).__init__()
  449.         self.transform = BertPredictionHeadTransform(config)
  450.  
  451.         # The output weights are the same as the input embeddings, but there is
  452.         # an output-only bias for each token.
  453.         self.decoder = nn.Linear(bert_model_embedding_weights.size(1),
  454.                                  bert_model_embedding_weights.size(0),
  455.                                  bias=False)
  456.         self.decoder.weight = bert_model_embedding_weights
  457.         self.bias = nn.Parameter(torch.zeros(bert_model_embedding_weights.size(0)))
  458.  
  459.     def forward(self, hidden_states):
  460.         hidden_states = self.transform(hidden_states)
  461.         hidden_states = self.decoder(hidden_states) + self.bias
  462.         return hidden_states
  463.  
  464.  
  465. class BertOnlyMLMHead(nn.Module):
  466.     def __init__(self, config, bert_model_embedding_weights):
  467.         super(BertOnlyMLMHead, self).__init__()
  468.         self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)
  469.  
  470.     def forward(self, sequence_output):
  471.         prediction_scores = self.predictions(sequence_output)
  472.         return prediction_scores
  473.  
  474.  
  475. class BertOnlyNSPHead(nn.Module):
  476.     def __init__(self, config):
  477.         super(BertOnlyNSPHead, self).__init__()
  478.         self.seq_relationship = nn.Linear(config.hidden_size, 2)
  479.  
  480.     def forward(self, pooled_output):
  481.         seq_relationship_score = self.seq_relationship(pooled_output)
  482.         return seq_relationship_score
  483.  
  484.  
  485. class BertPreTrainingHeads(nn.Module):
  486.     def __init__(self, config, bert_model_embedding_weights):
  487.         super(BertPreTrainingHeads, self).__init__()
  488.         self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)
  489.         self.seq_relationship = nn.Linear(config.hidden_size, 2)
  490.  
  491.     def forward(self, sequence_output, pooled_output):
  492.         prediction_scores = self.predictions(sequence_output)
  493.         seq_relationship_score = self.seq_relationship(pooled_output)
  494.         return prediction_scores, seq_relationship_score
  495.  
  496.  
  497. class BertPreTrainedModel(nn.Module):
  498.     """ An abstract class to handle weights initialization and
  499.         a simple interface for dowloading and loading pretrained models.
  500.     """
  501.     def __init__(self, config, *inputs, **kwargs):
  502.         super(BertPreTrainedModel, self).__init__()
  503.         if not isinstance(config, BertConfig):
  504.             raise ValueError(
  505.                 "Parameter config in `{}(config)` should be an instance of class `BertConfig`. "
  506.                 "To create a model from a Google pretrained model use "
  507.                 "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
  508.                     self.__class__.__name__, self.__class__.__name__
  509.                 ))
  510.         self.config = config
  511.  
  512.     def init_bert_weights(self, module):
  513.         """ Initialize the weights.
  514.         """
  515.         if isinstance(module, (nn.Linear, nn.Embedding)):
  516.             # Slightly different from the TF version which uses truncated_normal for initialization
  517.             # cf https://github.com/pytorch/pytorch/pull/5617
  518.             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
  519.         elif isinstance(module, BertLayerNorm):
  520.             module.bias.data.zero_()
  521.             module.weight.data.fill_(1.0)
  522.         if isinstance(module, nn.Linear) and module.bias is not None:
  523.             module.bias.data.zero_()
  524.  
  525.     @classmethod
  526.     def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
  527.         """
  528.         Instantiate a BertPreTrainedModel from a pre-trained model file or a pytorch state dict.
  529.         Download and cache the pre-trained model file if needed.
  531.         Params:
  532.             pretrained_model_name_or_path: either:
  533.                 - a str with the name of a pre-trained model to load selected in the list of:
  534.                     . `bert-base-uncased`
  535.                     . `bert-large-uncased`
  536.                     . `bert-base-cased`
  537.                     . `bert-large-cased`
  538.                     . `bert-base-multilingual-uncased`
  539.                     . `bert-base-multilingual-cased`
  540.                     . `bert-base-chinese`
  541.                 - a path or url to a pretrained model archive containing:
  542.                     . `bert_config.json` a configuration file for the model
  543.                     . `pytorch_model.bin` a PyTorch dump of a BertForPreTraining instance
  544.                 - a path or url to a pretrained model archive containing:
  545.                     . `bert_config.json` a configuration file for the model
  546.                     . `model.chkpt` a TensorFlow checkpoint
  547.             from_tf: should we load the weights from a locally saved TensorFlow checkpoint
  548.             cache_dir: an optional path to a folder in which the pre-trained models will be cached.
  549.             state_dict: an optional state dictionnary (collections.OrderedDict object) to use instead of Google pre-trained models
  550.             *inputs, **kwargs: additional input for the specific Bert class
  551.                 (ex: num_labels for BertForSequenceClassification)
  552.         """
  553.         state_dict = kwargs.get('state_dict', None)
  554.         kwargs.pop('state_dict', None)
  555.         cache_dir = kwargs.get('cache_dir', None)
  556.         kwargs.pop('cache_dir', None)
  557.         from_tf = kwargs.get('from_tf', False)
  558.         kwargs.pop('from_tf', None)
  559.  
  560.         if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
  561.             archive_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]
  562.         else:
  563.             archive_file = pretrained_model_name_or_path
  564.         # redirect to the cache, if necessary
  565.         try:
  566.             resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)
  567.         except EnvironmentError:
  568.             logger.error(
  569.                 "Model name '{}' was not found in model name list ({}). "
  570.                 "We assumed '{}' was a path or url but couldn't find any file "
  571.                 "associated to this path or url.".format(
  572.                     pretrained_model_name_or_path,
  573.                     ', '.join(PRETRAINED_MODEL_ARCHIVE_MAP.keys()),
  574.                     archive_file))
  575.             return None
  576.         if resolved_archive_file == archive_file:
  577.             logger.info("loading archive file {}".format(archive_file))
  578.         else:
  579.             logger.info("loading archive file {} from cache at {}".format(
  580.                 archive_file, resolved_archive_file))
  581.         tempdir = None
  582.         if os.path.isdir(resolved_archive_file) or from_tf:
  583.             serialization_dir = resolved_archive_file
  584.         else:
  585.             # Extract archive to temp dir
  586.             tempdir = tempfile.mkdtemp()
  587.             logger.info("extracting archive file {} to temp dir {}".format(
  588.                 resolved_archive_file, tempdir))
  589.             with tarfile.open(resolved_archive_file, 'r:gz') as archive:
  590.                 archive.extractall(tempdir)
  591.             serialization_dir = tempdir
  592.         # Load config
  593.         config_file = os.path.join(serialization_dir, CONFIG_NAME)
  594.         if not os.path.exists(config_file):
  595.             # Backward compatibility with old naming format
  596.             config_file = os.path.join(serialization_dir, BERT_CONFIG_NAME)
  597.         config = BertConfig.from_json_file(config_file)
  598.         logger.info("Model config {}".format(config))
  599.         # Instantiate model.
  600.         model = cls(config, *inputs, **kwargs)
  601.         if state_dict is None and not from_tf:
  602.             weights_path = os.path.join(serialization_dir, WEIGHTS_NAME)
  603.             state_dict = torch.load(weights_path, map_location='cpu')
  604.         if tempdir:
  605.             # Clean up temp dir
  606.             shutil.rmtree(tempdir)
  607.         if from_tf:
  608.             # Directly load from a TensorFlow checkpoint
  609.             weights_path = os.path.join(serialization_dir, TF_WEIGHTS_NAME)
  610.             return load_tf_weights_in_bert(model, weights_path)
  611.         # Load from a PyTorch state_dict
  612.         old_keys = []
  613.         new_keys = []
  614.         for key in state_dict.keys():
  615.             new_key = None
  616.             if 'gamma' in key:
  617.                 new_key = key.replace('gamma', 'weight')
  618.             if 'beta' in key:
  619.                 new_key = key.replace('beta', 'bias')
  620.             if new_key:
  621.                 old_keys.append(key)
  622.                 new_keys.append(new_key)
  623.         for old_key, new_key in zip(old_keys, new_keys):
  624.             state_dict[new_key] = state_dict.pop(old_key)
  625.  
  626.         missing_keys = []
  627.         unexpected_keys = []
  628.         error_msgs = []
  629.         # copy state_dict so _load_from_state_dict can modify it
  630.         metadata = getattr(state_dict, '_metadata', None)
  631.         state_dict = state_dict.copy()
  632.         if metadata is not None:
  633.             state_dict._metadata = metadata
  634.  
  635.         def load(module, prefix=''):
  636.             local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
  637.             module._load_from_state_dict(
  638.                 state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
  639.             for name, child in module._modules.items():
  640.                 if child is not None:
  641.                     load(child, prefix + name + '.')
  642.         start_prefix = ''
  643.         if not hasattr(model, 'bert') and any(s.startswith('bert.') for s in state_dict.keys()):
  644.             start_prefix = 'bert.'
  645.         load(model, prefix=start_prefix)
  646.         if len(missing_keys) > 0:
  647.             logger.info("Weights of {} not initialized from pretrained model: {}".format(
  648.                 model.__class__.__name__, missing_keys))
  649.         if len(unexpected_keys) > 0:
  650.             logger.info("Weights from pretrained model not used in {}: {}".format(
  651.                 model.__class__.__name__, unexpected_keys))
  652.         if len(error_msgs) > 0:
  653.             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
  654.                                model.__class__.__name__, "\n\t".join(error_msgs)))
  655.         return model
  656.  
  657.  
  658. class BertModel(BertPreTrainedModel):
  659.     """BERT model ("Bidirectional Embedding Representations from a Transformer").
  661.     Params:
  662.         config: a BertConfig class instance with the configuration to build a new model
  664.     Inputs:
  665.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  666.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  667.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  668.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  669.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  670.             a `sentence B` token (see BERT paper for more details).
  671.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  672.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  673.             input sequence length in the current batch. It's the mask that we typically use for attention when
  674.             a batch has varying length sentences.
  675.         `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
  677.     Outputs: Tuple of (encoded_layers, pooled_output)
  678.         `encoded_layers`: controled by `output_all_encoded_layers` argument:
  679.             - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
  680.                 of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
  681.                 encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
  682.             - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
  683.                 to the last attention block of shape [batch_size, sequence_length, hidden_size],
  684.         `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
  685.             classifier pretrained on top of the hidden state associated to the first character of the
  686.             input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
  688.     Example usage:
  689.     ```python
  690.     # Already been converted into WordPiece token ids
  691.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  692.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  693.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  695.     config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  696.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  698.     model = modeling.BertModel(config=config)
  699.     all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
  700.     ```
  701.     """
  702.     def __init__(self, config):
  703.         super(BertModel, self).__init__(config)
  704.         self.embeddings = BertEmbeddings(config)
  705.         self.encoder = BertEncoder(config)
  706.         self.pooler = BertPooler(config)
  707.         self.apply(self.init_bert_weights)
  708.  
  709.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, output_all_encoded_layers=True):
  710.         if attention_mask is None:
  711.             attention_mask = torch.ones_like(input_ids)
  712.         if token_type_ids is None:
  713.             token_type_ids = torch.zeros_like(input_ids)
  714.  
  715.         # We create a 3D attention mask from a 2D tensor mask.
  716.         # Sizes are [batch_size, 1, 1, to_seq_length]
  717.         # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
  718.         # this attention mask is more simple than the triangular masking of causal attention
  719.         # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
  720.         extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
  721.  
  722.         # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
  723.         # masked positions, this operation will create a tensor which is 0.0 for
  724.         # positions we want to attend and -10000.0 for masked positions.
  725.         # Since we are adding it to the raw scores before the softmax, this is
  726.         # effectively the same as removing these entirely.
  727.         extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
  728.         extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
  729.  
  730.         embedding_output = self.embeddings(input_ids, token_type_ids)
  731.         encoded_layers = self.encoder(embedding_output,
  732.                                       extended_attention_mask,
  733.                                       output_all_encoded_layers=output_all_encoded_layers)
  734.         sequence_output = encoded_layers[-1]
  735.         pooled_output = self.pooler(sequence_output)
  736.         if not output_all_encoded_layers:
  737.             encoded_layers = encoded_layers[-1]
  738.         return encoded_layers, pooled_output
  739.  
  740.  
  741. class BertForPreTraining(BertPreTrainedModel):
  742.     """BERT model with pre-training heads.
  743.     This module comprises the BERT model followed by the two pre-training heads:
  744.         - the masked language modeling head, and
  745.         - the next sentence classification head.
  747.     Params:
  748.         config: a BertConfig class instance with the configuration to build a new model.
  750.     Inputs:
  751.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  752.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  753.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  754.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  755.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  756.             a `sentence B` token (see BERT paper for more details).
  757.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  758.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  759.             input sequence length in the current batch. It's the mask that we typically use for attention when
  760.             a batch has varying length sentences.
  761.         `masked_lm_labels`: optional masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length]
  762.             with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
  763.             is only computed for the labels set in [0, ..., vocab_size]
  764.         `next_sentence_label`: optional next sentence classification loss: torch.LongTensor of shape [batch_size]
  765.             with indices selected in [0, 1].
  766.             0 => next sentence is the continuation, 1 => next sentence is a random sentence.
  768.     Outputs:
  769.         if `masked_lm_labels` and `next_sentence_label` are not `None`:
  770.             Outputs the total_loss which is the sum of the masked language modeling loss and the next
  771.             sentence classification loss.
  772.         if `masked_lm_labels` or `next_sentence_label` is `None`:
  773.             Outputs a tuple comprising
  774.             - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
  775.             - the next sentence classification logits of shape [batch_size, 2].
  777.     Example usage:
  778.     ```python
  779.     # Already been converted into WordPiece token ids
  780.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  781.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  782.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  784.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  785.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  787.     model = BertForPreTraining(config)
  788.     masked_lm_logits_scores, seq_relationship_logits = model(input_ids, token_type_ids, input_mask)
  789.     ```
  790.     """
  791.     def __init__(self, config):
  792.         super(BertForPreTraining, self).__init__(config)
  793.         self.bert = BertModel(config)
  794.         self.cls = BertPreTrainingHeads(config, self.bert.embeddings.word_embeddings.weight)
  795.         self.apply(self.init_bert_weights)
  796.  
  797.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None, next_sentence_label=None):
  798.         sequence_output, pooled_output = self.bert(input_ids, token_type_ids, attention_mask,
  799.                                                    output_all_encoded_layers=False)
  800.         prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
  801.  
  802.         if masked_lm_labels is not None and next_sentence_label is not None:
  803.             loss_fct = CrossEntropyLoss(ignore_index=-1)
  804.             masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
  805.             next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
  806.             total_loss = masked_lm_loss + next_sentence_loss
  807.             return total_loss
  808.         else:
  809.             return prediction_scores, seq_relationship_score
  810.  
  811.  
  812. class BertForMaskedLM(BertPreTrainedModel):
  813.     """BERT model with the masked language modeling head.
  814.     This module comprises the BERT model followed by the masked language modeling head.
  816.     Params:
  817.         config: a BertConfig class instance with the configuration to build a new model.
  819.     Inputs:
  820.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  821.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  822.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  823.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  824.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  825.             a `sentence B` token (see BERT paper for more details).
  826.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  827.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  828.             input sequence length in the current batch. It's the mask that we typically use for attention when
  829.             a batch has varying length sentences.
  830.         `masked_lm_labels`: masked language modeling labels: torch.LongTensor of shape [batch_size, sequence_length]
  831.             with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
  832.             is only computed for the labels set in [0, ..., vocab_size]
  834.     Outputs:
  835.         if `masked_lm_labels` is  not `None`:
  836.             Outputs the masked language modeling loss.
  837.         if `masked_lm_labels` is `None`:
  838.             Outputs the masked language modeling logits of shape [batch_size, sequence_length, vocab_size].
  840.     Example usage:
  841.     ```python
  842.     # Already been converted into WordPiece token ids
  843.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  844.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  845.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  847.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  848.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  850.     model = BertForMaskedLM(config)
  851.     masked_lm_logits_scores = model(input_ids, token_type_ids, input_mask)
  852.     ```
  853.     """
  854.     def __init__(self, config):
  855.         super(BertForMaskedLM, self).__init__(config)
  856.         self.bert = BertModel(config)
  857.         self.cls = BertOnlyMLMHead(config, self.bert.embeddings.word_embeddings.weight)
  858.         self.apply(self.init_bert_weights)
  859.  
  860.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None):
  861.         sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask,
  862.                                        output_all_encoded_layers=False)
  863.         prediction_scores = self.cls(sequence_output)
  864.  
  865.         if masked_lm_labels is not None:
  866.             loss_fct = CrossEntropyLoss(ignore_index=-1)
  867.             masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
  868.             return masked_lm_loss
  869.         else:
  870.             return prediction_scores
  871.  
  872.  
  873. class BertForNextSentencePrediction(BertPreTrainedModel):
  874.     """BERT model with next sentence prediction head.
  875.     This module comprises the BERT model followed by the next sentence classification head.
  877.     Params:
  878.         config: a BertConfig class instance with the configuration to build a new model.
  880.     Inputs:
  881.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  882.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  883.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  884.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  885.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  886.             a `sentence B` token (see BERT paper for more details).
  887.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  888.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  889.             input sequence length in the current batch. It's the mask that we typically use for attention when
  890.             a batch has varying length sentences.
  891.         `next_sentence_label`: next sentence classification loss: torch.LongTensor of shape [batch_size]
  892.             with indices selected in [0, 1].
  893.             0 => next sentence is the continuation, 1 => next sentence is a random sentence.
  895.     Outputs:
  896.         if `next_sentence_label` is not `None`:
  897.             Outputs the total_loss which is the sum of the masked language modeling loss and the next
  898.             sentence classification loss.
  899.         if `next_sentence_label` is `None`:
  900.             Outputs the next sentence classification logits of shape [batch_size, 2].
  902.     Example usage:
  903.     ```python
  904.     # Already been converted into WordPiece token ids
  905.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  906.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  907.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  909.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  910.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  912.     model = BertForNextSentencePrediction(config)
  913.     seq_relationship_logits = model(input_ids, token_type_ids, input_mask)
  914.     ```
  915.     """
  916.     def __init__(self, config):
  917.         super(BertForNextSentencePrediction, self).__init__(config)
  918.         self.bert = BertModel(config)
  919.         self.cls = BertOnlyNSPHead(config)
  920.         self.apply(self.init_bert_weights)
  921.  
  922.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, next_sentence_label=None):
  923.         _, pooled_output = self.bert(input_ids, token_type_ids, attention_mask,
  924.                                      output_all_encoded_layers=False)
  925.         seq_relationship_score = self.cls( pooled_output)
  926.  
  927.         if next_sentence_label is not None:
  928.             loss_fct = CrossEntropyLoss(ignore_index=-1)
  929.             next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
  930.             return next_sentence_loss
  931.         else:
  932.             return seq_relationship_score
  933.  
  934.  
  935. class BertForSequenceClassification(BertPreTrainedModel):
  936.     """BERT model for classification.
  937.     This module is composed of the BERT model with a linear layer on top of
  938.     the pooled output.
  940.     Params:
  941.         `config`: a BertConfig class instance with the configuration to build a new model.
  942.         `num_labels`: the number of classes for the classifier. Default = 2.
  944.     Inputs:
  945.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  946.             with the word token indices in the vocabulary. Items in the batch should begin with the special "CLS" token. (see the tokens preprocessing logic in the scripts
  947.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  948.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  949.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  950.             a `sentence B` token (see BERT paper for more details).
  951.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  952.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  953.             input sequence length in the current batch. It's the mask that we typically use for attention when
  954.             a batch has varying length sentences.
  955.         `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
  956.             with indices selected in [0, ..., num_labels].
  958.     Outputs:
  959.         if `labels` is not `None`:
  960.             Outputs the CrossEntropy classification loss of the output with the labels.
  961.         if `labels` is `None`:
  962.             Outputs the classification logits of shape [batch_size, num_labels].
  964.     Example usage:
  965.     ```python
  966.     # Already been converted into WordPiece token ids
  967.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  968.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  969.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  971.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  972.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  974.     num_labels = 2
  976.     model = BertForSequenceClassification(config, num_labels)
  977.     logits = model(input_ids, token_type_ids, input_mask)
  978.     ```
  979.     """
  980.     def __init__(self, config, num_labels):
  981.         super(BertForSequenceClassification, self).__init__(config)
  982.         self.num_labels = num_labels
  983.         self.bert = BertModel(config)
  984.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  985.         self.classifier = nn.Linear(config.hidden_size, num_labels)
  986.         self.apply(self.init_bert_weights)
  987.  
  988.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
  989.         _, pooled_output = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
  990.         pooled_output = self.dropout(pooled_output)
  991.         logits = self.classifier(pooled_output)
  992.  
  993.         if labels is not None:
  994.             loss_fct = CrossEntropyLoss()
  995.             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  996.             return loss
  997.         else:
  998.             return logits
  999.  
  1000.  
  1001. class BertForMultipleChoice(BertPreTrainedModel):
  1002.     """BERT model for multiple choice tasks.
  1003.     This module is composed of the BERT model with a linear layer on top of
  1004.     the pooled output.
  1006.     Params:
  1007.         `config`: a BertConfig class instance with the configuration to build a new model.
  1008.         `num_choices`: the number of classes for the classifier. Default = 2.
  1010.     Inputs:
  1011.         `input_ids`: a torch.LongTensor of shape [batch_size, num_choices, sequence_length]
  1012.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  1013.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  1014.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, num_choices, sequence_length]
  1015.             with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A`
  1016.             and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
  1017.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, num_choices, sequence_length] with indices
  1018.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  1019.             input sequence length in the current batch. It's the mask that we typically use for attention when
  1020.             a batch has varying length sentences.
  1021.         `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
  1022.             with indices selected in [0, ..., num_choices].
  1024.     Outputs:
  1025.         if `labels` is not `None`:
  1026.             Outputs the CrossEntropy classification loss of the output with the labels.
  1027.         if `labels` is `None`:
  1028.             Outputs the classification logits of shape [batch_size, num_labels].
  1030.     Example usage:
  1031.     ```python
  1032.     # Already been converted into WordPiece token ids
  1033.     input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]], [[12, 16, 42], [14, 28, 57]]])
  1034.     input_mask = torch.LongTensor([[[1, 1, 1], [1, 1, 0]],[[1,1,0], [1, 0, 0]]])
  1035.     token_type_ids = torch.LongTensor([[[0, 0, 1], [0, 1, 0]],[[0, 1, 1], [0, 0, 1]]])
  1036.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  1037.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  1039.     num_choices = 2
  1041.     model = BertForMultipleChoice(config, num_choices)
  1042.     logits = model(input_ids, token_type_ids, input_mask)
  1043.     ```
  1044.     """
  1045.     def __init__(self, config, num_choices):
  1046.         super(BertForMultipleChoice, self).__init__(config)
  1047.         self.num_choices = num_choices
  1048.         self.bert = BertModel(config)
  1049.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  1050.         self.classifier = nn.Linear(config.hidden_size, 1)
  1051.         self.apply(self.init_bert_weights)
  1052.  
  1053.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
  1054.         flat_input_ids = input_ids.view(-1, input_ids.size(-1))
  1055.         flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
  1056.         flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
  1057.         _, pooled_output = self.bert(flat_input_ids, flat_token_type_ids, flat_attention_mask, output_all_encoded_layers=False)
  1058.         pooled_output = self.dropout(pooled_output)
  1059.         logits = self.classifier(pooled_output)
  1060.         reshaped_logits = logits.view(-1, self.num_choices)
  1061.  
  1062.         if labels is not None:
  1063.             loss_fct = CrossEntropyLoss()
  1064.             loss = loss_fct(reshaped_logits, labels)
  1065.             return loss
  1066.         else:
  1067.             return reshaped_logits
  1068.  
  1069.  
  1070. class BertForTokenClassification(BertPreTrainedModel):
  1071.     """BERT model for token-level classification.
  1072.     This module is composed of the BERT model with a linear layer on top of
  1073.     the full hidden state of the last layer.
  1075.     Params:
  1076.         `config`: a BertConfig class instance with the configuration to build a new model.
  1077.         `num_labels`: the number of classes for the classifier. Default = 2.
  1079.     Inputs:
  1080.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  1081.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  1082.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  1083.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  1084.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  1085.             a `sentence B` token (see BERT paper for more details).
  1086.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  1087.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  1088.             input sequence length in the current batch. It's the mask that we typically use for attention when
  1089.             a batch has varying length sentences.
  1090.         `labels`: labels for the classification output: torch.LongTensor of shape [batch_size, sequence_length]
  1091.             with indices selected in [0, ..., num_labels].
  1093.     Outputs:
  1094.         if `labels` is not `None`:
  1095.             Outputs the CrossEntropy classification loss of the output with the labels.
  1096.         if `labels` is `None`:
  1097.             Outputs the classification logits of shape [batch_size, sequence_length, num_labels].
  1099.     Example usage:
  1100.     ```python
  1101.     # Already been converted into WordPiece token ids
  1102.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  1103.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  1104.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  1106.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  1107.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  1109.     num_labels = 2
  1111.     model = BertForTokenClassification(config, num_labels)
  1112.     logits = model(input_ids, token_type_ids, input_mask)
  1113.     ```
  1114.     """
  1115.     def __init__(self, config, num_labels):
  1116.         super(BertForTokenClassification, self).__init__(config)
  1117.         self.num_labels = num_labels
  1118.         self.bert = BertModel(config)
  1119.         self.dropout = nn.Dropout(config.hidden_dropout_prob)
  1120.         self.classifier = nn.Linear(config.hidden_size, num_labels)
  1121.         self.apply(self.init_bert_weights)
  1122.  
  1123.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
  1124.         sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
  1125.         sequence_output = self.dropout(sequence_output)
  1126.         logits = self.classifier(sequence_output)
  1127.  
  1128.         if labels is not None:
  1129.             loss_fct = CrossEntropyLoss()
  1130.             # Only keep active parts of the loss
  1131.             if attention_mask is not None:
  1132.                 active_loss = attention_mask.view(-1) == 1
  1133.                 active_logits = logits.view(-1, self.num_labels)[active_loss]
  1134.                 active_labels = labels.view(-1)[active_loss]
  1135.                 loss = loss_fct(active_logits, active_labels)
  1136.             else:
  1137.                 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
  1138.             return loss
  1139.         else:
  1140.             return logits
  1141.  
  1142.  
  1143. class BertForQuestionAnswering(BertPreTrainedModel):
  1144.     """BERT model for Question Answering (span extraction).
  1145.     This module is composed of the BERT model with a linear layer on top of
  1146.     the sequence output that computes start_logits and end_logits
  1148.     Params:
  1149.         `config`: a BertConfig class instance with the configuration to build a new model.
  1151.     Inputs:
  1152.         `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
  1153.             with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
  1154.             `extract_features.py`, `run_classifier.py` and `run_squad.py`)
  1155.         `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
  1156.             types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
  1157.             a `sentence B` token (see BERT paper for more details).
  1158.         `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
  1159.             selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
  1160.             input sequence length in the current batch. It's the mask that we typically use for attention when
  1161.             a batch has varying length sentences.
  1162.         `start_positions`: position of the first token for the labeled span: torch.LongTensor of shape [batch_size].
  1163.             Positions are clamped to the length of the sequence and position outside of the sequence are not taken
  1164.             into account for computing the loss.
  1165.         `end_positions`: position of the last token for the labeled span: torch.LongTensor of shape [batch_size].
  1166.             Positions are clamped to the length of the sequence and position outside of the sequence are not taken
  1167.             into account for computing the loss.
  1169.     Outputs:
  1170.         if `start_positions` and `end_positions` are not `None`:
  1171.             Outputs the total_loss which is the sum of the CrossEntropy loss for the start and end token positions.
  1172.         if `start_positions` or `end_positions` is `None`:
  1173.             Outputs a tuple of start_logits, end_logits which are the logits respectively for the start and end
  1174.             position tokens of shape [batch_size, sequence_length].
  1176.     Example usage:
  1177.     ```python
  1178.     # Already been converted into WordPiece token ids
  1179.     input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
  1180.     input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
  1181.     token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
  1183.     config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
  1184.         num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
  1186.     model = BertForQuestionAnswering(config)
  1187.     start_logits, end_logits = model(input_ids, token_type_ids, input_mask)
  1188.     ```
  1189.     """
  1190.     def __init__(self, config):
  1191.         super(BertForQuestionAnswering, self).__init__(config)
  1192.         self.bert = BertModel(config)
  1193.         # TODO check with Google if it's normal there is no dropout on the token classifier of SQuAD in the TF version
  1194.         # self.dropout = nn.Dropout(config.hidden_dropout_prob)
  1195.         self.qa_outputs = nn.Linear(config.hidden_size, 2)
  1196.         self.apply(self.init_bert_weights)
  1197.  
  1198.     def forward(self, input_ids, token_type_ids=None, attention_mask=None, start_positions=None, end_positions=None):
  1199.         sequence_output, _ = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
  1200.         logits = self.qa_outputs(sequence_output)
  1201.         start_logits, end_logits = logits.split(1, dim=-1)
  1202.         start_logits = start_logits.squeeze(-1)
  1203.         end_logits = end_logits.squeeze(-1)
  1204.  
  1205.         if start_positions is not None and end_positions is not None:
  1206.             # If we are on multi-GPU, split add a dimension
  1207.             if len(start_positions.size()) > 1:
  1208.                 start_positions = start_positions.squeeze(-1)
  1209.             if len(end_positions.size()) > 1:
  1210.                 end_positions = end_positions.squeeze(-1)
  1211.             # sometimes the start/end positions are outside our model inputs, we ignore these terms
  1212.             ignored_index = start_logits.size(1)
  1213.             start_positions.clamp_(0, ignored_index)
  1214.             end_positions.clamp_(0, ignored_index)
  1215.  
  1216.             loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
  1217.             start_loss = loss_fct(start_logits, start_positions)
  1218.             end_loss = loss_fct(end_logits, end_positions)
  1219.             total_loss = (start_loss + end_loss) / 2
  1220.             return total_loss
  1221.         else:
  1222.             return start_logits, end_logits

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/312372
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号