spark-nlp 4.2.6__py2.py3-none-any.whl → 6.2.1__py2.py3-none-any.whl

sparknlp/annotator/embeddings/mpnet_embeddings.py

@@ -0,0 +1,192 @@
+#  Copyright 2017-2022 John Snow Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""Contains classes for E5Embeddings."""
+
+from sparknlp.common import *
+
+
+class MPNetEmbeddings(AnnotatorModel,
+                           HasEmbeddingsProperties,
+                           HasCaseSensitiveProperties,
+                           HasStorageRef,
+                           HasBatchedAnnotate,
+                           HasMaxSentenceLengthLimit):
+    """Sentence embeddings using MPNet.
+
+    MPNet adopts a novel pre-training method, named masked and permuted language modeling,
+    to inherit the advantages of masked language modeling and permuted language modeling for
+    natural language understanding.
+
+    Note that this annotator is only supported for Spark Versions 3.4 and up.
+
+    Pretrained models can be loaded with :meth:`.pretrained` of the companion
+    object:
+
+    >>> embeddings = MPNetEmbeddings.pretrained() \\
+    ...     .setInputCols(["document"]) \\
+    ...     .setOutputCol("mpnet_embeddings")
+
+
+    The default model is ``"all_mpnet_base_v2"``, if no name is provided.
+
+    For available pretrained models please see the
+    `Models Hub <https://sparknlp.org/models?q=MPNet>`__.
+
+
+    ====================== ======================
+    Input Annotation types Output Annotation type
+    ====================== ======================
+    ``DOCUMENT``            ``SENTENCE_EMBEDDINGS``
+    ====================== ======================
+
+    Parameters
+    ----------
+    batchSize
+        Size of every batch , by default 8
+    dimension
+        Number of embedding dimensions, by default 768
+    caseSensitive
+        Whether to ignore case in tokens for embeddings matching, by default False
+    maxSentenceLength
+        Max sentence length to process, by default 512
+    configProtoBytes
+        ConfigProto from tensorflow, serialized into byte array.
+
+    References
+    ----------
+    `MPNet: Masked and Permuted Pre-training for Language Understanding <https://arxiv.org/pdf/2004.09297>`__
+
+    https://github.com/microsoft/MPNet
+
+    **Paper abstract**
+
+    *BERT adopts masked language modeling (MLM) for pre-training and is one of the most successful pre-training models.
+     Since BERT neglects dependency among predicted tokens, XLNet introduces permuted language modeling (PLM) for
+     pre-training to address this problem. However, XLNet does not leverage the full position information of a sentence
+     and thus suffers from position discrepancy between pre-training and fine-tuning. In this paper, we propose MPNet,
+     a novel pre-training method that inherits the advantages of BERT and XLNet and avoids their limitations. MPNet
+     leverages the dependency among predicted tokens through permuted language modeling (vs. MLM in BERT), and takes
+     auxiliary position information as input to make the model see a full sentence and thus reducing the position
+     discrepancy (vs. PLM in XLNet). We pre-train MPNet on a large-scale dataset (over 160GB text corpora) and fine-tune
+     on a variety of down-streaming tasks (GLUE, SQuAD, etc). Experimental results show that MPNet outperforms MLM and
+     PLM by a large margin, and achieves better results on these tasks compared with previous state-of-the-art
+     pre-trained methods (e.g., BERT, XLNet, RoBERTa) under the same model setting.*
+
+    Examples
+    --------
+    >>> import sparknlp
+    >>> from sparknlp.base import *
+    >>> from sparknlp.annotator import *
+    >>> from pyspark.ml import Pipeline
+    >>> documentAssembler = DocumentAssembler() \\
+    ...     .setInputCol("text") \\
+    ...     .setOutputCol("document")
+    >>> embeddings = MPNetEmbeddings.pretrained() \\
+    ...     .setInputCols(["document"]) \\
+    ...     .setOutputCol("mpnet_embeddings")
+    >>> embeddingsFinisher = EmbeddingsFinisher() \\
+    ...     .setInputCols(["mpnet_embeddings"]) \\
+    ...     .setOutputCols("finished_embeddings") \\
+    ...     .setOutputAsVector(True)
+    >>> pipeline = Pipeline().setStages([
+    ...     documentAssembler,
+    ...     embeddings,
+    ...     embeddingsFinisher
+    ... ])
+    >>> data = spark.createDataFrame([["This is an example sentence", "Each sentence is converted"]]).toDF("text")
+    >>> result = pipeline.fit(data).transform(data)
+    >>> result.selectExpr("explode(finished_embeddings) as result").show(5, 80)
+    +--------------------------------------------------------------------------------+
+    |                                                                          result|
+    +--------------------------------------------------------------------------------+
+    |[[0.022502584, -0.078291744, -0.023030775, -0.0051000593, -0.080340415, 0.039...|
+    |[[0.041702367, 0.0010974605, -0.015534201, 0.07092203, -0.0017729357, 0.04661...|
+    +--------------------------------------------------------------------------------+
+    """
+
+    name = "MPNetEmbeddings"
+
+    inputAnnotatorTypes = [AnnotatorType.DOCUMENT]
+
+    outputAnnotatorType = AnnotatorType.SENTENCE_EMBEDDINGS
+    configProtoBytes = Param(Params._dummy(),
+                             "configProtoBytes",
+                             "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
+                             TypeConverters.toListInt)
+
+
+    def setConfigProtoBytes(self, b):
+        """Sets configProto from tensorflow, serialized into byte array.
+
+        Parameters
+        ----------
+        b : List[int]
+            ConfigProto from tensorflow, serialized into byte array
+        """
+        return self._set(configProtoBytes=b)
+
+    @keyword_only
+    def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.MPNetEmbeddings", java_model=None):
+        super(MPNetEmbeddings, self).__init__(
+            classname=classname,
+            java_model=java_model
+        )
+        self._setDefault(
+            dimension=768,
+            batchSize=8,
+            maxSentenceLength=512,
+            caseSensitive=False,
+        )
+
+    @staticmethod
+    def loadSavedModel(folder, spark_session):
+        """Loads a locally saved model.
+
+        Parameters
+        ----------
+        folder : str
+            Folder of the saved model
+        spark_session : pyspark.sql.SparkSession
+            The current SparkSession
+
+        Returns
+        -------
+        MPNetEmbeddings
+            The restored model
+        """
+        from sparknlp.internal import _MPNetLoader
+        jModel = _MPNetLoader(folder, spark_session._jsparkSession)._java_obj
+        return MPNetEmbeddings(java_model=jModel)
+
+    @staticmethod
+    def pretrained(name="all_mpnet_base_v2", lang="en", remote_loc=None):
+        """Downloads and loads a pretrained model.
+
+        Parameters
+        ----------
+        name : str, optional
+            Name of the pretrained model, by default "all_mpnet_base_v2"
+        lang : str, optional
+            Language of the pretrained model, by default "en"
+        remote_loc : str, optional
+            Optional remote address of the resource, by default None. Will use
+            Spark NLPs repositories otherwise.
+
+        Returns
+        -------
+        MPNetEmbeddings
+            The restored model
+        """
+        from sparknlp.pretrained import ResourceDownloader
+        return ResourceDownloader.downloadModel(MPNetEmbeddings, name, lang, remote_loc)

