code: def ensure_dataset_eos(dataset, feature_keys=None):
      """Replaces the final token of features with EOS=1 if it is not PAD=0.
  
      Args:
        dataset: a tf.data.Dataset
        feature_keys: (optional) collection of strings, the feature names to ensure
          end with EOS or padding. Defaults to all features.
      Returns:
        a tf.data.Dataset where all specified features end with PAD=0 or EOS=1.
      """
      feature_keys = feature_keys or set(
          tf.data.get_output_shapes(dataset).keys()
      )
  
      def _ensure_eos(k, v):
          if k not in feature_keys:
              return v
          return tf.concat(
              [v[0:-1], tf.clip_by_value(v[-1:], 0, 1)], axis=0
          )
  
      return dataset.map(
          lambda ex: {k: _ensure_eos(k, v) for k, v in ex.items()},
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import gin
  import tensorflow.compat.v1
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: ensure_dataset_eos(dataset, feature_keys) and returns: dataset, tf, v