BERT TF-Lite conversion not working in TensorFlow 2.2.0

[r4ghu]

🐛 Bug

Model conversion succeeds but inputs and outputs are not recognized.

Information

Model I am using (Bert, XLNet ...): BERT (bert-base-uncased)

Language I am using the model on (English, Chinese ...): English

The problem arises when using:

• TensorFlow 2.1.0 + Transformers 2.8.0 - has no problem converting bert-base-uncased model to tflite version.
• TensorFlow 2.2.0rc3 + Transformers 2.8.0 - has issues with interoperability.

The tasks I am working on is:

• Convert BERT models to TF-Lite format to use it in mobile apps.
• Trying to use the latest TF-Lite package version for Android in the place of TF-Lite package provided in the repo huggingface/tflite-android-transformers.

To reproduce

Please execute the following code with TensorFlow versions 2.1.0 and 2.2.0-rc3

import transformers
from transformers import TFBertModel, BertConfig
import tensorflow as tf
print('TensorFlow version =', tf.__version__)
print('Transformers version =', transformers.__version__)

MODEL_DIR = 'bert-base-uncased'
MAX_SEQ_LEN = 50

# Read the model
config = BertConfig.from_pretrained(MODEL_DIR)
model = TFBertModel(config)


# Set input Spec
input_spec = [
    tf.TensorSpec([1, MAX_SEQ_LEN], tf.int32),
    tf.TensorSpec([1, MAX_SEQ_LEN], tf.int32),
    tf.TensorSpec([1, MAX_SEQ_LEN], tf.int32)
]

model._set_inputs(input_spec, training=False)

print(model.inputs)
print(model.outputs)

• For TensorFlow 2.2.0-rc3: Model outputs and inputs are None

TensorFlow version = 2.2.0-rc3
Transformers version = 2.8.0
None
None

• For TensorFlow 2.1.0:

TensorFlow version = 2.1.0
Transformers version = 2.8.0
...
[<tf.Tensor 'input_1:0' shape=(None, 50) dtype=int32>, <tf.Tensor 'input_2:0' shape=(None, 50) dtype=int32>, <tf.Tensor 'input_3:0' shape=(None, 50) dtype=int32>]
[<tf.Tensor 'tf_bert_model/Identity:0' shape=(None, 50, 768) dtype=float32>, <tf.Tensor 'tf_bert_model/Identity_1:0' shape=(None, 768) dtype=float32>]

Expected behavior

• I expect the BERT model conversion to work properly to TensorFlow 2.2.0-rc{1/2/3}
• Preferably BERT should use the default TF-Lite supported layers just like MobileBERT model provided by Google.
• Image - MobileBERT from Google's bert-qa android example (left) vs BERT converted using the above script using TensorFlow v2.1.0 (right)
  

Environment info

• transformers version: 2.8.0
• Platform: Windows
• Python version: 3.7.7
• PyTorch version (GPU?): 1.4.0 (No GPU)
• Tensorflow version (GPU?): 2.1.0 (working), 2.2.0-rc3 (not working) (no GPU for both versions)
• Using GPU in script?: No
• Using distributed or parallel set-up in script?: No

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.