sparknlp/annotator/embeddings/mxbai_embeddings.py

@@ -0,0 +1,184 @@
+#  Copyright 2017-2022 John Snow Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""Contains classes for MxbaiEmbeddings."""
+
+from sparknlp.common import *
+
+
+class MxbaiEmbeddings(AnnotatorModel,
+					HasEmbeddingsProperties,
+					HasCaseSensitiveProperties,
+					HasStorageRef,
+					HasBatchedAnnotate,
+					HasMaxSentenceLengthLimit):
+	"""Sentence embeddings using Mxbai Embeddings.
+
+
+	Pretrained models can be loaded with :meth:`.pretrained` of the companion
+	object:
+
+	>>> embeddings = MxbaiEmbeddings.pretrained() \\
+	...     .setInputCols(["document"]) \\
+	...     .setOutputCol("Mxbai_embeddings")
+
+
+	The default model is ``"mxbai_large_v1"``, if no name is provided.
+
+	For available pretrained models please see the
+	`Models Hub <https://sparknlp.org/models?q=Mxbai>`__.
+
+
+	====================== ======================
+	Input Annotation types Output Annotation type
+	====================== ======================
+	``DOCUMENT``            ``SENTENCE_EMBEDDINGS``
+	====================== ======================
+
+	Parameters
+	----------
+	batchSize
+		Size of every batch , by default 8
+	dimension
+		Number of embedding dimensions, by default 768
+	caseSensitive
+		Whether to ignore case in tokens for embeddings matching, by default False
+	maxSentenceLength
+		Max sentence length to process, by default 512
+	configProtoBytes
+		ConfigProto from tensorflow, serialized into byte array.
+
+
+
+	Examples
+	--------
+	>>> import sparknlp
+	>>> from sparknlp.base import *
+	>>> from sparknlp.annotator import *
+	>>> from pyspark.ml import Pipeline
+	>>> documentAssembler = DocumentAssembler() \\
+	...     .setInputCol("text") \\
+	...     .setOutputCol("document")
+	>>> embeddings = MxbaiEmbeddings.pretrained() \\
+	...     .setInputCols(["document"]) \\
+	...     .setOutputCol("embeddings")
+	>>> embeddingsFinisher = EmbeddingsFinisher() \\
+	...     .setInputCols("embeddings") \\
+	...     .setOutputCols("finished_embeddings") \\
+	...     .setOutputAsVector(True)
+	>>> pipeline = Pipeline().setStages([
+	...     documentAssembler,
+	...     embeddings,
+	...     embeddingsFinisher
+	... ])
+	>>> data = spark.createDataFrame([["hello world", "hello moon"]]).toDF("text")
+	>>> result = pipeline.fit(data).transform(data)
+	>>> result.selectExpr("explode(finished_embeddings) as result").show(5, 80)
+	+--------------------------------------------------------------------------------+
+	|                                                                          result|
+	+--------------------------------------------------------------------------------+
+	|[0.50387806, 0.5861606, 0.35129607, -0.76046336, -0.32446072, -0.117674336, 0...|
+	|[0.6660665, 0.961762, 0.24854276, -0.1018044, -0.6569202, 0.027635604, 0.1915...|
+	+--------------------------------------------------------------------------------+
+	"""
+
+	name = "MxbaiEmbeddings"
+
+	inputAnnotatorTypes = [AnnotatorType.DOCUMENT]
+
+	outputAnnotatorType = AnnotatorType.SENTENCE_EMBEDDINGS
+	poolingStrategy = Param(Params._dummy(),
+							"poolingStrategy",
+							"Pooling strategy to use for sentence embeddings",
+							TypeConverters.toString)
+
+	def setPoolingStrategy(self, value):
+		"""Pooling strategy to use for sentence embeddings.
+
+		Available pooling strategies for sentence embeddings are:
+		  - `"cls"`: leading `[CLS]` token
+		  - `"cls_avg"`: leading `[CLS]` token + mean of all other tokens
+		  - `"last"`: embeddings of the last token in the sequence
+		  - `"avg"`: mean of all tokens
+		  - `"max"`: max of all embedding features of the entire token sequence
+		  - `"int"`: An integer number, which represents the index of the token to use as the
+			embedding
+
+		Parameters
+		----------
+		value : str
+			Pooling strategy to use for sentence embeddings
+		"""
+
+		valid_strategies = {"cls", "cls_avg", "last", "avg", "max"}
+		if value in valid_strategies or value.isdigit():
+			return self._set(poolingStrategy=value)
+		else:
+			raise ValueError(f"Invalid pooling strategy: {value}. "
+							 f"Valid strategies are: {', '.join(self.valid_strategies)} or an integer.")
+
+	@keyword_only
+	def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.MxbaiEmbeddings", java_model=None):
+		super(MxbaiEmbeddings, self).__init__(
+			classname=classname,
+			java_model=java_model
+		)
+		self._setDefault(
+			dimension=1024,
+			batchSize=8,
+			maxSentenceLength=512,
+			caseSensitive=False,
+			poolingStrategy="cls"
+		)
+
+	@staticmethod
+	def loadSavedModel(folder, spark_session):
+		"""Loads a locally saved model.
+
+		Parameters
+		----------
+		folder : str
+			Folder of the saved model
+		spark_session : pyspark.sql.SparkSession
+			The current SparkSession
+
+		Returns
+		-------
+		MxbaiEmbeddings
+			The restored model
+		"""
+		from sparknlp.internal import _MxbaiEmbeddingsLoader
+		jModel = _MxbaiEmbeddingsLoader(folder, spark_session._jsparkSession)._java_obj
+		return MxbaiEmbeddings(java_model=jModel)
+
+	@staticmethod
+	def pretrained(name="mxbai_large_v1", lang="en", remote_loc=None):
+		"""Downloads and loads a pretrained model.
+
+		Parameters
+		----------
+		name : str, optional
+			Name of the pretrained model, by default "mxbai_large_v1"
+		lang : str, optional
+			Language of the pretrained model, by default "en"
+		remote_loc : str, optional
+			Optional remote address of the resource, by default None. Will use
+			Spark NLPs repositories otherwise.
+
+		Returns
+		-------
+		MxbaiEmbeddings
+			The restored model
+		"""
+		from sparknlp.pretrained import ResourceDownloader
+		return ResourceDownloader.downloadModel(MxbaiEmbeddings, name, lang, remote_loc)

