Merge pull request #3196 from AI-Hypercomputer:aireen/tokenizer_no_tf

PiperOrigin-RevId: 874631535

Author: Google-ML-Automation

src/maxtext/input_pipeline/grain_data_processing.py

@@ -213,7 +213,6 @@ def pretrain_preprocessing_pipeline(
       config.add_bos,
       config.add_eos,
       config.hf_access_token,
-      config.dataset_type,
   )
   if tokenizer_model.pad_id is not None:
     pad_id = tokenizer_model.pad_id
@@ -321,7 +320,6 @@ def dpo_preprocessing_pipeline(
       config.add_bos,
       config.add_eos,
       config.hf_access_token,
-      config.dataset_type,
   )
   if tokenizer_model.pad_id is not None:
     pad_id = tokenizer_model.pad_id

src/maxtext/input_pipeline/grain_tokenizer.py

@@ -30,7 +30,7 @@ class TokenizerTransformBase:
   # pylint: disable=attribute-defined-outside-init
   feature_names: str | Sequence[str]
   sequence_length: int | Sequence[int]
-  tokenizer: tokenizer.SentencePieceTokenizerGrain | tokenizer.HFTokenizer
+  tokenizer: tokenizer.SentencePieceTokenizer | tokenizer.HFTokenizer
 
   def __post_init__(self):
     self._processor = None

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -17,7 +17,7 @@
 import dataclasses
 import warnings
 from threading import current_thread
-from typing import Any, TYPE_CHECKING
+from typing import Any, Iterable, TYPE_CHECKING
 
 if TYPE_CHECKING:
   import datasets
@@ -40,11 +40,9 @@ def normalize_features(x, column_name):
   return {"inputs": x[column_name], "targets": x[column_name]}
 
 
-def get_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token=None, dataset_type="tfds"):
+def get_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token=None):
   # Load tokenizer
-  tokenizer_model = tokenizer.build_tokenizer(
-      tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token, dataset_type
-  )
+  tokenizer_model = tokenizer.build_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token)
   return tokenizer_model
 
 
@@ -67,6 +65,21 @@ def add_segmentation_and_position(x, data_columns, padding_token=0):
   return x
 
 
+def TokenizeOp(tokenizer_model, features: Features, data_keys: Iterable[str] = ("inputs", "targets")) -> Features:
+  """Op for tokenization"""
+
+  def _process_string(string_tensor):
+    # Extract string value and decode it if necessary
+    string_value = string_tensor.numpy().decode("utf-8")
+    # encode and extract the tokenized integers
+    modified_string = tokenizer_model.encode(string_value)
+    return [modified_string]
+
+  for k in data_keys:
+    features[k] = tf.py_function(_process_string, [features[k]], Tout=[tf.int32])[0]
+  return features
+
+
 ########## Functions used by HF pipeline
 
 

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -25,7 +25,6 @@
 import jax
 
 from maxtext.input_pipeline import multihost_dataloading
-from maxtext.input_pipeline import tokenizer
 from maxtext.input_pipeline.packing import sequence_packing
 from maxtext.input_pipeline import input_pipeline_utils
 
