code: def create_coverage_based_dataset(
      filenames: str,
      with_label: bool = True,
  ) -> tf.data.Dataset:
      """Creates a coverage based pileup dataset from a filepath.
  
      Args:
        filenames: The data directory/pattern containing data files.
        with_label: whether to load the labels or not.
  
      Returns:
        tf.data.Dataset
      """
      logging.info("Finding all data files matching the file pattern.")
      dataset_files = tf.io.matching_files(filenames)
  
      # Extract image shape from the first example
      shape = _extract_metadata_from_first_example(dataset_files[0])
  
      proto_features = {
          "variant/encoded": tf.io.FixedLenFeature(
              shape=(), dtype=tf.string
          ),
          "image/encoded": tf.io.FixedLenFeature(
              shape=(), dtype=tf.string
          ),
          "image/shape": tf.io.FixedLenFeature(
              shape=(len(shape),), dtype=tf.int64
          ),
      }
      if with_label:
          proto_features["label"] = tf.io.FixedLenFeature(
              shape=1, dtype=tf.int64
          )
  
      def _process_input(proto_string):
          """Helper function for input function that parses a serialized example."""
  
          parsed_features = tf.io.parse_single_example(
              serialized=proto_string, features=proto_features
          )
  
          features = {
              "image": tf.io.parse_tensor(
                  parsed_features["image/encoded"], tf.uint8
              ),
          }
  
          features["image"].set_shape(shape)
  
          if with_label:
              return features, parsed_features["label"]
          else:
              return features, None
  
      compression = _filename_to_compression(dataset_files[0])
  
      logging.info("Loading TFRecords as bytes.")
      dataset = tf.data.Dataset.from_tensor_slices(dataset_files)
  
      # pylint: disable=g-long-lambda
      # interleave parallelizes the data loading step by interleaving the I/O
      # operation to read the file. It speeds up the I/O step.
      dataset = dataset.interleave(
          lambda filename: tf.data.TFRecordDataset(
              filename,
              compression_type=compression,
          ).map(
              _process_input,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          ),
          cycle_length=len(dataset_files),
      )
  
      return dataset
  
file_path: scenic/scenic/projects/svvit/datasets/pileup_coverage_dataset.py
imports: import functools
  from typing import Optional, Tuple
  from absl import logging
  from flax import jax_utils
  import jax
  import jax.numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: create_coverage_based_dataset(filenames, with_label) and returns: dataset, features,  parsed_features, features,  None