sparknlp/annotator/embeddings/nomic_embeddings.py

@@ -0,0 +1,181 @@
+#  Copyright 2017-2022 John Snow Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""Contains classes for E5Embeddings."""
+
+from sparknlp.common import *
+
+
+class NomicEmbeddings(AnnotatorModel, HasEmbeddingsProperties, HasCaseSensitiveProperties, HasStorageRef,
+                      HasBatchedAnnotate, HasMaxSentenceLengthLimit):
+    """Sentence embeddings using NomicEmbeddings.
+
+    nomic-embed-text-v1 is 8192 context length text encoder that surpasses OpenAI
+    text-embedding-ada-002 and text-embedding-3-small performance on short and long context tasks.
+
+    Pretrained models can be loaded with :meth:`.pretrained` of the companion
+    object:
+
+    >>> embeddings = NomicEmbeddings.pretrained() \\
+    ...     .setInputCols(["document"]) \\
+    ...     .setOutputCol("nomic_embeddings")
+
+
+    The default model is ``"nomic_embed_v1"``, if no name is provided.
+
+    For available pretrained models please see the
+    `Models Hub <https://sparknlp.org/models?q=Nomic>`__.
+
+
+    ====================== ======================
+    Input Annotation types Output Annotation type
+    ====================== ======================
+    ``DOCUMENT``            ``SENTENCE_EMBEDDINGS``
+    ====================== ======================
+
+    Parameters
+    ----------
+    batchSize
+        Size of every batch , by default 8
+    dimension
+        Number of embedding dimensions, by default 768
+    caseSensitive
+        Whether to ignore case in tokens for embeddings matching, by default False
+    maxSentenceLength
+        Max sentence length to process, by default 512
+    configProtoBytes
+        ConfigProto from tensorflow, serialized into byte array.
+
+    References
+    ----------
+    `Text Embeddings by Weakly-Supervised Contrastive Pre-training <https://arxiv.org/pdf/2212.03533>`__
+
+    https://github.com/microsoft/unilm/tree/master/nomic
+
+    **Paper abstract**
+
+    *This technical report describes the training
+    of nomic-embed-text-v1, the first fully reproducible,
+    open-source, open-weights, opendata, 8192 context length
+    English text embedding model that outperforms both OpenAI
+    Ada-002 and OpenAI text-embedding-3-small
+    on short and long-context tasks. We release
+    the training code and model weights under
+    an Apache 2 license. In contrast with other
+    open-source models, we release a training data
+    loader with 235 million curated text pairs that
+    allows for the full replication of nomic-embedtext-v1.
+    You can find code and data to replicate the
+    model at https://github.com/nomicai/contrastors.*
+
+    Examples
+    --------
+    >>> import sparknlp
+    >>> from sparknlp.base import *
+    >>> from sparknlp.annotator import *
+    >>> from pyspark.ml import Pipeline
+    >>> documentAssembler = DocumentAssembler() \\
+    ...     .setInputCol("text") \\
+    ...     .setOutputCol("document")
+    >>> embeddings = NomicEmbeddings.pretrained() \\
+    ...     .setInputCols(["document"]) \\
+    ...     .setOutputCol("nomic_embeddings")
+    >>> embeddingsFinisher = EmbeddingsFinisher() \\
+    ...     .setInputCols(["nomic_embeddings"]) \\
+    ...     .setOutputCols("finished_embeddings") \\
+    ...     .setOutputAsVector(True)
+    >>> pipeline = Pipeline().setStages([
+    ...     documentAssembler,
+    ...     embeddings,
+    ...     embeddingsFinisher
+    ... ])
+    >>> data = spark.createDataFrame([["query: how much protein should a female eat",
+    ... "passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day." + \
+    ... "But, as you can see from this chart, you'll need to increase that if you're expecting or training for a" + \
+    ... "marathon. Check out the chart below to see how much protein you should be eating each day.",
+    ... ]]).toDF("text")
+    >>> result = pipeline.fit(data).transform(data)
+    >>> result.selectExpr("explode(finished_embeddings) as result").show(5, 80)
+    +--------------------------------------------------------------------------------+
+    |                                                                          result|
+    +--------------------------------------------------------------------------------+
+    |[[8.0190285E-4, -0.005974853, -0.072875895, 0.007944068, 0.026059335, -0.0080...|
+    |[[0.050514214, 0.010061974, -0.04340176, -0.020937217, 0.05170225, 0.01157857...|
+    +--------------------------------------------------------------------------------+
+    """
+
+    name = "NomicEmbeddings"
+
+    inputAnnotatorTypes = [AnnotatorType.DOCUMENT]
+
+    outputAnnotatorType = AnnotatorType.SENTENCE_EMBEDDINGS
+    configProtoBytes = Param(Params._dummy(), "configProtoBytes",
+                             "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
+                             TypeConverters.toListInt)
+
+    def setConfigProtoBytes(self, b):
+        """Sets configProto from tensorflow, serialized into byte array.
+
+        Parameters
+        ----------
+        b : List[int]
+            ConfigProto from tensorflow, serialized into byte array
+        """
+        return self._set(configProtoBytes=b)
+
+    @keyword_only
+    def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.NomicEmbeddings", java_model=None):
+        super(NomicEmbeddings, self).__init__(classname=classname, java_model=java_model)
+        self._setDefault(dimension=768, batchSize=8, maxSentenceLength=512, caseSensitive=False, )
+
+    @staticmethod
+    def loadSavedModel(folder, spark_session, use_openvino=False):
+        """Loads a locally saved model.
+
+        Parameters
+        ----------
+        folder : str
+            Folder of the saved model
+        spark_session : pyspark.sql.SparkSession
+            The current SparkSession
+
+        Returns
+        -------
+        NomicEmbeddings
+            The restored model
+        """
+        from sparknlp.internal import _NomicLoader
+        jModel = _NomicLoader(folder, spark_session._jsparkSession, use_openvino)._java_obj
+        return NomicEmbeddings(java_model=jModel)
+
+    @staticmethod
+    def pretrained(name="nomic_embed_v1", lang="en", remote_loc=None):
+        """Downloads and loads a pretrained model.
+
+        Parameters
+        ----------
+        name : str, optional
+            Name of the pretrained model, by default "nomic_embed_v1"
+        lang : str, optional
+            Language of the pretrained model, by default "en"
+        remote_loc : str, optional
+            Optional remote address of the resource, by default None. Will use
+            Spark NLPs repositories otherwise.
+
+        Returns
+        -------
+        NomicEmbeddings
+            The restored model
+        """
+        from sparknlp.pretrained import ResourceDownloader
+        return ResourceDownloader.downloadModel(NomicEmbeddings, name, lang, remote_loc)