@@ -116,7 +115,9 @@ def preprocessing_pipeline(
 
   if tokenize:
     dataset = dataset.map(
-        lambda x: tokenizer.TokenizeOp(tokenizer=tokenizer_model, features=x, data_keys=data_column_names),
+        lambda x: input_pipeline_utils.TokenizeOp(
+            tokenizer_model=tokenizer_model, features=x, data_keys=data_column_names
+        ),
         num_parallel_calls=AUTOTUNE,
     )
 

src/maxtext/input_pipeline/tfds_data_processing_c4_mlperf.py

@@ -27,10 +27,10 @@
 import jax.numpy as jnp
 from jax.experimental import multihost_utils
 
-from maxtext.input_pipeline import tokenizer
 from maxtext.input_pipeline import multihost_dataloading
 from maxtext.input_pipeline.packing import sequence_packing
 from maxtext.input_pipeline.input_pipeline_utils import get_tokenizer
+from maxtext.input_pipeline.input_pipeline_utils import TokenizeOp
 from maxtext.utils import max_logging
 
 AUTOTUNE = tf.data.experimental.AUTOTUNE
@@ -258,7 +258,7 @@ def preprocess_train_dataset(
   else:
     pad_id = -1
   train_ds = train_ds.map(
-      lambda x: tokenizer.TokenizeOp(tokenizer=sp_tokenizer, features=x, data_keys=("targets",)),
+      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)
@@ -283,7 +283,7 @@ def preprocess_eval_dataset(
   # group text up to max_target_length if the dataset is not pre-tokenized/pre-processed
   if not is_tokenized_dataset:
     eval_ds = eval_ds.map(
-        lambda x: tokenizer.TokenizeOp(tokenizer=sp_tokenizer, features=x, data_keys=("targets",)),
+        lambda x: TokenizeOp(tokenizer_model=sp_tokenizer, features=x, data_keys=("targets",)),
         num_parallel_calls=AUTOTUNE,
     )
     # hardcode batch_sizes 24567 i.e. the exp size in split validation_24567exp

src/maxtext/input_pipeline/tokenizer.py

@@ -14,20 +14,15 @@
 
 """Provides op for tokenizing a dataset."""
 
-from typing import Iterable, Literal, Sequence, Collection
+from typing import Literal, Sequence, Collection
 from pathlib import Path
-import tensorflow as tf
-import tensorflow_text as tftxt
 from maxtext.utils import max_logging
 import transformers
 import tiktoken
 from tiktoken.load import load_tiktoken_bpe
 from sentencepiece import SentencePieceProcessor
 
 
-Features = dict[str, tf.Tensor]
-
-
 class TikTokenTokenizer:
   """
   Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
@@ -184,33 +179,23 @@ def _split_whitespaces_or_nonwhitespaces(s: str, max_consecutive_slice_len: int)
 
 class SentencePieceTokenizer:
   """
-  Tokenizing and encoding/decoding text using the Sentencepiece tokenizer loaded with tensorflow_text
-  """
-
-  def __init__(self, model_path: str, add_bos: bool, add_eos: bool):
-    max_logging.log(f"Tokenizer path: {model_path}")
-    with tf.io.gfile.GFile(model_path, "rb") as model_fp:
-      sp_model = model_fp.read()
-    self.sp_tokenizer = tftxt.SentencepieceTokenizer(model=sp_model, add_bos=add_bos, add_eos=add_eos, reverse=False)
-    self.pad_id = self.sp_tokenizer.string_to_id("<pad>")
-    self.unk_id = self.sp_tokenizer.string_to_id("<unk>")
-
-  def encode(self, s: str) -> list[int]:
-    return self.sp_tokenizer.tokenize(s)
-
-  def decode(self, t: Sequence[int]) -> str:
-    return self.sp_tokenizer.detokenize(t)
-
-
-class SentencePieceTokenizerGrain:
-  """
-  Tokenizing and encoding/decoding text using the Sentencepiece tokenizer loaded with sentencepiece
+  Tokenizing and encoding/decoding text using the native sentencepiece library.
+  Supports both local and GCS (gs://) model paths.
   """
 
   def __init__(self, model_path: str, add_bos: bool, add_eos: bool):
     max_logging.log(f"Loading sentencepiece tokenizer: {model_path}")
     self._tokenizer_model = SentencePieceProcessor()
-    self._tokenizer_model.Load(model_path)
+    try:
+      if model_path.startswith("gs://"):
+        from maxtext.utils.gcs_utils import read_bytes_from_gcs  # pylint: disable=import-outside-toplevel
+
+        model_proto = read_bytes_from_gcs(model_path)
+        self._tokenizer_model.LoadFromSerializedProto(model_proto)
+      else:
+        self._tokenizer_model.Load(model_path)
+    except Exception as e:
+      raise ValueError(f"Failed to load sentencepiece tokenizer from {model_path}: {e}") from e
     self.pad_id = self._tokenizer_model.pad_id()
     self.unk_id = self._tokenizer_model.unk_id()
     self.bos_id = self._tokenizer_model.bos_id()
@@ -255,7 +240,7 @@ def decode(self, t: Sequence[int]) -> str:
     return self.tokenizer.decode(t)
 
 
-def build_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token, dataset_type):
+def build_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token):
   """Loads the tokenizer at `tokenizer_path`"""
   max_logging.log(f"Tokenizer path: {tokenizer_path}")
   if tokenizer_type == "tiktoken":
@@ -264,27 +249,6 @@ def build_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_
   elif tokenizer_type == "huggingface":
     return HFTokenizer(tokenizer_path, add_bos, add_eos, hf_access_token)
   elif tokenizer_type == "sentencepiece":
-    if dataset_type == "tfds":
-      return SentencePieceTokenizer(tokenizer_path, add_bos, add_eos)
-    else:
-      return SentencePieceTokenizerGrain(tokenizer_path, add_bos, add_eos)
+    return SentencePieceTokenizer(tokenizer_path, add_bos, add_eos)
   else:
     raise ValueError(f"Invalid tokenizer_type:{tokenizer_type} chosen in config")
-
-
-def TokenizeOp(tokenizer, features: Features, data_keys: Iterable[str] = ("inputs", "targets")) -> Features:
-  """Op for tokenization"""
-
-  def _process_string(string_tensor):
-    # Extract string value and decode it if necessary
-    string_value = string_tensor.numpy().decode("utf-8")
-    # encode and extract the tokenized integers
-    modified_string = tokenizer.encode(string_value)
-    return [modified_string]
-
-  for k in data_keys:
-    if isinstance(tokenizer, (TikTokenTokenizer, HFTokenizer)):
-      features[k] = tf.py_function(_process_string, [features[k]], Tout=[tf.int32])[0]
-    elif isinstance(tokenizer, SentencePieceTokenizer):
-      features[k] = tokenizer.encode(features[k])
-  return features

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -332,7 +332,6 @@ def train_distill(student_config: pyconfig.HyperParameters, teacher_config: pyco
       add_bos=student_config.add_bos,
       add_eos=student_config.add_eos,
       hf_access_token=student_config.hf_access_token,
-      dataset_type=student_config.dataset_type,
   )
   pad_id = tok.pad_id if tok.pad_id is not None else 0
 

src/maxtext/utils/gcs_utils.py

@@ -177,31 +177,45 @@ def gcs_glob_pattern(pattern):
   return data_files
 
 
-def read_json_from_gcs(file_path):
-  """
-  Read a json file from gcs bucket.
+def read_bytes_from_gcs(file_path):
+  """Read raw bytes from a GCS file.
 
   Args:
-    file_path: The gcs path of the json file.
+    file_path: The gcs path of the file (e.g. gs://bucket/path/to/file).
 
   Returns:
-    A dictionary with content from json file.
+    The file contents as bytes, or None if unavailable.
   """
-  if not _gcs_guard("read_json_from_gcs"):
+  if not _gcs_guard("read_bytes_from_gcs"):
     return None
   try:
     storage_client = storage.Client()
     bucket_name, file_prefix = parse_gcs_bucket_and_prefix(file_path)
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(file_prefix)
+    return blob.download_as_bytes()
+  except Exception as e:  # pylint: disable=broad-except
+    max_logging.log(f"Error reading bytes from GCS path {file_path}: {e}")
+    return None
 
-    json_string = blob.download_as_string()
 
-    data = json.loads(json_string)
+def read_json_from_gcs(file_path):
+  """
+  Read a json file from gcs bucket.
+
+  Args:
+    file_path: The gcs path of the json file.
 
-    return data
+  Returns:
+    A dictionary with content from json file.
+  """
+  try:
+    raw = read_bytes_from_gcs(file_path)
+    if raw is None:
+      return None
+    return json.loads(raw)
   except (ValueError, TypeError, json.JSONDecodeError) as e:
-    print(f"Error reading JSON file from GCS: {str(e)}")
+    max_logging.log(f"Error reading JSON file from GCS: {str(e)}")
     return None
 
 

tests/unit/tokenizer_test.py

@@ -67,7 +67,7 @@ def tearDownClass(cls):
   @pytest.mark.tpu_only
   def test_tokenize(self):
     text = "This is a test"
-    self.assertTrue(np.array_equal(self.source_tokenizer.encode(text).numpy(), self.test_tokenizer.encode(text).numpy()))
+    self.assertTrue(np.array_equal(self.source_tokenizer.encode(text), self.test_tokenizer.encode(text)))
 
   @pytest.mark.tpu_only
   def test_detokenize(self):