Integrating ONNX runtime (ORT) in Spark NLP 5.0.0 🎉

[maziyarpanahi]

Overview

This pull request (PR) aims to enhance the capabilities of Spark NLP by introducing the integration of the ONNX runtime ORT for Java. This integration will enable the users to import Transformers and LLM (Language Model) models in ONNX format into Spark NLP. 🎉

In the upcoming release of Spark NLP 5.0.0, users can work with models in both TensorFlow and ONNX formats. However, our team's default provision of pretrained models will be in ONNX format. This choice is driven by the fact that the ONNX models yield significantly better inference results, ranging from 3x to 5x, even without any optimization or quantization techniques on CPUs.

The integration of ORT in Spark NLP empowers users to further enhance the performance of their models. When exporting models to ONNX, users can leverage built-in features provided by libraries such as onnx-runtime, transformers, optimum, and pytorch. These features include optimization and quantization capabilities, which come ready to use out of the box.

Initial Annotators/Features to support ONNX Runtime

In the realm of Vector Databases, the quest for faster and more efficient Embeddings models has become an imperative pursuit. Models like BERT, DistilBERT, and DeBERTa have revolutionized natural language processing tasks by capturing intricate semantic relationships between words. However, their computational demands and slow inference times pose significant challenges in the game of Vector Databases.

In Vector Databases, the speed at which queries are processed and embeddings are retrieved directly impacts the overall performance and responsiveness of the system. As these databases store vast amounts of vectorized data, such as documents, sentences, or entities, swiftly retrieving relevant embeddings becomes paramount. It enables real-time applications like search engines, recommendation systems, sentiment analysis, and chat/instruct-like products similar to ChatGPT to deliver timely and accurate results, ensuring a seamless user experience.

With that in mind, we have started with the following annotators: BertEmbeddings, DistilBertEmebeddings, and DeBertaEmbeddings. We will identify all the existing models for these annotators on our Models Hub, re-exporting them in ONNX format, and re-uploading them with the same name to have a seamless transition for our community starting Spark NLP 5.0.0. (stuff will just get faster with each release starting Spark NLP 5.0.0 🚀)

Models converted to ONNX

• BERT (for Word Embeddings)
• DistilBERT (for Word Embeddings)
• DeBERTa (for Word Embeddings)
• RoBERTa (for Word Embeddings)

Tested platforms

• Test CPU (normal, optimized, and quantized)
  • Colab
  • EMR
  • Glue 3/4
  • AWS ECS Fargate
  • Databricks
• Test GPU/NVIDIA (normal, optimized, and quantized)
  • Colab
  • EMR
  • Glue 3/4
  • Databricks

[coveralls]

Pull Request Test Coverage Report for Build 5388321371

• 1 of 1 (100.0%) changed or added relevant line in 1 file are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage increased (+0.008%) to 65.887%

---

Totals
Change from base Build 5308728546:	0.008%
Covered Lines:	8647
Relevant Lines:	13124

---

💛 - Coveralls