def compute_bert_cls_embs(self, token_ids, segment_ids,
                            masks):
    """Return the embedding of CLS token in BERT.

    Args:
      token_ids: BERT token ids
      segment_ids: BERT segment ids
      masks: BERT mask

    Returns:
      BERT CLS embedding
    """
    bert_outputs = self.bert_module(
        inputs={
            'input_ids': token_ids,
            'segment_ids': segment_ids,
            'input_mask': masks
        },
        signature='tokens',
        as_dict=True)
    cls_embs = bert_outputs['sequence_output']
    # batch_size, max_seq_length, hidden_size
    cls_embs = cls_embs[:, 0, :]
    # batch_size, hidden_size
    return cls_embs
Return the embedding of CLS token in BERT.

    Args:
      token_ids: BERT token ids
      segment_ids: BERT segment ids
      masks: BERT mask

    Returns:
      BERT CLS embedding
    
language/emql/model.py