Merge pull request #3050 from AI-Hypercomputer:aireen/input_restructure

PiperOrigin-RevId: 864538350

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/rl/train_rl.py

@@ -77,7 +77,7 @@
 from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter
 from MaxText.rl.evaluate_rl import evaluate
 from MaxText.rl import utils_rl
-from MaxText.input_pipeline.instruction_data_processing import load_template_from_file
+from maxtext.input_pipeline.instruction_data_processing import load_template_from_file
 from maxtext.utils import max_logging, max_utils, maxtext_utils, model_creation_utils
 
 

src/maxtext/common/checkpointing.py

@@ -22,8 +22,8 @@
 from flax.training import train_state
 import jax
 from MaxText.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
-from MaxText.multihost_dataloading import MultiHostDataLoadIterator, RemoteIterator
-from MaxText.input_pipeline.input_pipeline_interface import PlaceHolderDataIterator
+from maxtext.input_pipeline.multihost_dataloading import MultiHostDataLoadIterator, RemoteIterator
+from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
 from maxtext.utils import exceptions
 from maxtext.utils import max_logging
 import numpy as np

src/maxtext/examples/demo_decoding.ipynb

@@ -136,8 +136,8 @@
         "import MaxText as mt\n",
         "from MaxText import common_types\n",
         "from MaxText import pyconfig\n",
-        "from MaxText.input_pipeline import _input_pipeline_utils\n",
-        "from MaxText.utils.ckpt_conversion import to_maxtext\n",
+        "from maxtext.input_pipeline import _input_pipeline_utils\n",
+        "from maxtext.src.maxtext.utils.ckpt_conversion import to_maxtext\n",
         "from maxtext.inference import inference_utils\n",
         "from maxtext.utils import maxtext_utils\n",
         "from maxtext.utils import max_logging\n",
@@ -437,4 +437,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 5
-}
+}

src/maxtext/examples/sft_train_and_evaluate.py

@@ -88,7 +88,7 @@
 
 from MaxText.globals import MAXTEXT_REPO_ROOT
 from MaxText import pyconfig
-from MaxText.input_pipeline import instruction_data_processing
+from maxtext.input_pipeline import instruction_data_processing
 from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter
 from maxtext.trainers.post_train.sft import train_sft
 from maxtext.utils import max_logging

src/maxtext/inference/inference_microbenchmark.py

@@ -23,9 +23,9 @@
 from collections.abc import MutableMapping
 
 from MaxText import maxengine
-from MaxText import prefill_packing
 from MaxText import pyconfig
 from maxtext.common import profiler
+from maxtext.input_pipeline.packing import prefill_packing
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils

src/maxtext/inference/mlperf/offline_inference.py

@@ -35,8 +35,8 @@
 # pylint: disable=no-name-in-module
 from MaxText.maxengine import MaxEngine
 from MaxText.maxengine import set_engine_vars_from_base_engine
-from MaxText.prefill_packing import PrefillProcessor
-from MaxText.prefill_packing import BatchedPrefillProcessor
+from maxtext.input_pipeline.packing.prefill_packing import PrefillProcessor
+from maxtext.input_pipeline.packing.prefill_packing import BatchedPrefillProcessor
 
 DecodeState = Any
 Params = Any

src/maxtext/inference/offline_engine.py

@@ -54,7 +54,7 @@
 from jax.experimental import mesh_utils
 
 from MaxText.maxengine import MaxEngine
-from MaxText.prefill_packing import PrefillProcessor, BatchedPrefillProcessor
+from maxtext.input_pipeline.packing.prefill_packing import PrefillProcessor, BatchedPrefillProcessor
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 

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–2025 Google LLC
+# Copyright 2023-2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

…ipeline/_distillation_data_processing.py …pipeline/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."