Utilizing TensorFlow Text with Whitespace Tokenizer in Python

Rate this post

πŸ’‘ Problem Formulation: In natural language processing, tokenization is a foundational step. Given a string of text, such as “TensorFlow is powerful and user-friendly!”, we want to split the text into tokens (words or symbols) based on whitespace to get an array of tokens: [“TensorFlow”, “is”, “powerful”, “and”, “user-friendly!”]. In Python, TensorFlow Text provides various tokenizers, including one for this exact purpose.

Method 1: Basic Usage of Whitespace Tokenizer

The TensorFlow Text library provides a WhitespaceTokenizer, which can tokenize a string of text on spaces, tabs, and newlines. This tokenizer is straightforward and ideal for texts where whitespace is the only delimiter between words.

Here’s an example:

import tensorflow as tf
import tensorflow_text as text

# Example text input
text_input = ["TensorFlow hub is amazing!"]

# Initialize the tokenizer
tokenizer = text.WhitespaceTokenizer()

# Tokenize the input text
tokens = tokenizer.tokenize(text_input)



[[b'TensorFlow', b'hub', b'is', b'amazing!']]

The code snippet initializes a WhitespaceTokenizer from the TensorFlow Text library and applies it to an array containing a single string. The text is tokenized based on whitespace, and the output is a list of tokens.

Method 2: Handling Multiple Text Inputs

The WhitespaceTokenizer can also handle multiple text inputs at once. It helps in processing batch data more efficiently when dealing with models expecting batch inputs.

Here’s an example:

texts = ["Welcome to TensorFlow world.", "Whitespace Tokenizer is useful!"]

# Tokenize batch inputs
batch_tokens = tokenizer.tokenize(texts)



[[b'Welcome', b'to', b'TensorFlow', b'world.'], [b'Whitespace', b'Tokenizer', b'is', b'useful!']]

The given example tokenizes a list of strings. Each string in the list is independently tokenized, and the result is a list of lists, each containing tokens for the corresponding input string.

Method 3: Integrating Tokenizer into TensorFlow Data Pipelines

WhitespaceTokenizer can be integrated within TensorFlow’s data pipeline, using it directly within a tf.data.Dataset map function for on-the-fly tokenization during the training or inference.

Here’s an example:

dataset = tf.data.Dataset.from_tensor_slices(["TensorFlow makes machine learning easy.", "Enjoy processing text with TensorFlow Text!"])

# Map the tokenizer onto the dataset
tokenized_dataset = dataset.map(lambda x: tokenizer.tokenize(x))

for tokens in tokenized_dataset:


[b'TensorFlow' b'makes' b'machine' b'learning' b'easy.']
[b'Enjoy' b'processing' b'text' b'with' b'TensorFlow' b'Text!']

The example tokenizes each element within a TensorFlow data dataset. The tokenizer is applied using the map function, so tokenization is performed as data passes through the pipeline, making it efficient and scalable.

Method 4: Customizing Tokenization with TensorFlow Ops

While the WhitespaceTokenizer is handy, sometimes it’s necessary to combine it with other TensorFlow ops for customized tokenization strategies. For example, you might remove punctuation or split tokens further.

Here’s an example:

import tensorflow as tf

# Define custom tokenization function
def custom_tokenization(text):
    # Use TensorFlow ops to replace punctuation with spaces
    text = tf.strings.regex_replace(text, '[^\\w\\s]', ' ')
    # Tokenize with the WhitespaceTokenizer
    return tokenizer.tokenize(text)

# Example usage
text_example = custom_tokenization("TensorFlow's tokenizer can handle complex text!")


[b'TensorFlow' b's' b'tokenizer' b'can' b'handle' b'complex' b'text']

The custom tokenization function first uses a TensorFlow string operation to replace punctuation with spaces. Then, it applies the WhitespaceTokenizer for final tokenization. This approach offers greater flexibility and can be fine-tuned for specific text preprocessing needs.

Bonus One-Liner Method 5: Inline Tokenization for Rapid Testing

Sometimes you might want to quickly tokenize text inline, especially while testing or prototyping. TensorFlow Text makes this surprisingly simple.

Here’s an example:

# One-liner tokenization
single_line_tokens = tokenizer.tokenize(["Quick example: TensorFlow text works!"]).numpy()


[[b'Quick' b'example:' b'TensorFlow' b'text' b'works!']]

With a single line of code, the text is tokenized into its component words. It’s particularly useful for quick, one-off tokenization tasks during development or while experimenting with different types of input text.


  • Method 1: Basic Usage of Whitespace Tokenizer. Essential for quick and simple tokenization. Does not handle punctuation or special characters.
  • Method 2: Handling Multiple Text Inputs. Efficient for batch processing. May require additional preprocessing if texts vary greatly in structure.
  • Method 3: Integrating Tokenizer into TensorFlow Data Pipelines. Suitable for seamless integration with TensorFlow pipelines, improving performance during model training or inference.
  • Method 4: Customizing Tokenization with TensorFlow Ops. Offers customization for complex tokenization tasks. Could introduce additional complexity and slow down processing if not optimized.
  • Bonus Method 5: Inline Tokenization for Rapid Testing. Quick and convenient for testing, less practical for production or complex data processing workflows.