code: def get_dataset(
      *,
      batch_size: int,
      eval_batch_size: int,
      num_shards: int,
      rng: Any,
      dataset_configs: ml_collections.ConfigDict,
      dtype_str: str = "float32",
      shuffle_seed: int = 0,
      dataset_service_address: Optional[str] = None,
  ) -> dataset_utils.Dataset:
      """Returns generators for image-text detection datasets.
  
      In addition to standard detection features this loader also produces textual
      queries (in features['queries']), which defines the per batch labelset. The
      queries are the (argus) tokenized string queries. Text queries from detection
      datasets are simply their class names.
  
      Args:
        batch_size: Determines the train batch size.
        eval_batch_size: Determines the evaluation batch size.
        num_shards: Unused; determined by CLU.
        rng: JAX RNG key, which can be used for augmentation, shuffling, etc.
        dataset_configs: Dataset configurations.
        dtype_str: Data type of the image. Only 'float32' is currently supported.
        shuffle_seed: Unsupported; use rng instead.
        dataset_service_address: Unsupported; must be None.
  
      Returns:
        A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
        a test_iter, and a dict of meta_data.
      """
      if dataset_service_address is not None:
          raise NotImplementedError(
              "This dataset does not support the DS service."
          )
      if rng is None:
          raise NotImplementedError("This dataset requires a JAX RNG.")
      if shuffle_seed:
          raise NotImplementedError(
              "This dataset requires a JAX RNG, do not use shuffle_seed."
          )
      if len(dataset_configs.eval.tfds_names) > 1:
          raise NotImplementedError(
              "Evaluation with more than one datasets is not supported."
          )
      if dtype_str != "float32":
          raise ValueError(f"Unsupported dtype_str: {dtype_str}")
  
      # Delete unused arguments (see docstring):
      del num_shards, shuffle_seed
  
      # Ensure a different key on each worker:
      rng = jax.random.fold_in(rng, jax.process_index())
  
      rng, train_rng = jax.random.split(rng)
  
      # Training dataset:
      train_config = _validate_and_normalize_config(
          dataset_configs, train=True
      )
      train_ds = _build_pipeline(
          config=train_config,
          batch_size=batch_size,
          rng=train_rng,
          shuffle=True,
      )
  
      example_batch = next(iter(train_ds))
      train_iter = iter(train_ds)
      train_iter = map(dataset_utils.tf_to_numpy, train_iter)
  
      if dataset_configs.get("prefetch_to_device"):
          # Async bind batch to device which speeds up training.
          train_iter = jax_utils.prefetch_to_device(
              train_iter, dataset_configs.get("prefetch_to_device")
          )
  
      rng, eval_rng = jax.random.split(rng)
      # Evaluation dataset:
      eval_config = _validate_and_normalize_config(
          dataset_configs, train=False
      )
      eval_ds = _build_pipeline(
          config=eval_config,
          batch_size=eval_batch_size,
          rng=eval_rng,
      )
      eval_iter = iter(eval_ds)
      eval_iter = map(dataset_utils.tf_to_numpy, eval_iter)
  
      total_examples = sum(
          dataset_utils.get_num_examples(
              name, split, train_config.data_dirs[0].get(name)
          )
          for name, split in zip(
              dataset_configs.train.tfds_names,
              dataset_configs.train.splits,
          )
      )
      eval_data_dir = eval_config.data_dirs[0].get(
          dataset_configs.eval.tfds_names[0]
      )
      total_eval_examples = dataset_utils.get_num_examples(
          dataset_configs.eval.tfds_names[0],
          dataset_configs.eval.splits[0],
          data_dir=eval_data_dir,
      )
  
      eval_builder = tfds.builder(
          dataset_configs.eval.tfds_names[0], data_dir=eval_data_dir
      )
  
      if "bobjects" in eval_builder.info.features:
          num_classes = eval_builder.info.features["bobjects"][
              "label"
          ].num_classes
      else:
          num_classes = eval_builder.info.features["objects"][
              "label"
          ].num_classes
      num_classes += 1
  
      meta_data = {
          "input_shape": (-1,)
          + tuple(example_batch["inputs"].shape[-3:]),
          "query_shape": (-1,)
          + tuple(example_batch["queries"].shape[-2:]),
          "num_train_examples": total_examples,
          "num_eval_examples": total_eval_examples,
          "input_dtype": jnp.float32,
          "target_is_onehot": True,  # We always use the one/multi-hot format.
          "eval_num_classes": num_classes,
      }
  
      return dataset_utils.Dataset(
          train_iter, eval_iter, None, meta_data
      )
  
imports: import dataclasses
  import functools
  from typing import Any, Dict, Optional, Sequence
  from clu import deterministic_data
  from clu import preprocess_spec
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.owl_vit.preprocessing import image_ops
  from scenic.projects.owl_vit.preprocessing import label_ops
  from scenic.projects.owl_vit.preprocessing import mosaic
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 7 for calling a global function with example usage: get_dataset() and returns: dataset_utils