Preprocessor model can't be saved and loaded if we have only one input_segment

[ianupamsingh]

I was trying to make a text-classification model using BERT and I wanted to use this preprocessor model. To do that I followed this notebook for creating the preprocessor model. In the notebook, we have the function `

def make_bert_preprocess_model(sentence_features, seq_length=128):
  """Returns Model mapping string features to BERT inputs.

  Args:
    sentence_features: a list with the names of string-valued features.
    seq_length: an integer that defines the sequence length of BERT inputs.

  Returns:
    A Keras Model that can be called on a list or dict of string Tensors
    (with the order or names, resp., given by sentence_features) and
    returns a dict of tensors for input to BERT.
  """

  input_segments = [
      tf.keras.layers.Input(shape=(), dtype=tf.string, name=ft)
      for ft in sentence_features]

  # Tokenize the text to word pieces.
  bert_preprocess = hub.load(tfhub_handle_preprocess)
  tokenizer = hub.KerasLayer(bert_preprocess.tokenize, name='tokenizer')
  segments = [tokenizer(s) for s in input_segments]

  packer = hub.KerasLayer(bert_preprocess.bert_pack_inputs,
                          arguments=dict(seq_length=seq_length),
                          name='packer')
  model_inputs = packer(segments)
  return tf.keras.Model(input_segments, model_inputs)

Since I have only a single text to classify I wanted to have just one input segment. Although it's possible to do that by passing sentence_features as ['text1'] but the model hence created cannot be loaded after saving it. See below:

model_two = make_bert_preprocess_model(['text1', 'text2'])
model_one = make_bert_preprocess_model(['text1'])
model_two.save('model_two')
model_one.save('model_one')

Now when I load both models, I can successfully load model_two but not model_one. Following is the traceback

model_two_load = tf.keras.models.load_model('model_two')
model_one_load = tf.keras.models.load_model('model_one')

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-12-f12cb48297bb> in <module>()
      1 model_two_load = tf.keras.models.load_model('model_two')
----> 2 model_one_load = tf.keras.models.load_model('model_one')

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/saved_model/function_deserialization.py in restored_function_body(*args, **kwargs)
    288           .format(index + 1, _pretty_format_positional(positional), keyword))
    289     raise ValueError(
--> 290         "Could not find matching concrete function to call loaded from the "
    291         f"SavedModel. Got:\n  {_pretty_format_positional(args)}\n  Keyword "
    292         f"arguments: {kwargs}\n\n Expected these arguments to match one of the "

ValueError: Exception encountered when calling layer "packer" (type KerasLayer).

Could not find matching concrete function to call loaded from the SavedModel. Got:
  Positional arguments (2 total):
    * tf.RaggedTensor(values=tf.RaggedTensor(values=Tensor("inputs:0", shape=(None,), dtype=int32), row_splits=Tensor("inputs_2:0", shape=(None,), dtype=int64)), row_splits=Tensor("inputs_1:0", shape=(None,), dtype=int64))
    * False
  Keyword arguments: {}

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

Option 1:
  Positional arguments (2 total):
    * [RaggedTensorSpec(TensorShape([None, None, None]), tf.int32, 2, tf.int64)]
    * False
  Keyword arguments: {}

Option 2:
  Positional arguments (2 total):
    * [RaggedTensorSpec(TensorShape([None, None, None]), tf.int32, 2, tf.int64)]
    * True
  Keyword arguments: {}

Call arguments received:
  • args=('tf.RaggedTensor(values=tf.RaggedTensor(values=Tensor("Placeholder:0", shape=(None,), dtype=int32), row_splits=Tensor("Placeholder_1:0", shape=(None,), dtype=int64)), row_splits=Tensor("Placeholder_2:0", shape=(None,), dtype=int64))',)
  • kwargs={'training': 'None'}

From a quick overview, I found out when the length of sentence_features is 1 there is some squeeze operation being done while saving and loading.

I solved the issue in my case by keeping a dummy feature which I always send as empty as follows.

preprocessor_model = make_bert_preprocess_model(['text', 'dummy`])

[gustheman]

Hi Ian, sorry for the delay

I've just tried saving a preprocess model with only one feature and it worked fine.

Maybe I'm missing some part of your test. What I did was just on the same notebook you mentioned, I've added these cells at the bottom just for testing:

save_options = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
model_one = make_bert_preprocess_model(['sentence'])
model_one.save('model_one', include_optimizer=False,
                      options=save_options)

with tf.device('/job:localhost'):
  test_model = tf.saved_model.load('model_one')

with tf.device('/job:localhost'):
  test_dataset = tf.data.Dataset.from_tensor_slices(in_memory_ds[test_split])
  for test_row in test_dataset.shuffle(1000).map(prepare).take(5):
    if len(sentence_features) == 1:
      result = test_model(test_row[0])
      print(result)

also, do you need a customisation of the preprocessing? if not, you don't need to run like that, you can use just the preprocess model from TFHub directly:

pp_model = hub.load(proprocess_model)

[pindinagesh]

Hi @ianupamsingh

Could you please respond to the above @gustheman's comment. Thanks!

[r-remus]

I believe the error that @ianupamsingh describes does not arise while (re-)loading the model, but when actually using the model for inference on a single input (I'm having the same issue currently).

[ianupamsingh]

Hi @gustheman,

I apologize for the late response.

I tried out your code and to my surprise, the loading of the model worked if I used tf.saved_model.load() instead of tf.keras.models.load_model(). But the problem still persists without it. To summarize, following code works fine:

model_one = make_bert_preprocess_model(['text1'])
model_one.save('model_one')
test_model_one = tf.saved_model.load('model_one')

Although, this does not:

model_one = make_bert_preprocess_model(['text1'])
model_one.save('model_one')
test_model_one = tf.keras.models.load_model('model_one')

One more noteworthy observation was that the originally created model model_one expects a single tensor inside a list while the model once saved and loaded into test_model_one expects just a tensor while making inference:

model_one([np.array(['some random sentence'])])  # notice np.array() is inside a list here
test_model_one(np.array(['some random sentece']))  # here it expects just a tensor-like object not within a list

This might have something to do with squeeze I mentioned earlier

[r-remus]

Right. I'm experiencing the same thing. I can't reload the saved Keras model as Keras model using tf.keras.models.load_model().

[ArthurD1]

Same issue here, I'd like to avoid inserting an empty string as a workaround. Does anyone have any better solution?

[WGierke]

I'll close this issue since the workaround provided by https://github.com/tensorflow/hub/issues/839#issuecomment-1055722317 seems to solve the problem. Feel free to re-open it if there are further issues or if you've found a better solution!