def get_transformer_encoder(bert_config,
                            type_vocab_sizes,
                            num_float_features,
                            sequence_length,
                            bert_init_ckpt=None):
  """Gets a 'TransformerEncoder' object.

  Args:
    bert_config: A 'modeling.BertConfig' or 'modeling.AlbertConfig' object.
    type_vocab_sizes: Type vocab sizes.
    num_float_features: Number of float features.
    sequence_length: Maximum sequence length of the training data.
    bert_init_ckpt: str, Location of the initial checkpoint.

  Returns:
    A networks.BertEncoder object.
  """

  type_vocab_sizes = type_vocab_sizes or ()

  kwargs = dict(
      vocab_size=bert_config.vocab_size,
      hidden_size=bert_config.hidden_size,
      num_layers=bert_config.num_hidden_layers,
      num_attention_heads=bert_config.num_attention_heads,
      intermediate_size=bert_config.intermediate_size,
      activation=tf_utils.get_activation(bert_config.hidden_act),
      dropout_rate=bert_config.hidden_dropout_prob,
      attention_dropout_rate=bert_config.attention_probs_dropout_prob,
      sequence_length=sequence_length,
      max_sequence_length=bert_config.max_position_embeddings,
      type_vocab_sizes=type_vocab_sizes,
      num_float_features=num_float_features,
      initializer=tf.keras.initializers.TruncatedNormal(
          stddev=bert_config.initializer_range),
      bert_init_ckpt=bert_init_ckpt,
  )
  assert isinstance(bert_config, configs.BertConfig)
  return TransformerEncoder(**kwargs)
Gets a 'TransformerEncoder' object.

  Args:
    bert_config: A 'modeling.BertConfig' or 'modeling.AlbertConfig' object.
    type_vocab_sizes: Type vocab sizes.
    num_float_features: Number of float features.
    sequence_length: Maximum sequence length of the training data.
    bert_init_ckpt: str, Location of the initial checkpoint.

  Returns:
    A networks.BertEncoder object.
  
language/search_agents/muzero/transformer_encoder.py