How to import fine-tune BERT model into Spark NLP

[maziyarpanahi]

Spark NLP offers most of the BERT models available on TF Hub, however, how do you import other BERT pre-trained models/weights or already fine-tuned BERT models to Spark NLP?

The BertEmbeddings annotator has a function called loadSavedModel which is responsible for this. But before that, the format of how to save the BERT model (fine-tuned, checkpoints, etc.) is important to match the same inputs/outputs names in Spark NLP:

Inputs:

bert_inputs = dict(
            input_ids=tf.placeholder(dtype=tf.int32, shape=(None, None), name="input_ids"),
            input_mask=tf.placeholder(dtype=tf.int32, shape=(None, None), name="input_mask"),
            segment_ids=tf.placeholder(dtype=tf.int32, shape=(None, None), name="segment_ids")
        )

Outputs:
Let's extract the sequence_output from BRET model provided by TF Hub (it should be the same in a stand-alone or fine-tuned BERT model):

model_output = bert_module(
            bert_inputs,
            signature="tokens",
            as_dict=True)

sequence_output = tf.identity(model_output["sequence_output"], name="sequence_output")

bert_outputs = dict(
     sequence_output=sequence_output
)

Save model:

tf.saved_model.simple_save(
                    session,
                    export_dir,
                    inputs=bert_inputs,
                    outputs=bert_outputs,
                    legacy_init_op=tf.initializers.tables_initializer(name='init_all_tables')
                )

[maziyarpanahi]

Here is the notebook to demonstrate how to import BERT models (checkpoints) to Spark NLP (BertEmbeddings)