def token_to_word_id_mapper(self, keys_to_map, suffix="_wid"):
    """Creates a mapping function to augment a tf.data.Dataset with word ids.

    Suppose we have a `dataset` with outputs `str1` and `str2` of arbitrary
    shape. Here is an example usage of this function:

    embeddings = PretrainedWordEmbeddings("/path/to/emb.txt", 10000)
    dataset = dataset.map(embeddings.token_to_word_id_mapper(['str1', 'str2']))

    Now the dataset will also include outputs `str1_wid` and `str2_wid` that
    can be used as features in a model.  The 'str1_wid' feature has the same
    shape as the 'str1' feature, except the string are replaced with their
    int32 word IDs.

    Args:
      keys_to_map: List of strings that are keys for tf.string Tensors to map to
          word ids.
      suffix: String to append to the given keys to indicate the mapped Tensors.

    Returns:
      _mapper: A mapping function that can be used with the tf.data.Dataset API.
    """
    return dataset_utils.string_to_int_mapper(
        keys_to_map=keys_to_map,
        mapping=self._idx2str,
        num_oov_buckets=self._num_oov_buckets,
        suffix=suffix)
Creates a mapping function to augment a tf.data.Dataset with word ids.

    Suppose we have a `dataset` with outputs `str1` and `str2` of arbitrary
    shape. Here is an example usage of this function:

    embeddings = PretrainedWordEmbeddings("/path/to/emb.txt", 10000)
    dataset = dataset.map(embeddings.token_to_word_id_mapper(['str1', 'str2']))

    Now the dataset will also include outputs `str1_wid` and `str2_wid` that
    can be used as features in a model.  The 'str1_wid' feature has the same
    shape as the 'str1' feature, except the string are replaced with their
    int32 word IDs.

    Args:
      keys_to_map: List of strings that are keys for tf.string Tensors to map to
          word ids.
      suffix: String to append to the given keys to indicate the mapped Tensors.

    Returns:
      _mapper: A mapping function that can be used with the tf.data.Dataset API.
    
language/common/inputs/embedding_utils.py