Import a Pytorch fine-tuned BERT model on a downstream task into Spark

[FHaouari]

Is your feature request related to a problem? Please describe.

I fine-tuned a BERT model for tweets classification and I need to load it to classify a stream of tweets.
Does Spark-NLP supports Pytorch BERT models or only tensorflow models? can we load our models fine-tuned outside Spark? Please advise if this is supported by Spark-NLP.

Describe the solution you'd like
Import a fine-tuned Pytorch BERT model and use it to classify tweets/text in real-time.

Additional context
I am using Spark 2.4.7 and PySpark.

[maziyarpanahi]

Hi,
We don't have any means to use Pytorch in JVM environment, unfortunately, unlike TensorFlow, they didn't offer anything official to support Java or any JNI.

If you can convert your Pytorch model to TensorFlow (like if you are using Transformers you should be able to save it as SavedModel) or using ONNX and onnx-tensorflow then you can use it easily in Spark NLP with your TF SavedModel.

This sounds like a very good topic to discuss, sharing ideas on how easily to fine-tune BERT, ELMO, ALBERT, ELECTRA, and XLNET which are available in Spark NLP by using libraries like HuggingFace (pytorch, tensorflow, or their native Trainer) and save it in TF SavedModel at the end.

[FHaouari]

Thanks I appreciate your quick response. Could you please direct me to an example on how I can load a TensorFlow saved model using SPARK NLP and use it for prediction.

[maziyarpanahi]

You are welcome. We just released Spark NLP 3.0.0 on TensorFlow v2.3.1 a week ago. I am preparing some notebooks to show this correctly and made some changes in the 3.0.1 release next week to make it even easier with 1 line.

I will update this discussion first with a simple import of any models on HuggingFace for those Transformers available in Spark NLP and then we can make another notebook as to how to simply fine-tune any of them, save them, and load them in Spark NLP.
(it's should be doable and easy)

[FHaouari]

Thanks a lot for your help.

[pulkitmalhotra1995]

Really interesting, @maziyarpanahi were you able to create an example notebook on fine tuning, saving and finally loading them in spark NLP?

[maziyarpanahi]

Hi @pulkitmalhotra1995
The complete support of some annotators like BERT, RoBERTa, DistilBERT, etc. for TF v2 are coming in the next release. I will make some notebooks for that, but it's basically 2 lines Tokenizer/model from HuggingFace, saved_mode, and then load in Spark NLP. We made it all happen without any modification to the models themselves.
It's WIP here: #2907