def embed_wiki_dataset(self, dataset):
    """Embed the wikipedia/evidence only dataset.

    Args:
      dataset: The wikipedia only dataset (e.g. wiki_tfds.py)

    Returns:
      Aligned wikipedia_urls, sentence_ids, and embeddings of model
    """

    # map_fn and tf_map_fn transform the dataset to the same format as
    # tfds_evidence/the one the model expects
    def map_fn(text, wikipedia_url, sentence_id):
      return ('a', text, wikipedia_url, str(sentence_id),
              json.dumps({
                  'sentence_id': int(sentence_id.numpy()),
                  'wikipedia_url': wikipedia_url.numpy().decode('utf8')
              }))

    def tf_map_fn(example):
      tensors = tf.py_function(
          map_fn,
          inp=[
              example['text'], example['wikipedia_url'], example['sentence_id']
          ],
          Tout=(tf.string, tf.string, tf.string, tf.string, tf.string))
      return {
          'claim_text': tensors[0],
          'evidence_text': tensors[1],
          'wikipedia_url': tensors[2],
          'sentence_id': tensors[3],
          'claim_label': tf.constant(0, dtype=tf.int64),
          'evidence_label': tf.constant(0, dtype=tf.int64),
          'metadata': tensors[4]
      }

    formatted_ds = dataset.map(tf_map_fn)
    batched_examples = self._encode_and_batch(
        formatted_ds, filter_claims=False, filter_evidence=False)
    preds = []
    wikipedia_urls = []
    sentence_ids = []
    for batch in tqdm.tqdm(batched_examples, mininterval=5):
      # model.predict() is broken after model load due to missing shapes, so
      # have to do our own batching/unbatching.
      inputs, _ = batch
      _, ev_encoding = self._inner_model(
          inputs, embed_claim=False, embed_evidence=True)
      for m in inputs['metadata'].numpy():
        key = json.loads(m.decode('utf8'))
        wikipedia_urls.append(key['wikipedia_url'])
        sentence_ids.append(key['sentence_id'])
      preds.append(ev_encoding)

    return np.array(wikipedia_urls), np.array(sentence_ids), np.vstack(preds)
Embed the wikipedia/evidence only dataset.

    Args:
      dataset: The wikipedia only dataset (e.g. wiki_tfds.py)

    Returns:
      Aligned wikipedia_urls, sentence_ids, and embeddings of model
    
language/serene/training.py