Support for BERT's hub.KerasLayer preprocessors and encoders

[thibaut-d]

System information

• TensorFlow.js version: 3.9.0
• tfjs-react-native version: 0.7.0

Current behavior/state.

I'm trying to make a pilot project to see how we can use BERT in React Native with TensorFlow.js.

The code is public: https://github.com/thibaut-d/Bert-in-React-Native-with-TensorFlow.js

I train the following model:

# Preprocessing
bert_preprocess_model = hub.KerasLayer('https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3')

# Bert encoder
bert_model = hub.KerasLayer('https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2')

# Model
def build_model():

  # Input
  text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')

  # Preprocessing 
  preprocessing_layer = hub.KerasLayer(tfhub_handle_preprocess, name='preprocessing')

  # Encoder
  encoder_inputs = preprocessing_layer(text_input)
  encoder = hub.KerasLayer(tfhub_handle_encoder, trainable=True, name='BERT_encoder')

  # Encoder's output
  outputs = encoder(encoder_inputs)
  net = outputs['pooled_output']
  net = tf.keras.layers.Dropout(0.1)(net)

  # Classifier
  regression = tf.keras.layers.Dense(1, name='regression', activation=None)(net)
  classifier = tf.keras.layers.Dense(1, name='classifier', activation='sigmoid')(net)

  # Final output
  outputs = {'regression': regression, 'classifier': classifier}

  # Return the model
  return tf.keras.Model(text_input, outputs)

I do not include the Python code for training and testing, as it is quite classic.

I save it to different formats:

# Save to Tensorflow SavedModel
model.save("./formats/tf_savedmodel",save_format='tf')

WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 165). These functions will not be directly callable after loading.

# Save to HDF5
model.save('./formats/tf_hdf5/model.h5',save_format='h5')

No warning

Then I try different exports:

# Keras HDF5 --> tfjs_layers_model
!tensorflowjs_converter --input_format keras --output_format tfjs_layers_model ./formats/tf_hdf5/model.h5 ./formats/tfjs_layers_model_from_keras_hdf5

No warning

# Keras HDF5 --> tfjs_graph_model
!tensorflowjs_converter --input_format keras --output_format tfjs_graph_model ./formats/tf_hdf5/model.h5 ./formats/tfjs_graph_model_from_keras_hdf5

I get the error message:

