Fix Python workflow for WordPieceTokenizer by maitry63 · Pull Request #2963 · keras-team/keras-hub

maitry63 · 2026-08-17T08:58:21Z

This PR adds a Python workflow for WordPieceTokenizer, allowing tokenization and detokenization to run without requiring tensorflow_text.

Fixes: #2948

Changes :

Added pure-Python pre-tokenization.
Added a pure-Python greedy longest-match-first WordPiece implementation.
Added pure-Python detokenization.
Added Python-side vocabulary and token-to-ID lookups that do not depend on
tensorflow_text.
Made tensorflow_text optional for the Python workflow.
Preserved the existing tensorflow_text workflow.
Made the Python workflow the default when not inside a tf.function.
Added separate Python input canonicalization for tokenization and
detokenization.
Handled scalar TensorFlow string tensors correctly when converting their
.numpy() values from bytes to Python str.

Tests :

Default Python workflow.
Python workflow without tensorflow_text.
Explicit TensorFlow workflow.
Python and TensorFlow workflow equivalence.
Python detokenization.
Special tokens.
Lowercasing and accent stripping.
CJK tokenization.
Integer and string output dtypes.
Scalar and batched inputs.

Checklist

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have followed the Keras Hub Model contribution guidelines in making these changes.
I have followed the Keras Hub API design guidelines in making these changes.
I have signed the Contributor License Agreement.

gemini-code-assist

Code Review

This pull request adds a pure Python tokenization and detokenization workflow to WordPieceTokenizer so it can run without tensorflow_text installed, and includes matching unit tests. The review feedback points out two robustness issues in the input canonicalization logic: strict type checks in _canonicalize_python_detokenize_inputs can fail with NumPy scalar types (like np.int32), and 0D NumPy arrays can bypass tensor checks in _canonicalize_python_inputs and cause iteration errors. Suggestions were provided to make these checks more robust.

gemini-code-assist · 2026-08-17T09:01:18Z

+        if isinstance(inputs, (int, str, bytes)):
+            return [[inputs]], True
+        if isinstance(inputs, (list, tuple)):
+            if len(inputs) == 0 or isinstance(inputs[0], (int, str, bytes)):
+                return [list(inputs)], True
+            return [list(x) for x in inputs], False


Using strict type checks like isinstance(inputs, (int, str, bytes)) and isinstance(inputs[0], (int, str, bytes)) is fragile because it fails for non-standard scalar types such as NumPy integer/string scalars (e.g., np.int32, np.str_).

Specifically, if inputs is a 1D list of NumPy integers (e.g., [np.int32(1), np.int32(2)]), isinstance(inputs[0], (int, str, bytes)) evaluates to False. This causes the function to incorrectly treat the input as a 2D batch and execute [list(x) for x in inputs], which tries to call list(np.int32(1)) and raises TypeError: 'numpy.int32' object is not iterable.

Using negative checks (i.e., checking if the input or its elements are NOT containers like list, tuple, or np.ndarray) provides a robust and future-proof way to identify scalar values and 1D sequences of scalars.

Suggested change

if isinstance(inputs, (int, str, bytes)):

return [[inputs]], True

if isinstance(inputs, (list, tuple)):

if len(inputs) == 0 or isinstance(inputs[0], (int, str, bytes)):

return [list(inputs)], True

return [list(x) for x in inputs], False

if not isinstance(inputs, (list, tuple, np.ndarray)):

return [[inputs]], True

if isinstance(inputs, (list, tuple)):

if len(inputs) == 0 or not isinstance(inputs[0], (list, tuple, np.ndarray)):

return [list(inputs)], True

return [list(x) for x in inputs], False

References

Poke Holes in the Implementation: Your goal is to critically test the logic. Actively search for and point out failing edge cases, race conditions, or unhandled exceptions in the implementation. ^(link)

gemini-code-assist · 2026-08-17T09:01:18Z

+                else:
+                    inputs = inputs.tolist()
+
+        if isinstance(inputs, bytes):


If inputs is a raw np.ndarray (e.g., when passed directly or from non-TF backends), a 0D array (like np.array("hello")) will bypass the tf.Tensor check and cause a TypeError: iteration over a 0-d array when iterating in the self.split block.

Converting any np.ndarray to a Python scalar or list at the entry point of the canonicalization path resolves this and ensures consistent behavior across all backends.

Suggested change

if isinstance(inputs, bytes):

if isinstance(inputs, np.ndarray):

inputs = inputs.item() if inputs.ndim == 0 else inputs.tolist()

if isinstance(inputs, bytes):

References

Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. ^(link)

When handling inputs that may be passed as tensors in a Python-only execution path, perform defensive type conversion (e.g., converting to int) at the entry point of the path rather than deep inside helper functions, especially if multiple downstream helper functions perform arithmetic on the value.