sparknlp/annotator/embeddings/roberta_embeddings.py

@@ -21,7 +21,8 @@ class RoBertaEmbeddings(AnnotatorModel,
                         HasCaseSensitiveProperties,
                         HasStorageRef,
                         HasBatchedAnnotate,
-                        HasEngine):
+                        HasEngine,
+                        HasMaxSentenceLengthLimit):
     """Creates word embeddings using RoBERTa.
 
     The RoBERTa model was proposed in `RoBERTa: A Robustly Optimized BERT
@@ -42,10 +43,10 @@ class RoBertaEmbeddings(AnnotatorModel,
 
     The default model is ``"roberta_base"``, if no name is provided. For
     available pretrained models please see the `Models Hub
-    <https://nlp.johnsnowlabs.com/models?task=Embeddings>`__.
+    <https://sparknlp.org/models?task=Embeddings>`__.
 
-    For extended examples of usage, see the `Spark NLP Workshop
-    <https://github.com/JohnSnowLabs/spark-nlp-workshop/blob/master/jupyter/transformers/HuggingFace%20in%20Spark%20NLP%20-%20RoBERTa.ipynb>`__.
+    For extended examples of usage, see the `Examples
+    <https://github.com/JohnSnowLabs/spark-nlp/blob/master/examples/python/transformers/HuggingFace%20in%20Spark%20NLP%20-%20RoBERTa.ipynb>`__.
     To see which models are compatible and how to import them see
     `Import Transformers into Spark NLP 🚀
     <https://github.com/JohnSnowLabs/spark-nlp/discussions/5669>`_.
@@ -151,11 +152,6 @@ class RoBertaEmbeddings(AnnotatorModel,
 
     outputAnnotatorType = AnnotatorType.WORD_EMBEDDINGS
 
-    maxSentenceLength = Param(Params._dummy(),
-                              "maxSentenceLength",
-                              "Max sentence length to process",
-                              typeConverter=TypeConverters.toInt)
-
     configProtoBytes = Param(Params._dummy(),
                              "configProtoBytes",
                              "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
@@ -171,16 +167,6 @@ class RoBertaEmbeddings(AnnotatorModel,
         """
         return self._set(configProtoBytes=b)
 
