NotImplementedError: Can only generate a valid config for `hub.KerasLayer(handle, ...)`that uses a string `handle`

Question

I am trying to convert a pre-trained BERT model from Tensorflow Hub to Keras H5 format, by following Google's provided colab notebooks for BERT fine-tuning. When I finally do:

save_options = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
model_for_export.save(saved_model_path, include_optimizer=False,
                      options=save_options, save_format="h5")

I get the following error message:

NotImplementedError: Can only generate a valid config for `hub.KerasLayer(handle, ...)`that uses a string `handle`.

The problem happens in the pre-processer.

PREPROCESS_MODEL = "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"

bert_preprocess = hub.load(PREPROCESS_MODEL)
tokenizer = hub.KerasLayer(bert_preprocess.tokenize, name='tokenizer')
# some more code
packer = hub.KerasLayer(bert_preprocess.bert_pack_inputs,
                          arguments=dict(seq_length=seq_length),
                          name='packer')
model_inputs = packer(truncated_segments)

It does not support conversion for KerasLayer with handlers other than string. It is expecting something like:

tokenizer = hub.KerasLayer(PREPROCESS_MODEL, name='tokenizer')
# some more code
packer = hub.KerasLayer(PREPROCESS_MODEL,
                          arguments=dict(seq_length=seq_length),
                          name='packer')

But we have to differentiate between the tokenize and the bert_pack_inputs sub-modules while building the pipeline, hence we first load the pre-processer and then use the two different sub-modules to the respective KerasLayer. Is there any way around this?

I have posted the full code as Github Gist for reference.