code: def _create_dataset(
      file_pattern: str,
      batch_size: int,
      randomize_input: bool = True,
      num_epochs: Optional[int] = None,
      drop_final_batch: bool = False,
      reader_num_threads: int = 10,
  ) -> tf.data.Dataset:
      """Returns `tf.data.Dataset` for training or evaluating the model.
  
      Args:
        file_pattern: (string) file pattern for input data.
        batch_size: (int) number of input examples to process per batch.
        randomize_input: (bool) if true, randomize input example order. It should
          almost always be true except for unittest/debug purposes.
        num_epochs: (int) Number of times the input dataset must be repeated. None
          to repeat the data indefinitely.
        drop_final_batch: (boolean) Whether to drop the last batch if there are
          fewer than batch_size elements left.
        reader_num_threads: (int) Number of parallel threads to read from.
  
      Returns:
        A `tf.data.Dataset` of features dictionary and label.
      """
      context_feature_spec, example_feature_spec = _create_feature_spec(
          with_label=True
      )
      dataset = tfr.data.build_ranking_dataset(
          data_format=tfr.data.ELWC,
          file_pattern=file_pattern,
          batch_size=batch_size,
          list_size=FLAGS.list_size,
          context_feature_spec=context_feature_spec,
          example_feature_spec=example_feature_spec,
          reader=tf.data.TFRecordDataset,
          num_epochs=num_epochs,
          shuffle=randomize_input,
          reader_num_threads=reader_num_threads,
          drop_final_batch=drop_final_batch,
          mask_feature_name=_MASK,
      )
  
      def _add_label(features):
          label = tf.squeeze(features.pop(_LABEL_FEATURE), axis=2)
          label = tf.cast(label, dtype=tf.float32)
          return features, label
  
      dataset = dataset.map(_add_label)
      return dataset
  
file_path: ranking/tensorflow_ranking/examples/keras/antique_kpl_din.py
imports: import os
  from typing import Dict, Optional, Tuple, Union
  from absl import app
  from absl import flags
  from absl import logging
  import tensorflow
  import tensorflow_ranking
label_desc: Use image node_id 2 for calling a global function with example usage: _create_dataset(file_pattern, batch_size, randomize_input, num_epochs, drop_final_batch, reader_num_threads) and returns: dataset, features,  label