code: def get_raw_dataset(
      dataset_builder: tfds.core.DatasetBuilder, split: str
  ) -> tf.data.Dataset:
      """Loads a raw text dataset and normalizes feature keys.
  
      Args:
        dataset_builder: TFDS dataset builder that can build `split`.
        split: Split to use. This must be the full split. We shard the split across
          multiple hosts and currently don't support sharding subsplits.
  
      Returns:
        Dataset with source and target language features mapped to 'inputs' and
        'targets'.
      """
      num_examples = dataset_builder.info.splits[split].num_examples
      per_host_split = deterministic_data.get_read_instruction_for_host(
          split, num_examples, drop_remainder=False
      )
      ds = dataset_builder.as_dataset(
          split=per_host_split, shuffle_files=False
      )
      ds = ds.map(
          NormalizeFeatureNamesOp(dataset_builder.info),
          num_parallel_calls=AUTOTUNE,
      )
      return ds
  
file_path: flax/examples/lm1b/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 1 for calling a global function with example usage: get_raw_dataset(dataset_builder, split) and returns: ds