Problem with a custom pre-processing layer for serving a NER model

TensorFlow General Discussion
,
1

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.

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:

  • 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

## 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

## 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

## MODEL

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

## 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?

5

@rory_gehman do you know how to solve the problem?