Traceback (most recent call last):
  File "C:\Users\thiba\anaconda3\envs\bert\lib\runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "C:\Users\thiba\anaconda3\envs\bert\Scripts\tensorflowjs_converter.exe\__main__.py", line 7, in <module>
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 813, in pip_main
    main([' '.join(sys.argv[1:])])
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 817, in main
    convert(argv[0].split(' '))
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 803, in convert
    _dispatch_converter(input_format, output_format, args, quantization_dtype_map,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 504, in _dispatch_converter
    dispatch_keras_h5_to_tfjs_graph_model_conversion(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 140, in dispatch_keras_h5_to_tfjs_graph_model_conversion
    model = tf.keras.models.load_model(h5_path, compile=False)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\saving\save.py", line 200, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\saving\hdf5_format.py", line 180, in load_model_from_hdf5
    model = model_config_lib.model_from_config(model_config,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\saving\model_config.py", line 52, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\layers\serialization.py", line 208, in deserialize
    return generic_utils.deserialize_keras_object(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\utils\generic_utils.py", line 674, in deserialize_keras_object
    deserialized_obj = cls.from_config(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\engine\functional.py", line 662, in from_config
    input_tensors, output_tensors, created_layers = reconstruct_from_config(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\engine\functional.py", line 1273, in reconstruct_from_config
    process_layer(layer_data)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\engine\functional.py", line 1255, in process_layer
    layer = deserialize_layer(layer_data, custom_objects=custom_objects)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\layers\serialization.py", line 208, in deserialize
    return generic_utils.deserialize_keras_object(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\utils\generic_utils.py", line 659, in deserialize_keras_object
    (cls, cls_config) = class_and_config_for_serialized_keras_object(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\keras\utils\generic_utils.py", line 556, in class_and_config_for_serialized_keras_object
    raise ValueError(
ValueError: Unknown layer: KerasLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.

I also tried to export form tf savedmodel, despite the warning:

# tf_saved_model --> tfjs_graph_model
!tensorflowjs_converter --input_format tf_saved_model --output_format=tfjs_graph_model ./formats/tf_savedmodel ./formats/tfjs_graph_model_from_tf_saved_model

2021-10-17 21:13:22.786947: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX AVX2
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-10-17 21:13:23.483860: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 1789 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1050, pci bus id: 0000:01:00.0, compute capability: 6.1
Traceback (most recent call last):
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\framework\ops.py", line 3962, in _get_op_def
    return self._op_def_cache[type]
KeyError: 'CaseFoldUTF8'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\saved_model\load.py", line 902, in load_internal
    loader = loader_cls(object_graph_proto, saved_model_proto, export_dir,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\saved_model\load.py", line 137, in __init__
    function_deserialization.load_function_def_library(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\saved_model\function_deserialization.py", line 388, in load_function_def_library
    func_graph = function_def_lib.function_def_to_graph(copy)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\framework\function_def_to_graph.py", line 63, in function_def_to_graph
    graph_def, nested_to_flat_tensor_name = function_def_to_graph_def(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\framework\function_def_to_graph.py", line 228, in function_def_to_graph_def
    op_def = default_graph._get_op_def(node_def.op)  # pylint: disable=protected-access
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\framework\ops.py", line 3966, in _get_op_def
    pywrap_tf_session.TF_GraphGetOpDef(self._c_graph, compat.as_bytes(type),
tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'CaseFoldUTF8' in binary running on IDEAPAD. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\thiba\anaconda3\envs\bert\lib\runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "C:\Users\thiba\anaconda3\envs\bert\Scripts\tensorflowjs_converter.exe\__main__.py", line 7, in <module>
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 813, in pip_main
    main([' '.join(sys.argv[1:])])
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 817, in main
    convert(argv[0].split(' '))
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 803, in convert
    _dispatch_converter(input_format, output_format, args, quantization_dtype_map,
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\converter.py", line 523, in _dispatch_converter
    tf_saved_model_conversion_v2.convert_tf_saved_model(
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\tf_saved_model_conversion_v2.py", line 599, in convert_tf_saved_model
    model = _load_model(saved_model_dir, saved_model_tags)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflowjs\converters\tf_saved_model_conversion_v2.py", line 536, in _load_model
    model = load(saved_model_dir, saved_model_tags)
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\saved_model\load.py", line 864, in load
    result = load_internal(export_dir, tags, options)["root"]
  File "C:\Users\thiba\anaconda3\envs\bert\lib\site-packages\tensorflow\python\saved_model\load.py", line 905, in load_internal
    raise FileNotFoundError(
FileNotFoundError: Op type not registered 'CaseFoldUTF8' in binary running on IDEAPAD. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
 If trying to load on a different device from the computational device, consider using setting the `experimental_io_device` option on tf.saved_model.LoadOptions to the io_device such as '/job:localhost'.

Since exporting to keras HDF5 and converting to tfjs_layers did not raise warning or errors, I tried to run it in React Native.

TensforFlow.js successfully loads, with backend: rn-webgl.

But when I load the model with:

const modelJson = require('./model/model.json');
const m1 = require('./model/group1-shard1of5.bin');
const m2 = require('./model/group1-shard2of5.bin');
const m3 = require('./model/group1-shard3of5.bin');
const m4 = require('./model/group1-shard4of5.bin');
const m5 = require('./model/group1-shard5of5.bin');

 const model = await tf.loadLayersModel(bundleResourceIO(modelJson, [m1,m2,m3,m4,m5]))

I get the following message:

[Error: Unknown layer: KerasLayer. This may be due to one of the following reasons:
1. The layer is defined in Python, in which case it needs to be ported to TensorFlow.js or your JavaScript code.
2. The custom layer is defined in JavaScript, but is not registered properly with tf.serialization.registerClass().]

So ?

My bet is that all these errors are linked to the fact that one or both of the hub.KerasLayer are not supported.

# Preprocessing
bert_preprocess_model = hub.KerasLayer('https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3')

# Bert encoder
bert_model = hub.KerasLayer('https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2')

Is there something I can do to load these layers?

Does another similar model support it?

The current option I see is to rely on HuggingFace transformers library on the Python side and to implement from scratch the preprocessing in JavaScript. Is it the only option?

[rthadur]

Did you get chance to check out this demo https://alexfi.dev/blog/tensorflowjs-bert-train ?

[google-ml-butler]

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

[thibaut-d]

Did you get chance to check out this demo https://alexfi.dev/blog/tensorflowjs-bert-train ?

Yes, this is what I mean by : "The current option I see is to rely on HuggingFace transformers library on the Python side and to implement from scratch the preprocessing in JavaScript".

[gaikwadrahul8]

Hi, @thibaut-d

Apologize for the delayed response and we are re-visiting our older feature requests and checking whether those feature requests implemented or not as of now and it seems like we haven't implemented this feature request yet if I'm not wrong and May I know are you still looking for this feature in TFJs ?

If someone wants to contribute for this feature then you're always welcome and please feel free to do and please refer these links Ref-1, Ref-2 . Thank you!

[thibaut-d]

Hi,

By the end I used PyTorch. I think this feature is welcomed when it comes to NLP.