code: def _get_datasets(self):
      config = default.get_config()
      config.per_device_batch_size = 1
      config.eval_per_device_batch_size = 2
      config.vocab_size = 32
      config.max_corpus_chars = 1000
      config.max_target_length = _TARGET_LENGTH
      config.max_eval_target_length = _EVAL_TARGET_LENGTH
      config.max_predict_length = _PREDICT_TARGET_LENGTH
  
      vocab_path = os.path.join(
          tempfile.mkdtemp(), "sentencepiece_model"
      )
  
      # Go two directories up to the root of the flax directory.
      flax_root_dir = pathlib.Path(__file__).parents[2]
      data_dir = (
          str(flax_root_dir) + "/.tfds/metadata"
      )  # pylint: disable=unused-variable
  
      with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):
          (
              train_ds,
              eval_ds,
              predict_ds,
              _,
          ) = input_pipeline.get_datasets(
              n_devices=2, config=config, vocab_path=vocab_path
          )
      return train_ds, eval_ds, predict_ds
  
file_path: flax/examples/lm1b/input_pipeline_test.py
imports: import os
  import pathlib
  import tempfile
  from absl.testing import absltest
  import tensorflow_datasets
  from configs import default
  import input_pipeline
label_desc: Use image node_id 2 for calling the InputPipelineTest obj's underlying member method code with example usage: obj._get_datasets() and returns: train_ds,  eval_ds,  predict_ds