Merge pull request #3811 from AI-Hypercomputer:input_fix

PiperOrigin-RevId: 910356301

Author: Google-ML-Automation

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -43,8 +43,6 @@
 
 
 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]}
 
 
@@ -87,8 +85,7 @@ def _process_string(string_tensor):
     return [modified_string]
 
   for k in data_keys:
-    if features[k].dtype == tf.string:
-      features[k] = tf.py_function(_process_string, [features[k]], Tout=[tf.int32])[0]
+    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,7 +249,6 @@ 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:
@@ -259,11 +258,10 @@ def preprocess_train_dataset(
   else:
     pad_id = -1
 
-  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 = 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)
@@ -325,33 +323,16 @@ def make_c4_mlperf_train_iterator(
     process_indices,
 ):
   """Make train iterator of customized C4 dataset for mlperf gpt3 training."""
-  train_split = "train"
-  if config.dataset_name == "c4/en:3.0.1":
-    # gs://max-datasets-rogue/c4/en/3.0.1
-    is_tokenized_dataset = False
-  elif config.dataset_name == "c4/en:3.0.5":
-    # gs://mlperf-6-submission/tfds/c4/en/3.0.5
-    is_tokenized_dataset = True
-  elif config.dataset_name == "c4/en:3.0.7":
-    # gs://max-datasets-rogue/c4/en/3.0.7
-    is_tokenized_dataset = False
-    train_split = "train2"
-  else:
-    raise ValueError(f"{config.dataset_name=} should be one of " "('c4/en:3.0.1', 'c4/en:3.0.5', 'c4/en:3.0.7')")
   train_ds = get_dataset(
       dataset_name=config.dataset_name,
-      split=train_split,
+      split="train2",
       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,
   )
 
-  if is_tokenized_dataset:
-    train_ds = rekey(train_ds, {"inputs": None, "targets": "ids"})
-  else:
-    train_ds = rekey(train_ds, {"inputs": None, "targets": "text"})
-
+  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
   )
@@ -362,7 +343,6 @@ 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