Reverts 885cd56

PiperOrigin-RevId: 864609609

Author: Google-ML-Automation

src/MaxText/experimental/rl/grpo_input_pipeline.py

@@ -34,8 +34,8 @@
 
 import grain.python as grain
 
-from maxtext.input_pipeline import input_pipeline_interface
-from maxtext.input_pipeline import input_pipeline_utils
+from MaxText.input_pipeline import input_pipeline_interface
+from MaxText.input_pipeline import _input_pipeline_utils
 
 
 class SingleHostDataLoader:
@@ -143,7 +143,7 @@ def preprocessing_pipeline(
     )
 
     dataset = dataset.map(
-        input_pipeline_utils.tokenization,
+        _input_pipeline_utils.tokenization,
         batched=True,
         fn_kwargs={
             "hf_tokenizer": tokenizer,
@@ -153,7 +153,7 @@ def preprocessing_pipeline(
         },
     )
   dataset = dataset.select_columns(data_column_names)
-  dataset = input_pipeline_utils.HFDataSource(
+  dataset = _input_pipeline_utils.HFDataSource(
       dataset,
       dataloading_host_index,
       dataloading_host_count,
@@ -168,7 +168,7 @@ def lists2array(x):
 
   operations = [
       grain.MapOperation(lists2array),
-      input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, add_true_length=True),
+      _input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, add_true_length=True),
       grain.Batch(batch_size=global_batch_size // jax.process_count(), drop_remainder=drop_remainder),
   ]
 

src/maxtext/input_pipeline/__init__.py src/MaxText/input_pipeline/__init__.pysrc/maxtext/input_pipeline/__init__.py renamed to src/MaxText/input_pipeline/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023-2026 Google LLC
+# Copyright 2023–2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

…pipeline/distillation_data_processing.py …ipeline/_distillation_data_processing.pysrc/maxtext/input_pipeline/distillation_data_processing.py renamed to src/MaxText/input_pipeline/_distillation_data_processing.py src/maxtext/input_pipeline/distillation_data_processing.py renamed to src/MaxText/input_pipeline/_distillation_data_processing.py

@@ -25,7 +25,7 @@
 
 import datasets
 
-from maxtext.input_pipeline import input_pipeline_utils
+from MaxText.input_pipeline import _input_pipeline_utils
 from maxtext.utils import max_logging
 
 
@@ -85,7 +85,7 @@ def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
   assert any(
       set(data_column_names) == set(supported) for supported in supported_columns
   ), f"Dataset column names mismatch. Expected columns to match one of {supported_columns}, but got {data_column_names}"
-  assert input_pipeline_utils.is_conversational(
+  assert _input_pipeline_utils.is_conversational(
       dataset.features, data_column_names
   ), "Dataset is not in conversational format."
 

…/input_pipeline/grain_data_processing.py …input_pipeline/_grain_data_processing.pysrc/maxtext/input_pipeline/grain_data_processing.py renamed to src/MaxText/input_pipeline/_grain_data_processing.py src/maxtext/input_pipeline/grain_data_processing.py renamed to src/MaxText/input_pipeline/_grain_data_processing.py

@@ -26,10 +26,10 @@
 from grain.experimental import BestFitPackIterDataset, pick_performance_config
 import grain.python as grain
 
-from maxtext.input_pipeline import input_pipeline_utils
-from maxtext.input_pipeline import grain_tokenizer
-from maxtext.input_pipeline import multihost_dataloading
-from maxtext.input_pipeline import tokenizer
+from MaxText.input_pipeline import _input_pipeline_utils
+from MaxText.input_pipeline import _grain_tokenizer
+from MaxText import multihost_dataloading
+from MaxText import tokenizer
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_logging
 
@@ -199,10 +199,10 @@ def pretrain_preprocessing_pipeline(
 ):
   """Use grain pipeline to pre-process the dataset and return iterators for pretrain"""
   if config.grain_file_type == "arrayrecord":
-    dataset = dataset.map(input_pipeline_utils.ParseFeatures(data_columns, tokenize))
-    dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
+    dataset = dataset.map(_input_pipeline_utils.ParseFeatures(data_columns, tokenize))
+    dataset = dataset.map(_input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
   else:
-    dataset = dataset.map(input_pipeline_utils.KeepFeatures(feature_names=data_columns))
+    dataset = dataset.map(_input_pipeline_utils.KeepFeatures(feature_names=data_columns))
 
   assert len(data_columns) == 1
   text_column = data_columns[0]
@@ -224,13 +224,13 @@ def pretrain_preprocessing_pipeline(
 
   if tokenize:
     if config.use_truncation:
-      dataset = dataset.map(grain_tokenizer.TokenizeAndTrim(text_column, config.max_target_length, tokenizer_model))
+      dataset = dataset.map(_grain_tokenizer.TokenizeAndTrim(text_column, config.max_target_length, tokenizer_model))
     else:
-      dataset = dataset.apply(grain_tokenizer.TokenizeAndChunk(text_column, config.max_target_length, tokenizer_model))
+      dataset = dataset.apply(_grain_tokenizer.TokenizeAndChunk(text_column, config.max_target_length, tokenizer_model))
 
   data_columns = ("inputs", "targets")
   rekey_dict = {col: text_column for col in data_columns}
-  dataset = dataset.map(input_pipeline_utils.Rekey(rekey_dict))
+  dataset = dataset.map(_input_pipeline_utils.Rekey(rekey_dict))
 
   # Pack and Batch examples.
   batch_size = config.global_batch_size_to_load // jax.process_count()
@@ -273,15 +273,15 @@ def pretrain_preprocessing_pipeline(
         "targets_position": "targets_positions",
         "inputs_position": "inputs_positions",
     }
-    dataset = dataset.map(input_pipeline_utils.Rekey(rekey_dict))
+    dataset = dataset.map(_input_pipeline_utils.Rekey(rekey_dict))
   else:
-    dataset = dataset.map(input_pipeline_utils.PadOrTrimToMaxLength(config.max_target_length, pad_id))
+    dataset = dataset.map(_input_pipeline_utils.PadOrTrimToMaxLength(config.max_target_length, pad_id))
   batch_fn = functools.partial(grain.experimental.batch_and_pad, batch_size=batch_size, pad_value=pad_id)
   dataset = dataset.batch(batch_size, batch_fn=batch_fn)
 
   # Shift inputs for teacher-forced training
   dataset = dataset.map(
-      input_pipeline_utils.ShiftData(
+      _input_pipeline_utils.ShiftData(
           ignored_ids=[pad_id],
           axis=1,
       )
@@ -313,8 +313,8 @@ def dpo_preprocessing_pipeline(
 ):
   """Use grain to pre-process the dataset and return iterators for dpo fine-tuning"""
   if config.grain_file_type == "arrayrecord":
-    dataset = dataset.map(input_pipeline_utils.ParseFeatures(data_columns, tokenize))
-    dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
+    dataset = dataset.map(_input_pipeline_utils.ParseFeatures(data_columns, tokenize))
+    dataset = dataset.map(_input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
   tokenizer_model = tokenizer.build_tokenizer(
       config.tokenizer_path,
       config.tokenizer_type,
@@ -331,9 +331,9 @@ def dpo_preprocessing_pipeline(
     pad_id = -1
 
   if tokenize:
-    dataset = dataset.map(grain_tokenizer.TokenizeAndTrim(data_columns, config.max_target_length, tokenizer_model))
+    dataset = dataset.map(_grain_tokenizer.TokenizeAndTrim(data_columns, config.max_target_length, tokenizer_model))
 
-  dataset = dataset.map(input_pipeline_utils.PadOrTrimToMaxLength(config.max_target_length, pad_id))
+  dataset = dataset.map(_input_pipeline_utils.PadOrTrimToMaxLength(config.max_target_length, pad_id))
   batch_size = config.global_batch_size_to_load // jax.process_count()
   batch_fn = functools.partial(grain.experimental.batch_and_pad, batch_size=batch_size, pad_value=pad_id)
   dataset = dataset.batch(batch_size, batch_fn=batch_fn)

…axtext/input_pipeline/grain_tokenizer.py …xText/input_pipeline/_grain_tokenizer.pysrc/maxtext/input_pipeline/grain_tokenizer.py renamed to src/MaxText/input_pipeline/_grain_tokenizer.py src/maxtext/input_pipeline/grain_tokenizer.py renamed to src/MaxText/input_pipeline/_grain_tokenizer.py

@@ -20,7 +20,7 @@
 from typing import Any
 import grain.python as grain
 import numpy as np
-from maxtext.input_pipeline import tokenizer
+from MaxText import tokenizer
 
 
 @dataclasses.dataclass

…ext/input_pipeline/hf_data_processing.py …xt/input_pipeline/_hf_data_processing.pysrc/maxtext/input_pipeline/hf_data_processing.py renamed to src/MaxText/input_pipeline/_hf_data_processing.py src/maxtext/input_pipeline/hf_data_processing.py renamed to src/MaxText/input_pipeline/_hf_data_processing.py

@@ -26,9 +26,9 @@
 
 import numpy as np
 
-from maxtext.input_pipeline import input_pipeline_utils
-from maxtext.input_pipeline import instruction_data_processing
-from maxtext.input_pipeline import multihost_dataloading
+from MaxText.input_pipeline import _input_pipeline_utils
+from MaxText.input_pipeline import instruction_data_processing
+from MaxText import multihost_dataloading
 
 
 def _get_pad_id(tokenizer):
@@ -61,7 +61,7 @@ def vision_sft_preprocessing_pipeline(
   # If multiple image columns are provided, merge them into a single 'images' column.
   if isinstance(image_column, list):
     dataset = dataset.map(
-        input_pipeline_utils.merge_image_columns,
+        _input_pipeline_utils.merge_image_columns,
         fn_kwargs={
             "image_columns": image_column,
             "max_num_images_per_example": config.max_num_images_per_example,
@@ -75,20 +75,20 @@ def vision_sft_preprocessing_pipeline(
     dataset = dataset.rename_column(image_column, "images")
 
   dataset = dataset.map(
-      input_pipeline_utils.reformat_prompt,
+      _input_pipeline_utils.reformat_prompt,
       fn_kwargs={
           "column": text_columns[0],
           "image_placeholder": config.image_placeholder,
           "model_name": config.model_name,
       },
   )
   dataset = dataset.map(
-      input_pipeline_utils.reformat_response,
+      _input_pipeline_utils.reformat_response,
       fn_kwargs={"column": text_columns[1], "model_name": config.model_name},
   )
 
   dataset = dataset.map(
-      input_pipeline_utils.pre_process_image_sft,
+      _input_pipeline_utils.pre_process_image_sft,
       fn_kwargs={"image_column": "images", "model_name": config.model_name},
   )
 
@@ -102,7 +102,7 @@ def vision_sft_preprocessing_pipeline(
   pad_id = _get_pad_id(tokenizer)
 
   dataset = dataset.map(
-      input_pipeline_utils.tokenization,
+      _input_pipeline_utils.tokenization,
       batched=True,
       batch_size=global_batch_size,
       fn_kwargs={
@@ -113,11 +113,11 @@ def vision_sft_preprocessing_pipeline(
       },
   )
   dataset = dataset.map(
-      input_pipeline_utils.prepare_text_for_image_fusion,
+      _input_pipeline_utils.prepare_text_for_image_fusion,
       fn_kwargs={"column_name": text_columns[0], "model_name": config.model_name},
   )
 
-  dataset = input_pipeline_utils.HFDataSource(
+  dataset = _input_pipeline_utils.HFDataSource(
       dataset=dataset,
       dataloading_host_index=dataloading_host_index,
       dataloading_host_count=dataloading_host_count,
@@ -127,7 +127,7 @@ def vision_sft_preprocessing_pipeline(
   )
   operations = []
   operations.append(
-      input_pipeline_utils.SFTPromptMaskingVision(
+      _input_pipeline_utils.SFTPromptMaskingVision(
           query_column=text_columns[0],
           response_column=text_columns[1],
           max_target_length=config.max_target_length,
@@ -136,17 +136,17 @@ def vision_sft_preprocessing_pipeline(
   )
   # TODO(aireenmei, hengtaoguo): support packing
   operations.append(
-      input_pipeline_utils.PadOrTrimToMaxLength(
+      _input_pipeline_utils.PadOrTrimToMaxLength(
           config.max_target_length,
           pad_id,
           model_name=config.model_name,
           max_num_images_per_example=config.max_num_images_per_example,
       )
   )
-  operations.append(input_pipeline_utils.ExtractImagesAndMasks())
+  operations.append(_input_pipeline_utils.ExtractImagesAndMasks())
   operations.append(grain.Batch(batch_size=batch_size, drop_remainder=True))
-  operations.append(input_pipeline_utils.FoldImagesIntoBatch(model_name=config.model_name))
-  operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id], axis=1))
+  operations.append(_input_pipeline_utils.FoldImagesIntoBatch(model_name=config.model_name))
+  operations.append(_input_pipeline_utils.ShiftData(ignored_ids=[pad_id], axis=1))
   dummy_index_sampler = grain.IndexSampler(
       num_records=len(dataset),
       num_epochs=1,
@@ -227,7 +227,7 @@ def preprocessing_pipeline(
         dataset=dataset, data_columns=data_column_names, chat_template_path=chat_template_path
     )
 
-    assert input_pipeline_utils.is_conversational(
+    assert _input_pipeline_utils.is_conversational(
         dataset.features, data_column_names
     ), "Dataset is not in conversational format."
 
@@ -237,15 +237,15 @@ def preprocessing_pipeline(
           {combined_column_name: [{"content": datasets.Value(dtype="string"), "role": datasets.Value(dtype="string")}]}
       )
       dataset = dataset.map(
-          input_pipeline_utils.combine_columns,
+          _input_pipeline_utils.combine_columns,
           fn_kwargs={"columns": data_column_names, "data_column": combined_column_name},
           remove_columns=data_column_names,
           features=dataset_features,
       )
 
     data_column_names = list(dataset.features.keys())
     dataset = dataset.map(
-        input_pipeline_utils.apply_chat_template,
+        _input_pipeline_utils.apply_chat_template,
         fn_kwargs={"tokenizer_model": tokenizer, "data_column_name": data_column_names[0]},
     )
   else:
@@ -255,7 +255,7 @@ def preprocessing_pipeline(
 
   if tokenize:
     dataset = dataset.map(
-        input_pipeline_utils.tokenization,
+        _input_pipeline_utils.tokenization,
         batched=True,
         fn_kwargs={
             "hf_tokenizer": tokenizer,
@@ -265,7 +265,7 @@ def preprocessing_pipeline(
         },
     )
 
-  dataset = input_pipeline_utils.HFDataSource(
+  dataset = _input_pipeline_utils.HFDataSource(
       dataset,
       dataloading_host_index,
       dataloading_host_count,
@@ -276,7 +276,7 @@ def preprocessing_pipeline(
   operations = []
   if use_sft:
     operations.append(
-        input_pipeline_utils.SFTPromptMasking(
+        _input_pipeline_utils.SFTPromptMasking(
             text_column_name=data_column_names[0],
             completion_only=sft_train_on_completion_only,
             max_target_length=max_target_length,
@@ -293,7 +293,7 @@ def lists2array(x):
     operations.append(grain.MapOperation(lists2array))
   else:
     assert len(data_column_names) == 1
-    operations.append(input_pipeline_utils.HFNormalizeFeatures(data_column_names[0]))
+    operations.append(_input_pipeline_utils.HFNormalizeFeatures(data_column_names[0]))
     data_column_names = ("inputs", "targets")
 
   if packing and not use_dpo:
@@ -308,13 +308,13 @@ def lists2array(x):
             max_sequences_per_bin=max_segments,
         )
     )
-    operations.append(input_pipeline_utils.ReformatPacking(data_column_names))
+    operations.append(_input_pipeline_utils.ReformatPacking(data_column_names))
   else:
-    operations.append(input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id))
+    operations.append(_input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id))
     operations.append(grain.Batch(batch_size=global_batch_size // jax.process_count(), drop_remainder=drop_remainder))
 
   if shift and not use_dpo:
-    operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1))
+    operations.append(_input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1))
 
   # Since HuggingFace IterableDataset does not support access through index
   # Indexes generated by dummy_index_sampler is not used.

…t/input_pipeline/input_pipeline_utils.py …/input_pipeline/_input_pipeline_utils.pysrc/maxtext/input_pipeline/input_pipeline_utils.py renamed to src/MaxText/input_pipeline/_input_pipeline_utils.py src/maxtext/input_pipeline/input_pipeline_utils.py renamed to src/MaxText/input_pipeline/_input_pipeline_utils.py

@@ -23,8 +23,8 @@
 import grain.python as grain
 import numpy as np
 import tensorflow as tf
+from MaxText import tokenizer
 from MaxText import multimodal_utils
-from maxtext.input_pipeline import tokenizer
 from maxtext.utils import max_logging
 
 Features = dict[str, tf.Tensor]