restructure input_pipeline

Author: aireenmei

src/MaxText/layers/engram.py

@@ -29,7 +29,7 @@
 from jax.sharding import Mesh
 from flax import nnx
 
-from MaxText.tokenizer import HFTokenizer
+from maxtext.input_pipeline.tokenizer import HFTokenizer
 from MaxText.common_types import MODEL_MODE_TRAIN, Array, Config
 from MaxText.layers.embeddings import Embed
 from MaxText.layers.initializers import nd_dense_init, NdInitializer

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
 
 
@@ -370,7 +370,6 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
     max_logging.log("Creating policy model with same config as reference model on trainer mesh")
     actor_model, actor_mesh = get_maxtext_model(trainer_config, trainer_devices)
 
-
   if trainer_config.debug.rl:
     max_logging.log("Policy Model initialized successfully")
     nnx.display(actor_model)

src/maxtext/common/checkpointing.py

@@ -23,8 +23,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/sft_train_and_evaluate.py

@@ -85,13 +85,12 @@
 
 from flax import nnx
 
-from MaxText.globals import MAXTEXT_REPO_ROOT
 from MaxText import pyconfig
-from MaxText.input_pipeline import instruction_data_processing
+from MaxText.globals import MAXTEXT_REPO_ROOT
 from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter
+from maxtext.input_pipeline import instruction_data_processing
 from maxtext.trainers.post_train.sft import train_sft
-from maxtext.utils import max_logging
-from maxtext.utils import max_utils
+from maxtext.utils import max_logging, max_utils
 
 # Suppress vLLM logging with a severity level below ERROR
 os.environ["VLLM_LOGGING_LEVEL"] = "ERROR"

src/maxtext/experimental/rl/grpo_input_pipeline.py

@@ -32,8 +32,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:
@@ -141,7 +141,7 @@ def preprocessing_pipeline(
     )
 
     dataset = dataset.map(
-        _input_pipeline_utils.tokenization,
+        input_pipeline_utils.tokenization,
         batched=True,
         fn_kwargs={
             "hf_tokenizer": tokenizer,
@@ -151,7 +151,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,
@@ -166,7 +166,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/inference/inference_microbenchmark.py

@@ -22,13 +22,10 @@
 from absl import app
 from collections.abc import MutableMapping
 
-from MaxText import maxengine
-from MaxText import prefill_packing
-from MaxText import pyconfig
+from MaxText import maxengine, pyconfig
 from maxtext.common import profiler
-from maxtext.utils import gcs_utils
-from maxtext.utils import max_utils
-from maxtext.utils import maxtext_utils
+from maxtext.input_pipeline.packing import prefill_packing
+from maxtext.utils import gcs_utils, max_utils, maxtext_utils
 
 import warnings
 

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

@@ -23,7 +23,7 @@
 
 from dataclasses import dataclass, field
 
-from MaxText.input_pipeline import _input_pipeline_utils
+from maxtext.input_pipeline import input_pipeline_utils
 from maxtext.utils import max_logging
 
 
@@ -83,7 +83,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."