def get_qry_embeddings(self, qrys, ents):
    """Run BERT to get features for queries.

    Args:
      qrys: List of list of tokens.
      ents: List of list of tokens.

    Returns:
      st_embeddings: Numpy array of token features.
      en_embeddings: Numpy array of token features.
      bow_embeddings: Numpy array of token features.
    """
    num_batches = (len(qrys) // self.batch_size) + 1
    tf.logging.info("Total batches for BERT = %d", num_batches)
    st_embeddings = np.empty((len(qrys), self.emb_dim), dtype=np.float32)
    en_embeddings = np.empty((len(qrys), self.emb_dim), dtype=np.float32)
    bow_embeddings = np.empty((len(qrys), 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(qrys):
        break
      if max_ > len(qrys):
        max_ = len(qrys)
      current_features = [
          self.get_features(["dummy"], qrys[ii], ents[ii], ii)
          for ii in range(min_, max_)
      ]
      results = self._run_on_features(current_features)
      for ir, rr in enumerate(results):
        st_embeddings[min_ + ir, :] = rr["qry_st_features"]
        en_embeddings[min_ + ir, :] = rr["qry_en_features"]
        bow_embeddings[min_ + ir, :] = rr["qry_bow_features"]
    return st_embeddings, en_embeddings, bow_embeddings
Run BERT to get features for queries.

    Args:
      qrys: List of list of tokens.
      ents: List of list of tokens.

    Returns:
      st_embeddings: Numpy array of token features.
      en_embeddings: Numpy array of token features.
      bow_embeddings: Numpy array of token features.
    
language/labs/drkit/bert_utils_v2.py