Fix text vectorization serialization for custom callables.

PiperOrigin-RevId: 604658738

tf_keras/layers/preprocessing/text_vectorization.py

@@ -24,6 +24,7 @@
 from tf_keras.layers.preprocessing import string_lookup
 from tf_keras.saving.legacy.saved_model import layer_serialization
 from tf_keras.saving.serialization_lib import deserialize_keras_object
+from tf_keras.saving.serialization_lib import serialize_keras_object
 from tf_keras.utils import layer_utils
 from tf_keras.utils import tf_utils
 
@@ -500,8 +501,8 @@ def vocabulary_size(self):
  def get_config(self):
  config = {
  "max_tokens": self._lookup_layer.max_tokens,
- "standardize": self._standardize,
- "split": self._split,
+ "standardize": serialize_keras_object(self._standardize),
+ "split": serialize_keras_object(self._split),
  "ngrams": self._ngrams_arg,
  "output_mode": self._output_mode,
  "output_sequence_length": self._output_sequence_length,

tf_keras/layers/preprocessing/text_vectorization_test.py

@@ -2408,6 +2408,27 @@ def test_serialization_with_custom_callables(self):
  new_output_dataset = new_model.predict(input_array)
  self.assertAllEqual(expected_output, new_output_dataset)
 
+ def test_cloning_with_custom_callable(self):
+ @register_keras_serializable(package="Test")
+ def pipe_split_fn(inp):
+ return tf.strings.split(inp, sep="|")
+
+ text_dataset = tf.data.Dataset.from_tensor_slices(
+ [
+ "this|is|some|pipe-delimited|text",
+ "some|more|pipe-delimited|text",
+ "yet|more|pipe-delimited|text",
+ ]
+ )
+ vectorizer = text_vectorization.TextVectorization(
+ max_tokens=10, standardize=None, split=pipe_split_fn
+ )
+ vectorizer.adapt(text_dataset)
+ input_data = keras.Input(shape=(), dtype=tf.string)
+ outputs = vectorizer(input_data)
+ model = keras.Model(inputs=input_data, outputs=outputs)
+ _ = keras.models.clone_model(model)
+
  @test_utils.run_v2_only()
  def test_saving_v3(self):
  vocab_data = ["earth", "wind", "and", "fire"]