modify data input pipline

Author: RissyRan

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -43,6 +43,8 @@
 
 
 def normalize_features(x, column_name):
+  if column_name not in x and "ids" in x:
+    column_name = "ids"
   return {"inputs": x[column_name], "targets": x[column_name]}
 
 
@@ -85,7 +87,8 @@ def _process_string(string_tensor):
     return [modified_string]
 
   for k in data_keys:
-    features[k] = tf.py_function(_process_string, [features[k]], Tout=[tf.int32])[0]
+    if features[k].dtype == tf.string:
+      features[k] = tf.py_function(_process_string, [features[k]], Tout=[tf.int32])[0]
   return features
 
 

src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py

@@ -249,6 +249,7 @@ def preprocess_train_dataset(
     max_target_length: int,
     shuffle_buffer_size: int,
     data_shuffle_seed: int,
+    is_tokenized_dataset: bool,
 ) -> tf.data.Dataset:
   """Preprocess the training dataset."""
   if sp_tokenizer.pad_id is not None:
@@ -257,10 +258,13 @@ def preprocess_train_dataset(
     pad_id = sp_tokenizer.unk_id
   else:
     pad_id = -1
-  train_ds = train_ds.map(
-      lambda x: TokenizeOp(tokenizer_model=sp_tokenizer, features=x, data_keys=("targets",)),
-      num_parallel_calls=AUTOTUNE,
-  )
+
+  if not is_tokenized_dataset:
+    train_ds = train_ds.map(
+        lambda x: TokenizeOp(tokenizer_model=sp_tokenizer, features=x, data_keys=("targets",)),
+        num_parallel_calls=AUTOTUNE,
+    )
+
   train_ds = reduce_concat_tokens(train_ds, feature_key="targets", batch_size=4096)
   train_ds = split_tokens_to_targets_length(train_ds, max_target_length)
   train_ds = train_ds.shuffle(shuffle_buffer_size, seed=data_shuffle_seed)
@@ -321,15 +325,33 @@ def make_c4_mlperf_train_iterator(
     process_indices,
 ):
   """Make train iterator of customized C4 dataset for mlperf gpt3 training."""
+  train_split = "train2"
+  if config.dataset_name == "c4/en:3.0.5":
+    is_tokenized_dataset = True
+    train_split = "train"
+  elif config.dataset_name == "c4/en:3.0.4":
+    is_tokenized_dataset = False
+  elif config.dataset_name in ["c4/en:3.0.1", "c4/en:3.0.8", "c4/en:3.0.9"]:
+    is_tokenized_dataset = False
+  else:
+    raise ValueError(
+        f"{config.dataset_name=} should be one of "
+        "('c4/en:3.0.1', 'c4/en:3.0.4', 'c4/en:3.0.5', "
+        "'c4/en:3.0.8', 'c4/en:3.0.9')"
+    )
   train_ds = get_dataset(
       dataset_name=config.dataset_name,
-      split="train2",
+      split=train_split,
       dataloading_host_index=process_indices.index(jax.process_index()),
       dataloading_host_count=len(process_indices),
       enable_data_shuffling=config.enable_data_shuffling,
       data_shuffle_seed=config.data_shuffle_seed,
   )
-  train_ds = rekey(train_ds, {"inputs": None, "targets": "text"})
+
+  if is_tokenized_dataset:
+    train_ds = rekey(train_ds, {"inputs": None, "targets": "ids"})
+  else:
+    train_ds = rekey(train_ds, {"inputs": None, "targets": "text"})
 
   sp_tokenizer = get_tokenizer(
       config.tokenizer_path, config.tokenizer_type, config.add_bos, config.add_eos, config.hf_access_token
@@ -341,6 +363,7 @@ def make_c4_mlperf_train_iterator(
       max_target_length=config.max_target_length,
       shuffle_buffer_size=128,
       data_shuffle_seed=config.data_shuffle_seed,
+      is_tokenized_dataset=is_tokenized_dataset,
   )
   train_multihost_gen = multihost_dataloading.MultiHostDataLoadIterator(train_ds, global_mesh)
   return train_multihost_gen
@@ -362,7 +385,11 @@ def make_c4_mlperf_eval_iterator(
     is_tokenized_dataset = False
     eval_slit = "validation"
   else:
-    raise ValueError(f"{config.eval_dataset_name=} should be one of ('c4/en:3.0.1', 'c4/en:3.0.4', 'c4/en:3.0.5')")
+    raise ValueError(
+        f"{config.dataset_name=} should be one of "
+        "('c4/en:3.0.1', 'c4/en:3.0.4', 'c4/en:3.0.5', "
+        "'c4/en:3.0.8', 'c4/en:3.0.9')"
+    )
 
   if is_tokenized_dataset:
     eval_ds = get_dataset(