code: def get_wmt_datasets(
      config: ml_collections.ConfigDict,
      *,
      n_devices: int,
      reverse_translation: bool = True,
      vocab_path: Optional[str] = None,
  ):
      """Load and return dataset of batched examples for use during training."""
      if vocab_path is None:
          vocab_path = os.path.expanduser("~/wmt_sentencepiece_model")
  
      train_ds_builder = tfds.builder(config.dataset_name)
      train_data = get_raw_dataset(
          train_ds_builder,
          "train",
          reverse_translation=reverse_translation,
      )
  
      if config.eval_dataset_name:
          eval_ds_builder = tfds.builder(config.eval_dataset_name)
      else:
          eval_ds_builder = train_ds_builder
      eval_data = get_raw_dataset(
          eval_ds_builder,
          config.eval_split,
          reverse_translation=reverse_translation,
      )
  
      # Tokenize data.
      sp_tokenizer = tokenizer.load_or_train_tokenizer(
          train_data,
          vocab_path=vocab_path,
          vocab_size=config.vocab_size,
          max_corpus_chars=config.max_corpus_chars,
      )
      train_data = train_data.map(
          tokenizer.TokenizeOp(sp_tokenizer),
          num_parallel_calls=AUTOTUNE,
      )
      eval_data = eval_data.map(
          tokenizer.TokenizeOp(sp_tokenizer),
          num_parallel_calls=AUTOTUNE,
      )
  
      batch_size = config.per_device_batch_size * n_devices
  
      train_ds = preprocess_wmt_data(
          train_data,
          shuffle=True,
          num_epochs=None,
          pack_examples=True,
          batch_size=batch_size,
          max_length=config.max_target_length,
      )
  
      eval_ds = preprocess_wmt_data(
          eval_data,
          shuffle=False,
          pack_examples=False,
          batch_size=batch_size,
          max_length=config.max_eval_target_length,
      )
  
      predict_ds = preprocess_wmt_data(
          eval_data,
          shuffle=False,
          pack_examples=False,
          batch_size=batch_size,
          max_length=config.max_predict_length,
          drop_remainder=False,
      )
  
      return train_ds, eval_ds, predict_ds, sp_tokenizer
  
file_path: flax/examples/wmt/input_pipeline.py
imports: import os
  from typing import Dict, Optional, List, Union
  from clu import deterministic_data
  import ml_collections
  import tensorflow
  import tensorflow_datasets
  import tokenizer
label_desc: Use image node_id 5 for calling a global function with example usage: get_wmt_datasets(config) and returns: train_ds,  eval_ds,  predict_ds,  sp_tokenizer