How to inference .pb model converted from huggingface?

Viewed 707

In general, I need to run DistilBERT in a browser. At first, I converted DistilBERT from huggingface to TensorFlow .pb format. However, I do not understand how to inference it.

Conversion code:

from transformers import TFAutoModel, AutoTokenizer

model = TFAutoModel.from_pretrained('distilbert-base-uncased')
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
dir = "distilbert_savedmodel"
model._set_inputs(tf.TensorSpec([1, 384], tf.int32))
tf.saved_model.save(model, dir)

Inference code:

encoded = tokenizer.encode('Hello, world!', add_special_tokens=True, return_tensors="tf")
model = tf.keras.models.load_model(dir)
model(encoded)

The error:

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (1 total):
    * Tensor("inputs:0", shape=(1, 1, 6), dtype=int32)
  Keyword arguments: {'training': False}

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

Option 1:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}
  Keyword arguments: {'training': False}

Option 2:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}
  Keyword arguments: {'training': True}

Option 3:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}
  Keyword arguments: {'training': False}

Option 4:
  Positional arguments (1 total):
    * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}
  Keyword arguments: {'training': True}

Notebook link: https://colab.research.google.com/drive/1otfNIYv8DRo2OZ0D2IpdoywrL0pN9k0I?usp=sharing

P. S. I am a novice at TensorFlow.

1 Answers

Looks like transformers are not compatible with saved_model format and produce garbage once loaded.

Related