Migrate ByteTokenizer to PyGrain

keras_hub/src/tokenizers/byte_tokenizer.py

@@ -1,8 +1,12 @@
+import unicodedata
+
+import keras
 import numpy as np
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+from keras_hub.src.utils.tensor_utils import in_tf_function
 from keras_hub.src.utils.tensor_utils import is_int_dtype
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 
@@ -16,6 +20,39 @@
     tf_text = None
 
 
+def _decode_with_replacement(byte_seq, errors, replacement_char):
+    byte_seq = bytes(byte_seq)
+
+    # If using a custom replacement character
+    if errors == "replace" and replacement_char != 65533:
+        result = []
+        start = 0
+        while start < len(byte_seq):
+            try:
+                # Try to decode the remainder of the sequence
+                decoded = byte_seq[start:].decode("utf-8", errors="strict")
+                result.append(decoded)
+                break
+            except UnicodeDecodeError as e:
+                # Decode the valid chunk before the error
+                valid_part = byte_seq[start : start + e.start].decode(
+                    "utf-8", errors="strict"
+                )
+                result.append(valid_part)
+                # Append the custom replacement character
+                result.append(chr(replacement_char))
+                # Skip past the invalid bytes reported by the error
+                start = start + e.end
+
+        return "".join(result)
+
+    # Standard behavior for all other cases
+    try:
+        return byte_seq.decode("utf-8", errors=errors)
+    except UnicodeDecodeError as e:
+        raise ValueError(f"Invalid byte sequence: {e}")
+
+
 @keras_hub_export("keras_hub.tokenizers.ByteTokenizer")
 class ByteTokenizer(tokenizer.Tokenizer):
     """Raw byte tokenizer.
@@ -71,7 +108,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> tokenizer = keras_hub.tokenizers.ByteTokenizer()
     >>> outputs = tokenizer("hello")
     >>> np.array(outputs)
-    array([104, 101, 108, 108, 111], dtype=int32)
+    array([104, 101, 108, 108, 111])
 
     Ragged outputs.
     >>> inputs = ["hello", "hi"]
@@ -87,9 +124,9 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> tokenizer = keras_hub.tokenizers.ByteTokenizer(sequence_length=8)
     >>> seq1, seq2 = tokenizer(inputs)
     >>> np.array(seq1)
-    array([104, 101, 108, 108, 111,   0,   0,   0], dtype=int32)
+    array([104, 101, 108, 108, 111,   0,   0,   0])
     >>> np.array(seq2)
-    array([104, 105,   0,   0,   0,   0,   0,   0], dtype=int32)
+    array([104, 105,   0,   0,   0,   0,   0,   0])
 
     Tokenize, then batch for ragged outputs.
     >>> tokenizer = keras_hub.tokenizers.ByteTokenizer()
@@ -170,17 +207,24 @@ def __init__(
                 f"Received: errors={errors}"
             )
 
-        super().__init__(dtype=dtype, **kwargs)
+        _allow_python_workflow = kwargs.pop("_allow_python_workflow", True)
+        super().__init__(
+            dtype=dtype, _allow_python_workflow=_allow_python_workflow, **kwargs
+        )
 
         self.lowercase = lowercase
         self.sequence_length = sequence_length
         self.normalization_form = normalization_form
         self.errors = errors
         self.replacement_char = replacement_char
 
-        self._char_lst = tf.constant(
-            [i.tobytes() for i in np.arange(256, dtype=np.uint8)]
-        )
+        self._char_lst = [i.tobytes() for i in np.arange(256, dtype=np.uint8)]
+
+        # Build the tensor once, safely outside of any graph compilation
+        self._char_lst_tensor = None
+        if tf is not None:
+            self._char_lst_tensor = tf.constant(self._char_lst)
+
         self._update_special_token_ids()
 
     def vocabulary_size(self):
@@ -193,8 +237,14 @@ def get_vocabulary(self):
             vocab[chr(i)] = i
         return vocab
 
-    @preprocessing_function
     def tokenize(self, inputs):
+        if not self._allow_python_workflow or in_tf_function():
+            return self._tokenize_tf(inputs)
+        else:
+            return self._tokenize_python(inputs)
+
+    @preprocessing_function
+    def _tokenize_tf(self, inputs):
         unbatched = inputs.shape.rank == 0
         if unbatched:
             inputs = tf.expand_dims(inputs, 0)
@@ -224,15 +274,70 @@ def tokenize(self, inputs):
             tokens = tf.squeeze(tokens, 0)
         return tokens
 
-    @preprocessing_function
+    def _tokenize_python(self, inputs):
+        def _canonicalize_tokenize_inputs(inputs):
+            if isinstance(inputs, str):
+                return [inputs], False
+            elif isinstance(inputs, (tuple, list)):
+                return list(inputs), True
+            elif tf is not None and isinstance(inputs, tf.Tensor):
+                unbatched = inputs.shape.rank == 0
+                if unbatched:
+                    inputs = tf.expand_dims(inputs, 0)
+                inputs = inputs.numpy().tolist()
+                inputs = keras.tree.map_structure(
+                    lambda x: x.decode("utf-8"), inputs
+                )
+                return inputs, not unbatched
+            else:
+                raise ValueError(
+                    "Input should be a string or a list of strings. "
+                    f"Received: {inputs}"
+                )
+
+        inputs, batched = _canonicalize_tokenize_inputs(inputs)
+
+        def tokenize_single_string(text):
+            if isinstance(text, bytes):
+                text = text.decode("utf-8")
+            if self.lowercase:
+                text = text.casefold()
+            if self.normalization_form is not None:
+                text = unicodedata.normalize(self.normalization_form, text)
+            # Convert to byte integers
+            seq = list(text.encode("utf-8"))
+
+            # Handle sequence_length truncation and padding
+            if self.sequence_length:
+                pad_token_id = getattr(self, "pad_token_id", 0)
+                seq = seq[: self.sequence_length] + [pad_token_id] * max(
+                    0, self.sequence_length - len(seq)
+                )
+            return seq
+
+        batched_tokens = keras.tree.map_structure(
+            tokenize_single_string, inputs
+        )
+
+        if not batched:
+            batched_tokens = batched_tokens[0]
+        return batched_tokens
+
     def detokenize(self, inputs):
+        if not self._allow_python_workflow or in_tf_function():
+            return self._detokenize_tf(inputs)
+        else:
+            return self._detokenize_python(inputs)
+
+    @preprocessing_function
+    def _detokenize_tf(self, inputs):
         inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
         # Remove trailing padding tokens, so that trailing "\x00" bytes don't
         # show up in the detokenized output.
         inputs = tf.ragged.boolean_mask(inputs, tf.not_equal(inputs, 0))
 
         outputs = tf.strings.reduce_join(
-            tf.gather(self._char_lst, inputs), axis=-1
+            tf.gather(self._char_lst_tensor, inputs), axis=-1
         )
 
         # Handle errors if an invalid byte sequence is encountered.
@@ -247,6 +352,65 @@ def detokenize(self, inputs):
             outputs = tf.squeeze(outputs, 0)
         return outputs
 
+    def _detokenize_python(self, inputs):
+        def _canonicalize_detokenize_inputs(inputs):
+            if tf is not None and isinstance(
+                inputs, (tf.Tensor, tf.RaggedTensor)
+            ):
+                if isinstance(inputs, tf.RaggedTensor):
+                    inputs = inputs.to_list()
+                else:
+                    inputs = inputs.numpy().tolist()
+            is_batched = True
+            if isinstance(inputs, int):
+                inputs = [[inputs]]
+                is_batched = False
+            elif isinstance(inputs, (tuple, list)):
+                if not inputs or isinstance(inputs[0], int):
+                    inputs = [list(inputs)]
+                    is_batched = False
+                else:
+                    inputs = [list(seq) for seq in inputs]
+            elif isinstance(inputs, np.ndarray) or keras.ops.is_tensor(inputs):
+                inputs = keras.ops.convert_to_numpy(inputs)
+                if inputs.ndim == 0:
+                    inputs = [[inputs.item()]]
+                    is_batched = False
+                elif inputs.ndim == 1:
+                    inputs = [inputs.tolist()]
+                    is_batched = False
+                elif inputs.ndim == 2:
+                    inputs = inputs.tolist()
+                else:
+                    raise ValueError(
+                        "Array must be 0, 1 or 2 dimensional. "
+                        f"Received: {inputs.shape}"
+                    )
+            else:
+                raise ValueError(
+                    "Input should be an integer, a list of integers, backend "
+                    f"tensor or numpy array. Received: {inputs}"
+                )
+            return inputs, is_batched
+
+        inputs, batched = _canonicalize_detokenize_inputs(inputs)
+
+        outputs = []
+        for seq in inputs:
+            # Remove padding tokens, so that trailing "\x00" bytes don't
+            # show up in the detokenized output.
+            # Using bytes().replace() executes directly in C for maximum speed
+            seq_bytes = bytes(seq).replace(b"\x00", b"")
+
+            decoded = _decode_with_replacement(
+                seq_bytes, self.errors, self.replacement_char
+            )
+            outputs.append(decoded)
+
+        if not batched:
+            outputs = outputs[0]
+        return outputs
+
     def id_to_token(self, id):
         """Convert an integer id to a string token."""
         if id >= self.vocabulary_size() or id < 0:

keras_hub/src/tokenizers/byte_tokenizer_test.py

@@ -86,9 +86,41 @@ def test_detokenize_strict_error(self):
         input_data = [[104, 101, 226, 150, 108, 108, 111]]
 
         tokenizer = ByteTokenizer(errors="strict")
-        with self.assertRaises(tf.errors.InvalidArgumentError):
+        expected_errors = (ValueError,)
+        if tf is not None:
+            expected_errors = (ValueError, tf.errors.InvalidArgumentError)
+        with self.assertRaises(expected_errors):
             _ = tokenizer.detokenize(input_data)
 
+    def test_detokenize_replace_valid_chars(self):
+        # 255 is invalid, 239,191,189 is valid U+FFFD.
+        # The invalid byte should be replaced by 'H' (72), but the valid
+        # U+FFFD should remain.
+        input_data = [[104, 101, 255, 108, 108, 111, 239, 191, 189]]
+        tokenizer = ByteTokenizer(errors="replace", replacement_char=72)
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, ["heHllo\ufffd"])
+
+    def test_workflow_parity(self):
+        if tf is None:
+            return  # Skip if TensorFlow is not available
+
+        input_data = ["hello", "fun", "▀▁▂▃", "haha"]
+        tokenizer = ByteTokenizer(sequence_length=12)
+
+        # Force TF Workflow
+        tokenizer._allow_python_workflow = False
+        tf_out = tokenizer(input_data)
+        tf_detok = tokenizer.detokenize(tf_out)
+
+        # Force Python Workflow
+        tokenizer._allow_python_workflow = True
+        python_out = tokenizer(input_data)
+        python_detok = tokenizer.detokenize(python_out)
+
+        self.assertAllEqual(tf_out, python_out)
+        self.assertAllEqual(tf_detok, python_detok)
+
     def test_vocab_size(self):
         tokenizer = ByteTokenizer()
         self.assertEqual(tokenizer.vocabulary_size(), 256)