class_docstr: BERT model ("Bidirectional Encoder Representations from Transformers").
class_docstr_tok: ["BERT","model","(","``","Bidirectional","Encoder","Representations","from","Transformers","''",")","."]
code: def __init__(
      self,
      config,
      is_training,
      input_ids,
      input_mask=None,
      token_type_ids=None,
      scope=None,
      mesh_shape="",
      layout="",
  ):
      self.config = copy.deepcopy(config)
      del config
      if not is_training:
          self.config.layer_output_dropout_prob = 0.0
          self.config.attention_probs_dropout_prob = 0.0
          self.config.feedforward_intermediate_dropout_prob = 0.0
      input_shape = input_ids.shape
      assert input_shape.ndims == 2
  
      self._seq_dim = input_shape.dims[1]
      self._memory_seq_dim = mtf.Dimension(
          "memory_seq", self.seq_dim.size
      )
      self._extra_losses = []
      mesh = input_ids.mesh
  
      if token_type_ids is None:
          token_type_ids = mtf.zeros(mesh, input_shape, dtype=tf.int32)
  
      with tf.variable_scope(scope, default_name="bert"):
          with tf.variable_scope("embeddings"):
              # Perform embedding lookup on the word ids.
              self.embedding_table = mtf.get_variable(
                  mesh,
                  "word_embeddings",
                  mtf.Shape([self.vocab_dim, self.model_dim]),
                  initializer=self.embedding_initializer,
              )
              self.word_embedding_output = mtf.gather(
                  self.embedding_table, input_ids, self.vocab_dim
              )
  
              # Add positional embeddings and token type embeddings, then layer
              # normalize and perform dropout.
              self.embedding_output = self.word_embedding_output
  
              token_type_table = mtf.get_variable(
                  mesh,
                  "token_type_embeddings",
                  mtf.Shape(
                      [self.token_type_vocab_dim, self.model_dim]
                  ),
                  initializer=self.embedding_initializer,
              )
              if token_type_ids is not None:
                  self.embedding_output += mtf.gather(
                      token_type_table,
                      token_type_ids,
                      self.token_type_vocab_dim,
                  )
              if self.config.position_signal == "embedding":
                  full_position_table = mtf.get_variable(
                      mesh,
                      "position_embeddings",
                      mtf.Shape(
                          [
                              self.max_position_embeddings_dim,
                              self.model_dim,
                          ]
                      ),
                      initializer=self.embedding_initializer,
                  )
                  short_position_table = mtf.rename_dimension(
                      mtf.slice(
                          full_position_table,
                          0,
                          self.seq_dim.size,
                          self.max_position_embeddings_dim.name,
                      ),
                      self.max_position_embeddings_dim.name,
                      self.seq_dim.name,
                  )
                  self.embedding_output += short_position_table
              self.embedding_output = self.normalize(
                  self.embedding_output
              )
              self.embedding_output = mtf.dropout(
                  self.embedding_output,
                  is_training,
                  keep_prob=1.0 - self.config.layer_output_dropout_prob,
              )
  
          with tf.variable_scope("encoder"):
              attention_biases = []
              if input_mask:
                  # [batch_dim, memory_seq_dim]
                  attention_biases.append(
                      (
                          1.0
                          - mtf.to_float(
                              mtf.replace_dimensions(
                                  input_mask,
                                  self.seq_dim,
                                  self.memory_seq_dim,
                              )
                          )
                      )
                      * -10000.0
                  )
              if (
                  self.config.position_signal
                  == "relative_attention_bias"
              ):
                  buckets_dim = mtf.Dimension("buckets", 32)
                  rp_bucket = _relative_position_bucket(
                      mtf.range(mesh, self.memory_seq_dim, tf.int32)
                      - mtf.range(mesh, self.seq_dim, tf.int32),
                      num_buckets=buckets_dim.size,
                  )
                  bias_var = mtf.get_variable(
                      mesh,
                      "relative_attention_bias",
                      [self.num_heads_dim, buckets_dim],
                      initializer=tf.zeros_initializer(),
                  )
                  attention_biases.append(
                      mtf.gather(bias_var, rp_bucket, buckets_dim)
                  )
              attention_bias = mtf.add_n(attention_biases)
              prev_layer_output = self.embedding_output
              self.all_encoder_layers = []
              for block_num in range(self.config.num_blocks):
                  with tf.variable_scope("block_%d" % block_num):
                      for layer_idx, layer_type in enumerate(
                          self.config.block_layers
                      ):
                          layer_name = layer_type
                          count = self.config.block_layers[
                              :layer_idx
                          ].count(layer_type)
                          if count:
                              layer_name += "_%d" % count
                          with tf.variable_scope(layer_name):
                              x = prev_layer_output
                              if (
                                  self.config.residual_structure
                                  == "direct"
                              ):
                                  x = self.normalize(x)
                              if layer_type == "attention":
                                  x = self.self_attention(
                                      x, attention_bias
                                  )
                              elif layer_type == "feedforward":
                                  x = self.feedforward(x)
                              elif layer_type == "moe":
                                  x = self.moe(
                                      x,
                                      layout,
                                      mesh_shape,
                                      input_mask,
                                      is_training,
                                  )
                              else:
                                  raise ValueError(
                                      "unknown layer type " + layer_type
                                  )
                              x = mtf.dropout(
                                  x,
                                  is_training,
                                  keep_prob=1.0
                                  - self.config.layer_output_dropout_prob,
                              )
                              layer_output = prev_layer_output + x
                              if (
                                  self.config.residual_structure
                                  == "original"
                              ):
                                  layer_output = self.normalize(
                                      layer_output
                                  )
                              prev_layer_output = layer_output
                  self.all_encoder_layers.append(layer_output)
  
          self.sequence_output = prev_layer_output
          if self.config.residual_structure == "direct":
              self.sequence_output = self.normalize(
                  self.sequence_output
              )
  
          # The "pooler" converts the encoded sequence tensor of shape
          # [batch_dim, seq_dim, hidden_size] to a tensor of shape
          # [batch_dim, hidden_size]. This is necessary for segment-level
          # (or segment-pair-level) classification tasks where we need a fixed
          # dimensional representation of the segment.
          with tf.variable_scope("pooler"):
              # We "pool" the model by simply taking the hidden state corresponding
              # to the first token. We assume that this has been pre-trained
              first_token_tensor = mtf.gather(
                  self.sequence_output, 0, self.seq_dim
              )
              self.pooled_output = mtf.layers.dense(
                  first_token_tensor,
                  reduced_dims=[self.model_dim],
                  new_dims=[self.model_dim],
                  activation=mtf.tanh,
                  kernel_initializer=self.dense_initializer,
                  use_bias=self.config.use_bias,
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import copy
  import json
  import math
  import re
  import mesh_tensorflow
  import mesh_tensorflow.transformer.moe
  import six
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new BertModel object from inherited base classes: object with example: obj = BertModel(config, is_training, input_ids, input_mask, token_type_ids, scope, mesh_shape, layout)