StringLookup layer does not retrieve vocabulary after saving and loading the model
Hi @brunolucatto,

Pickling isn't natively supported in keras as it is generally [not secure](https://docs.python.org/3/library/pickle.html). Therefore I don't think this is actionable on Keras side, and overriding methods for `tf.keras.layers.serialize` is necessary. 

Please refer to the similar issue [here](https://github.com/keras-team/tf-keras/issues/626).

Please look at the [gist](https://colab.research.google.com/gist/Varsha-anjanappa/f87d4145f9d2a2a9e46069790154b4f6/61369.ipynb) here.

Please let us know if it works.

Thank you!!
 Hi @Varsha-anjanappa, thank you so much for the prompt reply.

I am afraid your example on gist missed the part where it fails, which is in the deserialization, not in the serialization. I added the tf.keras.layers.deserialize function, and it throws the same error as the pickle one.

```
import tensorflow as tf
import pickle

model_input = tf.keras.Input(shape=(1,), dtype=tf.int64)
lookup = tf.keras.layers.StringLookup(vocabulary=['a', 'b'])(model_input)
lookup = tf.keras.layers.Flatten()(lookup)
output = tf.keras.layers.Dense(10)(lookup)
full_model = tf.keras.Model(model_input, output)

# this part should work
model_bytes = pickle.dumps(full_model)
model_recovered = pickle.loads(model_bytes)


# this part should throw an error
full_model.save("/tmp/temp_model")
full_model_loaded = tf.keras.models.load_model("/tmp/temp_model")
model_bytes_2 = tf.keras.layers.serialize(full_model_loaded)
model_recovered_2 = tf.keras.layers.deserialize(model_bytes_2)
```

Moreover, I also added the proposed solution for the problem (i.e., changing the get_config method of the StringLookup class) to show that it also works using Keras' proper serialization.

```
import tensorflow as tf
import pickle

@tf.keras.utils.register_keras_serializable()
class MyStringLookup(tf.keras.layers.StringLookup):
    def get_config(self):
        base_config = super().get_config()
        custom = {"vocabulary": self.get_vocabulary()}
        return {**base_config, **custom}

model_input = tf.keras.Input(shape=(1,), dtype=tf.int64)
lookup = MyStringLookup(vocabulary=['a', 'b'])(model_input)
lookup = tf.keras.layers.Flatten()(lookup)
output = tf.keras.layers.Dense(10)(lookup)
full_model = tf.keras.Model(model_input, output)

# this part should work
model_bytes = pickle.dumps(full_model)
model_recovered = pickle.loads(model_bytes)


# this part should throw an error
full_model.save("/tmp/temp_model")
full_model_loaded = tf.keras.models.load_model("/tmp/temp_model")
model_bytes_2 = tf.keras.layers.serialize(full_model_loaded)
model_recovered_2 = tf.keras.layers.deserialize(model_bytes_2)
```

#-----------------

As a side note, I am using pickle because I want to distribute the model on spark using broadcast, so I can make predictions using a pandas UDF more efficiently. The implementation of pyspark.broadcast itself uses pickle, I tried to change that by creating a custom broadcast function but it was not an easy task for me. Loading the model from the disk inside the pandas UDF is also an alternative to broadcast, but it is many times slower.

 Hi @Varsha-anjanappa,

Just following up to know if you were able to reproduce the error with the changes above.

Please let me know if I can provide any further information. Hi @brunolucatto ,

Please find the attached [gist](https://colab.research.google.com/gist/Varsha-anjanappa/89c9def29103fc1b3aa3735f427362a7/61369_v2.ipynb).

As you can see in the gist, when we serialize, the vocabulary values aren't captured, it gives an empty list. Whereas the vocab size is getting captured. There is some issue in the source code. We'll check and get back to you.

Thank you  @SuryanarayanaY , can you please look into it we are seeing the same issue using tf 2.6.2 with python 3.6.9 @brunolucatto ,

It seems a valid issue for me. It needs to dig more for the root cause. I have done some work, on this bug. The issue seems to be in this file load.py. 
![image](https://github.com/tensorflow/tensorflow/assets/46747610/811dbd24-2090-45ea-95a6-989b6997c2a2)

As I have highlighted, parse from string is generating the data as nodes via proto buff.

![image](https://github.com/tensorflow/tensorflow/assets/46747610/150c769d-f668-4c6e-911e-3d628c3643e5)

The output nodes are as follows, 

nodes {
  node_id: 1
  node_path: "root.layer-0"
  identifier: "_tf_keras_input_layer"
  metadata: "{\"class_name\": \"InputLayer\", \"name\": \"input_1\", \"dtype\": \"int64\", \"sparse\": false, \"ragged\": false, \"batch_input_shape\": {\"class_name\": \"__tuple__\", \"items\": [null, 1]}, \"config\": {\"batch_input_shape\": {\"class_name\": \"__tuple__\", \"items\": [null, 1]}, \"dtype\": \"int64\", \"sparse\": false, \"ragged\": false, \"name\": \"input_1\"}}"
  version {
    producer: 2
    min_consumer: 1
  }
}
nodes {
  node_id: 2
  node_path: "root.layer-1"
  identifier: "_tf_keras_layer"
  metadata: "{\"name\": \"string_lookup\", \"trainable\": true, \"expects_training_arg\": false, \"dtype\": \"int64\", \"batch_input_shape\": null, \"stateful\": false, \"must_restore_from_config\": true, \"preserve_input_structure_in_config\": false, \"autocast\": true, \"class_name\": \"StringLookup\", \"config\": {\"name\": \"string_lookup\", \"trainable\": true, \"dtype\": \"int64\", \"invert\": false, \"max_tokens\": null, \"num_oov_indices\": 1, \"oov_token\": \"[UNK]\", \"mask_token\": null, \"output_mode\": \"int\", \"sparse\": false, \"pad_to_max_tokens\": false, \"idf_weights\": null, \"vocabulary\": null, \"vocabulary_size\": 3, \"encoding\": \"utf-8\", \"has_input_vocabulary\": true}, \"inbound_nodes\": [[[\"input_1\", 0, 0, {}]]], \"shared_object_id\": 1, \"build_input_shape\": {\"class_name\": \"TensorShape\", \"items\": [null, 1]}}"
  version {
    producer: 2
    min_consumer: 1
  }
}
nodes {
  node_id: 3
  node_path: "root.layer_with_weights-0"
  identifier: "_tf_keras_layer"
  metadata: "{\"name\": \"dense\", \"trainable\": true, \"expects_training_arg\": false, \"dtype\": \"float32\", \"batch_input_shape\": null, \"stateful\": false, \"must_restore_from_config\": false, \"preserve_input_structure_in_config\": false, \"autocast\": true, \"class_name\": \"Dense\", \"config\": {\"name\": \"dense\", \"trainable\": true, \"dtype\": \"float32\", \"units\": 10, \"activation\": \"linear\", \"use_bias\": true, \"kernel_initializer\": {\"class_name\": \"GlorotUniform\", \"config\": {\"seed\": null}, \"shared_object_id\": 2}, \"bias_initializer\": {\"class_name\": \"Zeros\", \"config\": {}, \"shared_object_id\": 3}, \"kernel_regularizer\": null, \"bias_regularizer\": null, \"activity_regularizer\": null, \"kernel_constraint\": null, \"bias_constraint\": null}, \"inbound_nodes\": [[[\"string_lookup\", 0, 0, {}]]], \"shared_object_id\": 4, \"input_spec\": {\"class_name\": \"InputSpec\", \"config\": {\"dtype\": null, \"shape\": null, \"ndim\": null, \"max_ndim\": null, \"min_ndim\": 2, \"axes\": {\"-1\": 1}}, \"shared_object_id\": 7}, \"build_input_shape\": {\"class_name\": \"TensorShape\", \"items\": [null, 1]}}"
  version {
    producer: 2
    min_consumer: 1
  }

in node_id:2 the vocabulary is getting dropped where as it is present in file content. Issue seems to be with protobuff and parsing. It's 2 Am and I have worked on this for straight 4 hours and am taking a break. If you can wait will resume working on it 🙂.


 My bad, I think that there is some mistake in how model is saving metadata for the layers.

Opened the model and looked at metadata:
![image](https://github.com/tensorflow/tensorflow/assets/46747610/b4934ce9-bb39-45be-9ac9-99fa561ff1cc)

![image](https://github.com/tensorflow/tensorflow/assets/46747610/1df343b8-31dc-4a0a-92ea-2fe5d84b99b5) Hi folks, has there been any work on this by any chance or is there a workaround for this? I am currently running into this issue as well.  @sharyar Isn't the workaround proposed in the original post working in your case?

> We were able to circumvent the issue by creating a new class as follows:
> 
> ```
> @tf.keras.utils.register_keras_serializable()
> class MyStringLookup(tf.keras.layers.StringLookup):
>     def get_config(self):
>         base_config = super().get_config()
>         custom = {"vocabulary": self.get_vocabulary()}
>         return {**base_config, **custom}
> ``` > We were able to circumvent the issue by creating a new class as follows:
> 
> ```
> @tf.keras.utils.register_keras_serializable()
> class MyStringLookup(tf.keras.layers.StringLookup):
>     def get_config(self):
>         base_config = super().get_config()
>         custom = {"vocabulary": self.get_vocabulary()}
>         return {**base_config, **custom}
> ```
> 
> However, it would be nice if we didn't have to create this wrapper.

Hi @brunolucatto ,
The decorator tf.keras.utils.register_keras_serializable() is one of the way for serializing a class/method into Keras object which can then be saved as .keras model and the able to reload successfully.

There are 3 ways we can do this as mentioned here in [custom_objects](https://www.tensorflow.org/guide/keras/serialization_and_saving#custom_objects) serialization_and_saving guide . You can choose any one.


 Hi @SuryanarayanaY, thanks for the reply!

I think the main issue is not the decorator, but the fact that we need to create this class to make the serialization work. The only thing this new class is doing is "saving again" the vocabulary attribute of the instance, and all of this would be unnecessary if the StringLookup layer serialization was working properly. (*)

Creating this class is a somewhat ugly workaround for the problem, as it does not address the root cause directly. This is the reason I did not submit it as a potential fix, although it has been working nicely so far.

(*) _In the original post, I referred to this new class as a wrapper, which in hindsight makes it look like it was only the decorator that was making me unhappy. Sorry for the ambiguity there._ I ran into this issue recently when passing a keras model object to sklearn [permutation_importance()](https://scikit-learn.org/stable/modules/generated/sklearn.inspection.permutation_importance.html) and was hitting this error when using `n_jobs` in that function. 

In my case i had a model object that has `StringLookup` layers i need to replace with the `MyStringLookup` layers. Using [`clone_model`](https://www.tensorflow.org/api_docs/python/tf/keras/models/clone_model) helped like this: 

```
import tensorflow as tf
import pickle

tf.keras.utils.get_custom_objects().clear()


@tf.keras.utils.register_keras_serializable()
class MyStringLookup(tf.keras.layers.StringLookup):
    def get_config(self):
        base_config = super().get_config()
        vocabulary = self.get_vocabulary()
        custom = {"vocabulary": vocabulary}
        return {**base_config, **custom}


def my_clone_function(layer):
    if isinstance(layer, tf.keras.layers.StringLookup):
        clone_layer = MyStringLookup(vocabulary=layer.get_vocabulary())
        return clone_layer
    return layer


# create model
model_input = tf.keras.Input(shape=(1,), dtype=tf.int64)
lookup = tf.keras.layers.StringLookup(vocabulary=['a', 'b'])(model_input)
output = tf.keras.layers.Dense(10)(lookup)
model = tf.keras.Model(model_input, output)

# model StringLookup has vocabulary set
print('\ncorrect:\n')
print(model.layers[1].get_config()['vocabulary'])
# ListWrapper(['a', 'b'])

# save model
model.save("/tmp/temp_model")

# load model
model_orginal_reloaded = tf.keras.models.load_model("/tmp/temp_model")

# model StringLookup now has no vocabulary set
print('\nbroken:\n')
print(model_orginal_reloaded.layers[1].get_config()['vocabulary'])
# ListWrapper([])

# fix the model
model_fixed = tf.keras.models.clone_model(
    model,
    clone_function=my_clone_function
)

# we now see a correct model vocab in the MyStringLookup layer
print('\nfixed:\n')
print(model_fixed.layers[1].get_config()['vocabulary'])
```

[Here is a google colab](https://colab.research.google.com/drive/1qmdgR7a06oHQgFs1_jiOWggQJJJyYrqV?usp=sharing) to show what i mean.

Just sharing this in case is useful for anyone else in a similar situation.
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.11, 2.12, 2.13

### Custom code

No

### OS platform and distribution

Linux Ubuntu 20.04.1

### Mobile device

_No response_

### Python version

3.9.5

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

We noticed we could pickle the model right after building it, but unpickling would fail after saving and loading it from the disk. Upon further investigation, we realized the error was due to the vocabulary of the StringLookup layer, which was becoming an empty list after the tf.keras.models.load_model operation. 

The unpickling issue happens on TF 2.11 onwards. The unpickling worked on TF 2.8, 2.9, and 2.10, even though the empty vocabulary issue was still there.

#--------------------

Using the minimal reproducible example below, before saving the model, if we inspect the StringLookup layer we get:
```
full_model.layers[1].get_config()

Out[9]: {'name': 'string_lookup',
 'trainable': True,
 'dtype': 'int64',
 'invert': False,
 'max_tokens': None,
 'num_oov_indices': 1,
 'oov_token': '[UNK]',
 'mask_token': None,
 'output_mode': 'int',
 'sparse': False,
 'pad_to_max_tokens': False,
 'vocabulary': ListWrapper(['a', 'b']),
 'idf_weights': None,
 'encoding': 'utf-8'}
```

After saving and loading to the disk, we get:
```
full_model_loaded.layers[1].get_config()

Out[10]: {'name': 'string_lookup',
 'trainable': True,
 'dtype': 'int64',
 'invert': False,
 'max_tokens': None,
 'num_oov_indices': 1,
 'oov_token': '[UNK]',
 'mask_token': None,
 'output_mode': 'int',
 'sparse': False,
 'pad_to_max_tokens': False,
 'vocabulary': ListWrapper([]),
 'idf_weights': None,
 'encoding': 'utf-8'}
 ```

#-----------------------


We were able to circumvent the issue by creating a new class as follows:
```
@tf.keras.utils.register_keras_serializable()
class MyStringLookup(tf.keras.layers.StringLookup):
    def get_config(self):
        base_config = super().get_config()
        custom = {"vocabulary": self.get_vocabulary()}
        return {**base_config, **custom}
```

However, it would be nice if we didn't have to create this wrapper.


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import pickle

model_input = tf.keras.Input(shape=(1,), dtype=tf.int64)
lookup = tf.keras.layers.StringLookup(vocabulary=['a', 'b'])(model_input)
output = tf.keras.layers.Dense(10)(lookup)
full_model = tf.keras.Model(model_input, output)

# this part should work
model_bytes = pickle.dumps(full_model)
model_recovered = pickle.loads(model_bytes)


# this part should throw an error
full_model.save("/tmp/temp_model")
full_model_loaded = tf.keras.models.load_model("/tmp/temp_model")
model_bytes_2 = pickle.dumps(full_model_loaded)
model_recovered_2 = pickle.loads(model_bytes_2)
```


### Relevant log output

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File <command-901450068412846>:1
----> 1 model_recovered_2 = pickle.loads(model_bytes_2)

File /databricks/python/lib/python3.9/site-packages/keras/saving/pickle_utils.py:48, in deserialize_model_from_bytecode(serialized_model)
     46     model = saving_lib.load_model(filepath)
     47 except Exception as e:
---> 48     raise e
     49 else:
     50     return model

File /databricks/python/lib/python3.9/site-packages/keras/saving/pickle_utils.py:46, in deserialize_model_from_bytecode(serialized_model)
     40         f.write(serialized_model)
     41     # When loading, direct import will work for most custom objects
     42     # though it will require get_config() to be implemented.
     43     # Some custom objects (e.g. an activation in a Dense layer,
     44     # serialized as a string by Dense.get_config()) will require
     45     # a custom_object_scope.
---> 46     model = saving_lib.load_model(filepath)
     47 except Exception as e:
     48     raise e

File /databricks/python/lib/python3.9/site-packages/keras/saving/experimental/saving_lib.py:196, in load_model(filepath, custom_objects)
    194     h5_file.close()
    195 except Exception as e:
--> 196     raise e
    197 else:
    198     return model

File /databricks/python/lib/python3.9/site-packages/keras/saving/experimental/saving_lib.py:183, in load_model(filepath, custom_objects)
    181 config_dict = json.loads(config_json)
    182 # Construct the model from the configuration file in the archive.
--> 183 model = deserialize_keras_object(config_dict, custom_objects)
    184 h5_file = h5py.File(tf.io.gfile.join(temp_path, _VARS_FNAME), "r")
    185 _print_h5_file(h5_file, action="loading")

File /databricks/python/lib/python3.9/site-packages/keras/saving/experimental/serialization_lib.py:318, in deserialize_keras_object(config, custom_objects)
    315 # Instantiate the class from its config inside a custom object scope
    316 # so that we can catch any custom objects that the config refers to.
    317 with object_registration.custom_object_scope(custom_objects):
--> 318     return cls.from_config(inner_config)

File /databricks/python/lib/python3.9/site-packages/keras/engine/training.py:3114, in Model.from_config(cls, config, custom_objects)
   3107 functional_model_keys = [
   3108     "name",
   3109     "layers",
   3110     "input_layers",
   3111     "output_layers",
   3112 ]
   3113 if all(key in config for key in functional_model_keys):
-> 3114     inputs, outputs, layers = functional.reconstruct_from_config(
   3115         config, custom_objects
   3116     )
   3117     model = cls(
   3118         inputs=inputs, outputs=outputs, name=config.get("name")
   3119     )
   3120     functional.connect_ancillary_layers(model, layers)

File /databricks/python/lib/python3.9/site-packages/keras/engine/functional.py:1470, in reconstruct_from_config(config, custom_objects, created_layers)
   1468 # First, we create all layers and enqueue nodes to be processed
   1469 for layer_data in config["layers"]:
-> 1470     process_layer(layer_data)
   1471 # Then we process nodes in order of layer depth.
   1472 # Nodes that cannot yet be processed (if the inbound node
   1473 # does not yet exist) are re-enqueued, and the process
   1474 # is repeated until all nodes are processed.
   1475 while unprocessed_nodes:

File /databricks/python/lib/python3.9/site-packages/keras/engine/functional.py:1451, in reconstruct_from_config.<locals>.process_layer(layer_data)
   1447 else:
   1448     # Instantiate layer.
   1449     from keras.layers import deserialize as deserialize_layer
-> 1451     layer = deserialize_layer(layer_data, custom_objects=custom_objects)
   1452     created_layers[layer_name] = layer
   1454 node_count_by_layer[layer] = int(_should_skip_first_node(layer))

File /databricks/python/lib/python3.9/site-packages/keras/layers/serialization.py:252, in deserialize(config, custom_objects)
    215 """Instantiates a layer from a config dictionary.
    216 
    217 Args:
   (...)
    249 
    250 """
    251 populate_deserializable_objects()
--> 252 return serialization.deserialize_keras_object(
    253     config,
    254     module_objects=LOCAL.ALL_OBJECTS,
    255     custom_objects=custom_objects,
    256     printable_module_name="layer",
    257 )

File /databricks/python/lib/python3.9/site-packages/keras/saving/legacy/serialization.py:527, in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    525     else:
    526         with object_registration.CustomObjectScope(custom_objects):
--> 527             deserialized_obj = cls.from_config(cls_config)
    528 else:
    529     # Then `cls` may be a function returning a class.
    530     # in this case by convention `config` holds
    531     # the kwargs of the function.
    532     custom_objects = custom_objects or {}

File /databricks/python/lib/python3.9/site-packages/keras/engine/base_layer.py:860, in Layer.from_config(cls, config)
    844 @classmethod
    845 def from_config(cls, config):
    846     """Creates a layer from its config.
    847 
    848     This method is the reverse of `get_config`,
   (...)
    858         A layer instance.
    859     """
--> 860     return cls(**config)

File /databricks/python/lib/python3.9/site-packages/keras/layers/preprocessing/string_lookup.py:333, in StringLookup.__init__(self, max_tokens, num_oov_indices, mask_token, oov_token, vocabulary, idf_weights, encoding, invert, output_mode, sparse, pad_to_max_tokens, **kwargs)
    329     del kwargs["dtype"]
    331 self.encoding = encoding
--> 333 super().__init__(
    334     max_tokens=max_tokens,
    335     num_oov_indices=num_oov_indices,
    336     mask_token=mask_token,
    337     oov_token=oov_token,
    338     vocabulary=vocabulary,
    339     vocabulary_dtype=tf.string,
    340     idf_weights=idf_weights,
    341     invert=invert,
    342     output_mode=output_mode,
    343     sparse=sparse,
    344     pad_to_max_tokens=pad_to_max_tokens,
    345     **kwargs
    346 )
    347 base_preprocessing_layer.keras_kpl_gauge.get_cell("StringLookup").set(
    348     True
    349 )

File /databricks/python/lib/python3.9/site-packages/keras/layers/preprocessing/index_lookup.py:323, in IndexLookup.__init__(self, max_tokens, num_oov_indices, mask_token, oov_token, vocabulary_dtype, vocabulary, idf_weights, invert, output_mode, sparse, pad_to_max_tokens, **kwargs)
    320     self.idf_weights_const = self.idf_weights.value()
    322 if vocabulary is not None:
--> 323     self.set_vocabulary(vocabulary, idf_weights)
    324 else:
    325     # When restoring from a keras SavedModel, the loading code will
    326     # expect to find and restore a lookup_table attribute on the layer.
    327     # This table needs to be uninitialized as a StaticHashTable cannot
    328     # be initialized twice.
    329     self.lookup_table = self._uninitialized_lookup_table()

File /databricks/python/lib/python3.9/site-packages/keras/layers/preprocessing/index_lookup.py:510, in IndexLookup.set_vocabulary(self, vocabulary, idf_weights)
    507     idf_weights = np.array(idf_weights)
    509 if vocabulary.size == 0:
--> 510     raise ValueError(
    511         f"Cannot set an empty vocabulary, you passed {vocabulary}."
    512     )
    514 oov_start = self._oov_start_index()
    515 token_start = self._token_start_index()

ValueError: Cannot set an empty vocabulary, you passed [].
```
