Merge pull request #3117 from AI-Hypercomputer:aireen/local_import

PiperOrigin-RevId: 867841433

Author: Google-ML-Automation

src/MaxText/experimental/rl/grpo_input_pipeline.py

@@ -28,8 +28,6 @@
 import jax
 from jax.sharding import Mesh
 
-import datasets
-
 import transformers
 
 import grain.python as grain
@@ -217,6 +215,8 @@ def make_hf_train_iterator(
   Returns:
     A local data iterator for the training set.
   """
+  import datasets  # pylint: disable=import-outside-toplevel
+
   train_ds = datasets.load_dataset(
       config.hf_path,
       data_dir=config.hf_data_dir,

src/MaxText/input_pipeline/_distillation_data_processing.py

@@ -23,8 +23,6 @@
 
 from dataclasses import dataclass, field
 
-import datasets
-
 from MaxText.input_pipeline import _input_pipeline_utils
 from maxtext.utils import max_logging
 
@@ -101,6 +99,8 @@ def process_dataset(config, dataset):  # pylint: disable=redefined-outer-name
 
 def load_dataset(config):  # pylint: disable=redefined-outer-name
   """Loads dataset from Hugging Face."""
+  import datasets  # pylint: disable=import-outside-toplevel
+
   assert config.dataset_type == "huggingface", "Only dataset from Hugging Face is supported."
 
   return datasets.load_dataset(

src/MaxText/input_pipeline/_hf_data_processing.py

@@ -18,8 +18,6 @@
 
 import jax
 
-import datasets
-
 import transformers
 
 import grain.python as grain
@@ -62,6 +60,8 @@ def vision_sft_preprocessing_pipeline(
     batch_size = global_batch_size // jax.process_count()
 
   # for multi-epoch with shuffle, shuffle each epoch with different seeds then concat
+  import datasets  # pylint: disable=import-outside-toplevel
+
   if config.enable_data_shuffling and config.num_epoch > 1:
     epoch_datasets = [dataset.shuffle(seed=config.data_shuffle_seed + i) for i in range(config.num_epoch)]
     dataset = datasets.concatenate_datasets(epoch_datasets)
@@ -215,6 +215,7 @@ def preprocessing_pipeline(
     num_epoch=1,
 ):
   """pipeline for preprocessing HF dataset"""
+  import datasets  # pylint: disable=import-outside-toplevel
 
   assert global_batch_size % global_mesh.size == 0, "Batch size should be divisible by number of global devices."
   # Tunix GA requires per-micro-batch slicing at the data level,
@@ -377,6 +378,8 @@ def make_hf_train_iterator(
     process_indices_train,
 ):
   """Load, preprocess dataset and return iterators"""
+  import datasets  # pylint: disable=import-outside-toplevel
+
   train_ds = datasets.load_dataset(
       config.hf_path,
       name=config.hf_name,
@@ -433,6 +436,8 @@ def make_hf_eval_iterator(
     process_indices_eval,
 ):
   """Make Hugging Face evaluation iterator. Load and preprocess eval dataset: and return iterator."""
+  import datasets  # pylint: disable=import-outside-toplevel
+
   eval_ds = datasets.load_dataset(
       config.hf_path,
       name=config.hf_name,

src/MaxText/input_pipeline/_input_pipeline_utils.py

@@ -17,9 +17,11 @@
 import dataclasses
 import warnings
 from threading import current_thread
-from typing import Any
-import datasets
-from datasets.distributed import split_dataset_by_node
+from typing import Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+  import datasets
+
 import grain.python as grain
 import numpy as np
 import tensorflow as tf
@@ -145,6 +147,8 @@ def is_conversational(features, data_columns):
   data_columns = ["prompt", "completion"]
   is_conversational(features, data_columns) returns False.
   """
+  import datasets  # pylint: disable=import-outside-toplevel
+
   for column in data_columns:
     messages = features[column]
     if isinstance(messages, datasets.Sequence):
@@ -293,13 +297,16 @@ class HFDataSource(grain.RandomAccessDataSource):
 
   def __init__(
       self,
-      dataset: datasets.IterableDataset,
+      dataset: "datasets.IterableDataset",
       dataloading_host_index: int,
       dataloading_host_count: int,
       num_threads: int,
       max_target_length: int,
       data_column_names: list[str],
   ):
+    from datasets.distributed import split_dataset_by_node  # pylint: disable=import-outside-toplevel
+
+    self._split_dataset_by_node = split_dataset_by_node
     self.dataset = dataset
     self.num_threads = num_threads
     self.dataloading_host_count = dataloading_host_count
@@ -312,7 +319,7 @@ def __init__(
       self.n_shards = 1
     self._check_shard_count()
     self.dataset_shards = [dataloading_host_index * self.num_threads + i for i in range(self.num_threads)]
-    self.datasets = [split_dataset_by_node(dataset, world_size=self.n_shards, rank=x) for x in self.dataset_shards]
+    self.datasets = [self._split_dataset_by_node(dataset, world_size=self.n_shards, rank=x) for x in self.dataset_shards]
     self.data_iters = []
 
   def _check_shard_count(self):
@@ -333,7 +340,9 @@ def _update_shard(self, idx):
       )
       max_logging.log(f"New shard is {new_shard}")
       self.dataset_shards[idx] = new_shard
-      self.datasets[idx] = split_dataset_by_node(self.dataset, world_size=self.n_shards, rank=self.dataset_shards[idx])
+      self.datasets[idx] = self._split_dataset_by_node(
+          self.dataset, world_size=self.n_shards, rank=self.dataset_shards[idx]
+      )
       self.data_iters[idx] = iter(self.datasets[idx])
     else:
       raise StopIteration(f"Run out of shards on host {self.dataloading_host_index}, shard {new_shard} is not available")

src/MaxText/input_pipeline/instruction_data_processing.py

@@ -14,7 +14,6 @@
 
 """Preprocessing for instruction dataset."""
 
-import datasets
 import json
 import os
 import re
@@ -117,6 +116,8 @@ def convert_to_conversational_format(
     chat_template_path,
 ):
   """Converts instruction dataset to conversational format."""
+  import datasets  # pylint: disable=import-outside-toplevel
+
   template_config = None
   if chat_template_path:
     template_config = load_template_from_file(chat_template_path)

src/MaxText/pyconfig.py

@@ -20,6 +20,9 @@
 from typing import Any
 import copy
 
+# Disable dill to avoid conflict with gfile (dill requires buffering=0, which gfile forbids)
+os.environ["HF_DATASETS_DISABLE_DILL"] = "1"
+
 import jax
 import jax.numpy as jnp
 

src/maxtext/examples/sft_train_and_evaluate.py

@@ -78,7 +78,6 @@
 from tqdm.auto import tqdm
 from typing import Sequence
 
-import datasets
 import grain
 import os
 import re
@@ -140,6 +139,8 @@ def get_test_dataset(config, tokenizer):
     A grain.MapDataset instance for the test split, with prompts and target
     answers.
   """
+  import datasets  # pylint: disable=import-outside-toplevel
+
   template_config = instruction_data_processing.load_template_from_file(config.chat_template_path)
   dataset = datasets.load_dataset(
       DATASET_NAME,