def _pool(self, seq_to_pool: tf.Tensor) -> tf.Tensor:
    """Grab the [CLS] molecule for use in classification tasks."""
    # The "pooler" converts the encoded sequence tensor of shape
    # [batch_size, seq_length, hidden_size] to a tensor of shape
    # [batch_size, 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.
      # This snippet is taken from vanilla BERT.
      first_token_tensor = tf.squeeze(seq_to_pool[:, 0:1, :], axis=1)
      return tf.layers.dense(
          first_token_tensor,
          self.config.hidden_size,
          activation=tf.tanh,
          kernel_initializer=bert_modeling.create_initializer(
              self.config.initializer_range))
Grab the [CLS] molecule for use in classification tasks.
language/canine/modeling.py