def _bert_stack(self, molecules_in: tf.Tensor,
                  attention_mask: tf.Tensor) -> Sequence[tf.Tensor]:
    """Encode the molecules using a deep transformer stack."""
    with tf.variable_scope("bert"):
      return bert_modeling.transformer_model(
          input_tensor=molecules_in,
          attention_mask=attention_mask,
          hidden_size=self.config.hidden_size,
          num_hidden_layers=self.config.num_hidden_layers,
          num_attention_heads=self.config.num_attention_heads,
          intermediate_size=self.config.intermediate_size,
          intermediate_act_fn=bert_modeling.get_activation(
              self.config.hidden_act),
          hidden_dropout_prob=self.config.hidden_dropout_prob,
          attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
          initializer_range=self.config.initializer_range,
          do_return_all_layers=True)
Encode the molecules using a deep transformer stack.
language/canine/modeling.py