def _string_to_tokens_dataset_mapper(keys_to_map, suffix="_tok"):
  """Wrapper for mapper that tokenizes and truncates by length."""

  def _mapper(dataset):
    """Tokenizes strings using tf.string_split and truncates by length."""
    for k in keys_to_map:
      # pylint: disable=g-explicit-length-test
      if len(dataset[k].get_shape()) == 0:  # Used for questions.
        # pylint: enable=g-explicit-length-test
        # <string> [num_tokens]
        tokens = tf.string_split([dataset[k]]).values
      else:  # Used for contexts.
        # <string> [num_context, num_tokens] (sparse)
        sparse_tokens = tf.string_split(dataset[k])

        # <string>[num_tokens, max_num_tokens] (dense)
        tokens = tf.sparse_tensor_to_dense(sparse_tokens, default_value="")

      dataset[k + suffix] = tokens
      # Compute exact length of each context.
      dataset[k + suffix + "_len"] = tf.count_nonzero(
          tokens, axis=-1, dtype=tf.int32)
    return dataset

  return _mapper
Wrapper for mapper that tokenizes and truncates by length.
language/question_answering/decatt_docreader/models/nq_long_model.py