…kenizer

divyashreepathihalli · 2026-08-26T20:27:54Z

+        split_regex = re.compile(
+            f"({keep_split_pattern})|(?:{WHITESPACE_REGEX})"
+        )
+        words = [piece for piece in split_regex.split(text) if piece]


Unicode whitespace — Python's \s is Unicode-wide; RE2's is ASCII-only. "the\u00a0fox" → Python [the, fox], TF [UNK] (same for U+2009, U+3000). Fix: a Python-side whitespace pattern using the regex module's scoped ASCII flag, (?a:\s)|\p{Cc}|\p{Cf}.

Hi @divyashreepathihalli , Thank you for your Review!

Separated whitespace patterns into WHITESPACE_REGEX_TF(using standard ASCII characters for RE2/TensorFlow Text compatibility) and
WHITESPACE_REGEX_PYTHON(using the regex module with the scoped ASCII flag (?a:\s)|\p{Cc}|\p{Cf}) to ensure exact parity across both workflows without breaking RE2 compilation.

divyashreepathihalli · 2026-08-26T20:29:09Z

+                for word in words
+            ]
+        else:
+            words = [word.casefold() for word in words]


str.casefold() ≠ tf_text.case_fold_utf8 (NFKC_CaseFold). "ＦＯＸ fox" → Python [UNK, fox], TF [fox, fox]. Fix: unicodedata.normalize("NFKC", word.casefold())

Updated the Python case-folding logic to apply NFKC normalization via unicodedata.normalize("NFKC", word.casefold()), matching the behavior of tf_text.case_fold_utf8.

divyashreepathihalli · 2026-08-26T20:29:59Z

+    vocabulary_set,
+    unknown_token,
+    suffix_indicator,
+    max_input_chars_per_word=200,


Python bails at 200 chars; FastWordpieceTokenizer defaults to max_bytes_per_word=100. A 120-byte word tokenizes on Python, [UNK] on TF. Fix: len(word.encode("utf-8")) > 100

Replaced character length bounds with UTF-8 byte-length checks (len(word.encode("utf-8")) > 100) in the pure Python WordPiece loop.

divyashreepathihalli · 2026-08-26T20:30:58Z

+        output = tokenizer(np.bytes_("the quick"))
+        self.assertAllEqual(output, [1, 2, 3])
+
+    def test_python_workflow_with_seq_len_returns_dense_output(self):


missing WordPieceTokenizerTFTest subclass re-running the suite with _allow_python_workflow=False

Added the WordPieceTokenizerTFTest test case subclass to the test suite to ensure the entire test set runs and passes with _allow_python_workflow=False.

divyashreepathihalli · 2026-08-26T20:31:22Z

+        output = tokenizer.detokenize([1, 2, 3])
+        self.assertAllEqual(output, "the quick")
+
+    def test_python_and_tf_workflows_match(self):


parity test too thin (add CJK, accents, special tokens, unicode whitespace, >100-byte words)

Expanded test_python_and_tf_workflows_match to validate parity across CJK characters, accent stripping, special tokens, unicode whitespace(\u00a0), and >100-byte words.

divyashreepathihalli · 2026-08-26T20:32:35Z

+        # workflow (`_tokenize_python`/`_detokenize_python`), which does not
+        # require `tensorflow_text` to be installed.
+        self._vocabulary_set = set(self.vocabulary)
+        self._token_to_id_python = {


duplicate-token token_to_id semantics

divyashreepathihalli · 2026-08-26T20:33:42Z

+        else:
+            tokens = batched_tokens
+        if output_is_int and self.sequence_length:
+            return np.asarray(tokens, dtype=self.compute_dtype)


np.asarray → backend tensor via call

Replaced raw NumPy conversion (np.asarray) with Keras tensor (keras.ops.convert_to_tensor).

…kenizer

Fix Python workflow for WordPieceTokenizer

14610c0

github-actions Bot requested review from divyashreepathihalli and laxmareddyp August 17, 2026 08:58

gemini-code-assist Bot reviewed Aug 17, 2026

View reviewed changes

maitry63 added 4 commits August 17, 2026 16:34

fix: word_piece detokenization workflow

dfdf0f7

fix:dense output

e1d34bb

fix: docstring test

e1c0657

Merge branch 'keras-team:master' into Fix_python_workflow_wordpieceto…

34284df

…kenizer

divyashreepathihalli requested changes Aug 26, 2026

View reviewed changes

maitry63 added 2 commits August 27, 2026 12:26

Merge branch 'keras-team:master' into Fix_python_workflow_wordpieceto…

dc3375f

…kenizer

Addressed review comments

923ce2c

Uh oh!

Conversation

maitry63 commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes :

Tests :

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 17, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

maitry63 commented Aug 17, 2026 •

edited

Loading