def get_doc_embeddings(self, docs):
    """Run BERT to get features for docs.

    Args:
      docs: List of list of tokens.

    Returns:
      embeddings: Numpy array of token features.
    """
    num_batches = (len(docs) // self.batch_size) + 1
    tf.logging.info("Total batches for BERT = %d", num_batches)
    embeddings = np.empty((len(docs), self.max_seq_length, self.emb_dim),
                          dtype=np.float32)
    for nb in tqdm(range(num_batches)):
      min_ = nb * self.batch_size
      max_ = (nb + 1) * self.batch_size
      if min_ >= len(docs):
        break
      if max_ > len(docs):
        max_ = len(docs)
      current_features = [
          self.get_features(docs[ii], ["dummy"], ["dummy"], ii)
          for ii in range(min_, max_)
      ]
      results = self._run_on_features(current_features)
      for ir, rr in enumerate(results):
        embeddings[min_ + ir, :, :] = rr["doc_features"]
    return embeddings[:, 1:, :]  # remove [CLS]
Run BERT to get features for docs.

    Args:
      docs: List of list of tokens.

    Returns:
      embeddings: Numpy array of token features.
    
language/labs/drkit/bert_utils.py