-    def setMaxSentenceLength(self, value):
-        """Sets max sentence length to process.
-
-        Parameters
-        ----------
-        value : int
-            Max sentence length to process
-        """
-        return self._set(maxSentenceLength=value)
-
     @keyword_only
     def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.RoBertaEmbeddings", java_model=None):
         super(RoBertaEmbeddings, self).__init__(
@@ -195,7 +181,7 @@ class RoBertaEmbeddings(AnnotatorModel,
         )
 
     @staticmethod
-    def loadSavedModel(folder, spark_session):
+    def loadSavedModel(folder, spark_session, use_openvino=False):
         """Loads a locally saved model.
 
         Parameters
@@ -204,6 +190,8 @@ class RoBertaEmbeddings(AnnotatorModel,
             Folder of the saved model
         spark_session : pyspark.sql.SparkSession
             The current SparkSession
+        use_openvino: bool
+            Use OpenVINO backend
 
         Returns
         -------
@@ -211,7 +199,7 @@ class RoBertaEmbeddings(AnnotatorModel,
             The restored model
         """
         from sparknlp.internal import _RoBertaLoader
-        jModel = _RoBertaLoader(folder, spark_session._jsparkSession)._java_obj
+        jModel = _RoBertaLoader(folder, spark_session._jsparkSession, use_openvino)._java_obj
         return RoBertaEmbeddings(java_model=jModel)
 
     @staticmethod

sparknlp/annotator/embeddings/roberta_sentence_embeddings.py

@@ -17,11 +17,12 @@ from sparknlp.common import *
 
 
 class RoBertaSentenceEmbeddings(AnnotatorModel,
-                                HasEmbeddingsProperties,
-                                HasCaseSensitiveProperties,
-                                HasStorageRef,
-                                HasBatchedAnnotate,
-                                HasEngine):
+                                     HasEmbeddingsProperties,
+                                     HasCaseSensitiveProperties,
+                                     HasStorageRef,
+                                     HasBatchedAnnotate,
+                                     HasEngine,
+                                     HasMaxSentenceLengthLimit):
     """Sentence-level embeddings using RoBERTa. The RoBERTa model was proposed in RoBERTa: A Robustly Optimized BERT
     Pretraining Approach  by Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy,
     Mike Lewis, Luke Zettlemoyer, Veselin Stoyanov. It is based on Google's BERT model released in 2018. It builds on
@@ -39,7 +40,7 @@ class RoBertaSentenceEmbeddings(AnnotatorModel,
     The default model is ``"sent_roberta_base"``, if no name is provided.
 
     For available pretrained models please see the
-    `Models Hub <https://nlp.johnsnowlabs.com/models?task=Embeddings>`__.
+    `Models Hub <https://sparknlp.org/models?task=Embeddings>`__.
 
     ====================== =======================
     Input Annotation types Output Annotation type
@@ -119,11 +120,6 @@ class RoBertaSentenceEmbeddings(AnnotatorModel,
 
     outputAnnotatorType = AnnotatorType.SENTENCE_EMBEDDINGS
 
-    maxSentenceLength = Param(Params._dummy(),
-                              "maxSentenceLength",
-                              "Max sentence length to process",
-                              typeConverter=TypeConverters.toInt)
-
     configProtoBytes = Param(Params._dummy(),
                              "configProtoBytes",
                              "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
@@ -139,16 +135,6 @@ class RoBertaSentenceEmbeddings(AnnotatorModel,
         """
         return self._set(configProtoBytes=b)
 
-    def setMaxSentenceLength(self, value):
-        """Sets max sentence length to process.
-
-        Parameters
-        ----------
-        value : int
-            Max sentence length to process
-        """
-        return self._set(maxSentenceLength=value)
-
     @keyword_only
     def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.RoBertaSentenceEmbeddings", java_model=None):
         super(RoBertaSentenceEmbeddings, self).__init__(