Why I'm not able to pass a string to my custom layer?

I’m building a NER model with a custom pre-processing layer, where this layer converts the input string into a sequence of ints (each int mapped to a word of the vocabulary), so that I could then deploy the model on a server and call the predict method through APIs.

However I’m having some problems with the input passed to the custom layer:

  • by running predict right after compile, it works only if run_eagerly=True was passed to the compile.
  • if the model is exported and then loaded, the predict method raises an error even if run_eagerly=True was used in the compilation of the model.

Full code below; colab notebook here.

Sample data

import pandas as pd
import numpy as np
import io, os, traceback

data = '''Sentence #,Word,Tag
Sentence: 0,It,O
Sentence: 0,was,O
Sentence: 0,23/10/1991,B-dat
Sentence: 0,when,O
Sentence: 0,the,O
Sentence: 0,life,O
Sentence: 0,of,O
Sentence: 0,James,B-nam
Sentence: 0,Parker,I-nam
Sentence: 0,suddenly,O
Sentence: 0,changed,O
Sentence: 1,George,B-nam
Sentence: 1,is,O
Sentence: 1,going,O
Sentence: 1,to,O
Sentence: 1,buy,O
Sentence: 1,a,O
Sentence: 1,gift,O
Sentence: 1,for,O
Sentence: 1,his,O
Sentence: 1,dad,O
Sentence: 1,Jim,B-nam
Sentence: 1,since,O
Sentence: 1,on,O
Sentence: 1,11/11/2021,B-dat
Sentence: 1,its,O
Sentence: 1,his,O
Sentence: 1,birthday,O
Sentence: 2,There,O
Sentence: 2,is,O
Sentence: 2,no,O
Sentence: 2,evidence,O
Sentence: 2,that,O
Sentence: 2,Alice,B-nam
Sentence: 2,Jackson,I-nam
Sentence: 2,was,O
Sentence: 2,at,O
Sentence: 2,home,O
Sentence: 2,the,O
Sentence: 2,night,O
Sentence: 2,of,O
Sentence: 2,10/10/2010,B-dat'''

data = pd.read_csv(io.StringIO(data), sep=',')

Create vocabulary

padding_value = ""
vocab_token = list(set(data['Word'])) + [padding_value]
idx2token = {idx:tok for idx, tok in enumerate(vocab_token)}
token2idx = {tok:idx for idx, tok in enumerate(vocab_token)}

vocab_tag = list(set(data['Tag']))
idx2tag = {idx:tok for idx, tok in enumerate(vocab_tag)}
tag2idx = {tok:idx for idx, tok in enumerate(vocab_tag)}

data['Word_idx'] = data['Word'].map(token2idx)
data['Tag_idx']  = data['Tag'].map(tag2idx)

n_words = len(vocab_token)
n_tags  = len(vocab_tag)
n_sentences = len(data['Sentence #'].unique())
max_len = max(data['Sentence #'].value_counts())

Padding

from sklearn.model_selection import train_test_split
from keras_preprocessing.sequence import pad_sequences
import tensorflow as tf

data_group = data.groupby(['Sentence #'], as_index=False)[['Word', 'Tag', 'Word_idx', 'Tag_idx']].agg(lambda x: list(x))

tokens = data_group['Word_idx'].tolist()
pad_tokens = pad_sequences(tokens, maxlen=max_len, dtype='int32', padding='post', value=token2idx[padding_value])

tags = data_group['Tag_idx'].tolist()
pad_tags = pad_sequences(tags, maxlen=max_len, dtype='int32', padding='post', value=tag2idx["O"])
pad_tags = [tf.keras.utils.to_categorical(i, num_classes=n_tags) for i in pad_tags]

Model definition and fit

import tensorflow as tf
from tensorflow.keras import Sequential, Input
from tensorflow.keras.layers import LSTM, Embedding, Dense, TimeDistributed, Bidirectional

output_dim = 32

model = Sequential()
model.add(Embedding(input_dim=n_words + 1, output_dim=output_dim, input_length=max_len))
model.add(Bidirectional(LSTM(units=output_dim, return_sequences=True, dropout=0.2, recurrent_dropout=0.2), merge_mode = 'concat'))
model.add(TimeDistributed(Dense(n_tags, activation="softmax")))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

model.fit(pad_tokens, np.array(pad_tags), batch_size=16, epochs=1)

Pre-processing layer

@tf.keras.utils.register_keras_serializable(name='pre_processing_layer')
class InputProcessingLayer(tf.keras.layers.Layer):
    def __init__(self, vocab, **kwargs):
        super(InputProcessingLayer, self).__init__(**kwargs)
        self.vocab = vocab

    def call(self, input_str, training=False):
        # a print useful for debugging
        print('-->', input_str, type(input_str))
        # from tf.Tensor([b'...']) to tf.Tensor(b'...')
        input_str = input_str[0]
        try:
            # lowers the input string and splits it on whitespaces
            tokens = self.vocab(tf.strings.split(tf.strings.lower(input_str)))
            left = 0
            right = max(0, max_len - len(tokens))
            paddings = tf.constant([[left, right]])
            # pad the list so that its length will be equal to max_len
            padded_list = tf.pad(tokens, paddings, mode="CONSTANT", constant_values=token2idx[padding_value])
            input_words = tf.reshape(padded_list, [1,max_len])
        except:
            print('\n- - - - ERROR MSG - - - -\n')
            traceback.print_exc()
            print('\n- - - - END MSG - - - -\n')
            padded_list = [0]*max_len
            # a dummy list but necessary because the output's shape of InputProcessingLayer must be compatible with the embedding layer's shape of the model 
            input_words = np.array(padded_list, dtype='int32').reshape(1,max_len)
        return input_words
    
    def get_config(self):
        config = super().get_config()
        config['vocab'] = self.vocab
        return config

