How to save TextVectorization to disk in tensorflow?

Instead of pickling the object, pickle the configuration and weights. Later unpickle it and use configuration to create the object and load the saved weights. Official docs here.

Code

text_dataset = tf.data.Dataset.from_tensor_slices([
                                                   "this is some clean text", 
                                                   "some more text", 
                                                   "even some more text"]) 
# Fit a TextVectorization layer
vectorizer = TextVectorization(max_tokens=10, output_mode='tf-idf',ngrams=None)    
vectorizer.adapt(text_dataset.batch(1024))

# Vector for word "this"
print (vectorizer("this"))

# Pickle the config and weights
pickle.dump({'config': vectorizer.get_config(),
             'weights': vectorizer.get_weights()}
            , open("tv_layer.pkl", "wb"))

print ("*"*10)
# Later you can unpickle and use 
# `config` to create object and 
# `weights` to load the trained weights. 

from_disk = pickle.load(open("tv_layer.pkl", "rb"))
new_v = TextVectorization.from_config(from_disk['config'])
# You have to call `adapt` with some dummy data (BUG in Keras)
new_v.adapt(tf.data.Dataset.from_tensor_slices(["xyz"]))
new_v.set_weights(from_disk['weights'])

# Lets see the Vector for word "this"
print (new_v("this"))

Output:

tf.Tensor(
[[0.         0.         0.         0.         0.91629076 0.
  0.         0.         0.         0.        ]], shape=(1, 10), dtype=float32)
**********
tf.Tensor(
[[0.         0.         0.         0.         0.91629076 0.
  0.         0.         0.         0.        ]], shape=(1, 10), dtype=float32)

Borrowing @jakub's model-vehicle trick - with which I couldn't get the model to load - I did in the end go via the JSON serialization route, as follows.

Note that you need to have tensorflow>=2.7 for the TextVectorization layer, and you need to use the same version to save and load the layer/model.

So, following on from the middle of @jakub's brilliant sample,

# Save.
model_json = model.to_json()
with open(filepath, "w") as model_json_fh:
    model_json_fh.write(model_json)

# Load.
with open(filepath, 'r') as model_json_fh:
    loaded_model = tf.keras.models.model_from_json(model_json_fh.read())
    vectorization_layer = loaded_model.layers[0]

loaded_model = tf.keras.models.load_model(filepath)
loaded_vectorizer = loaded_model.layers[0]

That's it.

I am not sure of the advantages of the one route over the other.

This also shows how it goes: https://machinelearningmastery.com/save-load-keras-deep-learning-models

And this helps with a JSON error you may encounter on your travels in these parts:

https://github.com/keras-team/keras/issues/6971

One can use a bit of a hack to do this. Construct your TextVectorization object, then put it in a model. Save the model to save the vectorizer. Loading the model will reproduce the vectorizer. See the example below.

import tensorflow as tf
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization

data = [
    "The sky is blue.",
    "Grass is green.",
    "Hunter2 is my password.",
]

# Create vectorizer.
text_dataset = tf.data.Dataset.from_tensor_slices(data)
vectorizer = TextVectorization(
    max_tokens=100000, output_mode='tf-idf', ngrams=None,
)
vectorizer.adapt(text_dataset.batch(1024))

# Create model.
model = tf.keras.models.Sequential()
model.add(tf.keras.Input(shape=(1,), dtype=tf.string))
model.add(vectorizer)

# Save.
filepath = "tmp-model"
model.save(filepath, save_format="tf")

# Load.
loaded_model = tf.keras.models.load_model(filepath)
loaded_vectorizer = loaded_model.layers[0]

Here is a test that both vectorizers (original and loaded) produce the same output.

import numpy as np

np.testing.assert_allclose(loaded_vectorizer("blue"), vectorizer("blue"))

For one anyone asking themselves how to get a dense tensor instead of ragged tensor when loading the config of a TextVectorization layer, try explicitly setting the output_mode. The problem is related to a very recent bug, where the output_mode is not set correctly when it comes from a saved configuration.

This results in a dense Tensor:

text_dataset = tf.data.Dataset.from_tensor_slices([
                                                   "this is some clean text", 
                                                   "some more text", 
                                                   "even some more text"]) 
vectorizer = TextVectorization(max_tokens=10, output_mode='int', output_sequence_length = 10)   
vectorizer.adapt(text_dataset.batch(1024))

print(vectorizer("this"))
pickle.dump({'config': vectorizer.get_config(),
             'weights': vectorizer.get_weights()}
            , open("tv_layer.pkl", "wb"))

from_disk = pickle.load(open("tv_layer.pkl", "rb"))
new_vectorizer = TextVectorization(max_tokens=from_disk['config']['max_tokens'],
                                          output_mode='int',
                                          output_sequence_length=from_disk['config']['output_sequence_length'])
new_vectorizer.adapt(tf.data.Dataset.from_tensor_slices(["xyz"]))
new_vectorizer.set_weights(from_disk['weights'])

print(new_vectorizer("this"))

tf.Tensor([5 0 0 0 0 0 0 0 0 0], shape=(10,), dtype=int64)
tf.Tensor([5 0 0 0 0 0 0 0 0 0], shape=(10,), dtype=int64)

This results in a ragged Tensor when loading:

import tensorflow as tf

text_dataset = tf.data.Dataset.from_tensor_slices([
                                                   "this is some clean text", 
                                                   "some more text", 
                                                   "even some more text"]) 
vectorizer = TextVectorization(max_tokens=10, output_mode='int', output_sequence_length = 10)   
vectorizer.adapt(text_dataset.batch(1024))

print(vectorizer("this"))
pickle.dump({'config': vectorizer.get_config(),
             'weights': vectorizer.get_weights()}
            , open("tv_layer.pkl", "wb"))

from_disk = pickle.load(open("tv_layer.pkl", "rb"))
new_vectorizer = TextVectorization(max_tokens=from_disk['config']['max_tokens'],
                                          output_mode=from_disk['config']['output_mode'],
                                          output_sequence_length=from_disk['config']['output_sequence_length'])
new_vectorizer.adapt(tf.data.Dataset.from_tensor_slices(["xyz"]))
new_vectorizer.set_weights(from_disk['weights'])

print(new_vectorizer("this"))

tf.Tensor([5 0 0 0 0 0 0 0 0 0], shape=(10,), dtype=int64)
tf.Tensor([5], shape=(1,), dtype=int64)