Help protect the Great Barrier Reef with TensorFlow on Kaggle Join Challenge

Converting TensorFlow Text operators to TensorFlow Lite

View on TensorFlow.org Run in Google Colab View on GitHub Download notebook

Overview

Machine learning models are frequently deployed using TensorFlow Lite to mobile, embedded, and IoT devices to improve data privacy and lower response times. These models often require support for text processing operations. TensorFlow Text version 2.7 and higher provides improved performance, reduced binary sizes, and operations specifically optimized for use in these environments.

Text operators

The following TensorFlow Text classes can be used from within a TensorFlow Lite model.

  • FastWordpieceTokenizer
  • WhitespaceTokenizer

Model Example

pip install -q -U tf-nightly
pip install -q -U tensorflow-text-nightly
from absl import app
import numpy as np
import tensorflow as tf
import tensorflow_text as tf_text

from tensorflow.lite.python import interpreter

The following code example shows the conversion process and interpretation in Python using a simple test model. Note that the output of a model cannot be a tf.RaggedTensor object when you are using TensorFlow Lite. However, you can return the components of a tf.RaggedTensor object or convert it using its to_tensor function. See the RaggedTensor guide for more details.

class TokenizerModel(tf.keras.Model):

  def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self.tokenizer = tf_text.WhitespaceTokenizer()

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[None], dtype=tf.string, name='input')
  ])
  def call(self, input_tensor):
    return { 'tokens': self.tokenizer.tokenize(input_tensor).flat_values }
# Test input data.
input_data = np.array(['Some minds are better kept apart'])

# Define a Keras model.
model = TokenizerModel()

# Perform TensorFlow Text inference.
tf_result = model(tf.constant(input_data))
print('TensorFlow result = ', tf_result['tokens'])
TensorFlow result =  tf.Tensor([b'Some' b'minds' b'are' b'better' b'kept' b'apart'], shape=(6,), dtype=string)

Convert the TensorFlow model to TensorFlow Lite

When converting a TensorFlow model with TensorFlow Text operators to TensorFlow Lite, you need to indicate to the TFLiteConverter that there are custom operators using the allow_custom_ops attribute as in the example below. You can then run the model conversion as you normally would. Review the TensorFlow Lite converter documentation for a detailed guide on the basics of model conversion.

# Convert to TensorFlow Lite.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
converter.allow_custom_ops = True
tflite_model = converter.convert()
2021-11-12 12:21:28.882819: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
INFO:tensorflow:Assets written to: /tmp/tmp6ev3e0l6/assets
2021-11-12 12:21:30.642172: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:351] Ignored output_format.
2021-11-12 12:21:30.642219: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:354] Ignored drop_control_dependency.
WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded
2021-11-12 12:21:30.854523: W tensorflow/compiler/mlir/lite/flatbuffer_export.cc:1903] The following operation(s) need TFLite custom op implementation(s):
Custom ops: TFText>WhitespaceTokenizeWithOffsetsV2
Details:
    tf.TFText>WhitespaceTokenizeWithOffsetsV2(tensor<?x!tf_type.string>, tensor<!tf_type.string>) -> (tensor<?x!tf_type.string>, tensor<?xi64>, tensor<?xi32>, tensor<?xi32>) : {device = ""}
See instructions: https://www.tensorflow.org/lite/guide/ops_custom

Inference

For the TensorFlow Lite interpreter to properly read your model containing TensorFlow Text operators, you must configure it to use these custom operators, and provide registration methods for them. Use tf_text.tflite_registrar.SELECT_TFTEXT_OPS to provide the full suite of registration functions for the supported TensorFlow Text operators to InterpreterWithCustomOps.

Note, that while the example below shows inference in Python, the steps are similar in other languages with some minor API translations, and the necessity to build the tflite_registrar into your binary. See TensorFlow Lite Inference for more details.

# Perform TensorFlow Lite inference.
interp = interpreter.InterpreterWithCustomOps(
    model_content=tflite_model,
    custom_op_registerers=tf_text.tflite_registrar.SELECT_TFTEXT_OPS)
interp.get_signature_list()
{'serving_default': {'inputs': ['input'], 'outputs': ['tokens']} }

Next, the TensorFlow Lite interpreter is invoked with the input, providing a result which matches the TensorFlow result from above.

tokenize = interp.get_signature_runner('serving_default')
output = tokenize(input=input_data)
print('TensorFlow Lite result = ', output['tokens'])
TensorFlow Lite result =  [b'Some' b'minds' b'are' b'better' b'kept' b'apart']