Add layer to the model

vocab = tf.keras.layers.StringLookup(vocabulary=[x for x in token2idx.keys()])

temp_model = Sequential([
    Input(shape=(max_len,), dtype=tf.string, name='description'),
    InputProcessingLayer(vocab),
    model
])
temp_model.compile('adam', loss=None, run_eagerly=False)

Test the final model

sample_input = "James birthday is on 11/11/2021"
pred = temp_model.predict([sample_input])
print(pred)

With run_eagerly=False I get (notice that 17 is the max_len)

WARNING:tensorflow:Model was constructed with shape (None, 17) for input KerasTensor(type_spec=TensorSpec(shape=(None, 17), dtype=tf.string, name='description'), name='description', description="created by layer 'description'"), but it was called on an input with incompatible shape (None,).
--> Tensor("IteratorGetNext:0", shape=(None,), dtype=string) <class 'tensorflow.python.framework.ops.Tensor'>

Notice that what follows the --> is a print from inside the InputProcessingLayer, and in this case it says that the string received by the layer is "IteratorGetNext:0" instead of "James birthday is on 11/11/2021".

By using run_eagerly=True then string received by the layer is correct

--> tf.Tensor([b'James birthday is on 11/11/2021'], shape=(1,), dtype=string) <class 'tensorflow.python.framework.ops.EagerTensor'>

However, if I save the model, i.e. temp_model.save(EXPORT_PATH, overwrite=True), and then load it for a prediction, then it doesn’t work

import tensorflow as tf

loaded_model = tf.keras.models.load_model(EXPORT_PATH)

sample_input = "James birthday is on 11/11/2021"
pred = loaded_model.predict([sample_input])
print(pred)

raises an error

WARNING:tensorflow:Model was constructed with shape (None, 17) for input KerasTensor(type_spec=TensorSpec(shape=(None, 17), dtype=tf.string, name='description'), name='description', description="created by layer 'description'"), but it was called on an input with incompatible shape (None,).
Traceback (most recent call last):
  File "C:\Users\gtu\Desktop\test_signature_inp.py", line 12, in <module>
    pred = loaded_model.predict([sample_input])
  File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1147, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\engine\training.py", line 1801, in predict_function  *
        return step_function(self, iterator)
    File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\engine\training.py", line 1790, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\engine\training.py", line 1783, in run_step  **
        outputs = model.predict_step(data)
    File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\engine\training.py", line 1751, in predict_step
        return self(x, training=False)
    File "C:\Users\gtu\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
        raise e.with_traceback(filtered_tb) from None

    ValueError: Exception encountered when calling layer "input_processing_layer" (type Custom>pre_processing_layer).

    Could not find matching concrete function to call loaded from the SavedModel. Got:
      Positional arguments (2 total):
        * Tensor("input_str:0", shape=(None,), dtype=string)
        * False
      Keyword arguments: {}

     Expected these arguments to match one of the following 2 option(s):

    Option 1:
      Positional arguments (2 total):
        * TensorSpec(shape=(None, 17), dtype=tf.string, name='input_str')
        * False
      Keyword arguments: {}

    Option 2:
      Positional arguments (2 total):
        * TensorSpec(shape=(None, 17), dtype=tf.string, name='input_str')
        * True
      Keyword arguments: {}

    Call arguments received:
      • args=('tf.Tensor(shape=(None,), dtype=string)',)
      • kwargs={'training': 'False'}

From the error we see that the string received by the layer is "input_str:0", instead of the one passed to the predict method. What’s going on here?

The issues with passing a string to your custom layer in TensorFlow and encountering errors during model prediction, especially after serialization and deserialization, can be complex. Here’s a condensed version of the potential problems and solutions:

Problems:

  1. Custom Layer Input Handling: Your layer’s handling of string inputs might not be compatible with TensorFlow’s expectations, especially after model loading.
  2. Eager vs. Graph Execution: The layer works with run_eagerly=True due to TensorFlow’s eager execution mode, but faces issues in the default graph execution mode.
  3. Serialization and Deserialization: Custom layers can have issues when the model is saved and loaded, particularly if TensorFlow cannot correctly serialize and deserialize the layer’s behavior.

Solutions:

  1. Adjust Custom Layer: Make sure your custom layer properly implements get_config for serialization. Consider preprocessing inputs outside the model.
  2. Use tf.function: Wrap complex input processing in your layer with tf.function to help TensorFlow optimize the execution.
  3. Explicit Input Shapes and Types: Ensure the model’s input layer clearly defines expected input shapes and types.
  4. Model Architecture Review: If persistent, consider redesigning the model architecture or moving preprocessing outside the TensorFlow model.
  5. Consult Documentation: Check TensorFlow’s documentation and community resources for guidance on handling custom layers and model serialization.

These steps can help address compatibility issues with custom layers in TensorFlow models, especially concerning serialization and prediction.