spark-nlp 5.4.2__py2.py3-none-any.whl → 5.5.0__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of spark-nlp might be problematic. Click here for more details.

@@ -0,0 +1,211 @@
1
+ # Copyright 2017-2024 John Snow Labs
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """Contains classes for AlbertForZeroShotClassification."""
16
+
17
+ from sparknlp.common import *
18
+
19
+
20
+ class AlbertForZeroShotClassification(AnnotatorModel,
21
+ HasCaseSensitiveProperties,
22
+ HasBatchedAnnotate,
23
+ HasClassifierActivationProperties,
24
+ HasCandidateLabelsProperties,
25
+ HasEngine,
26
+ HasMaxSentenceLengthLimit):
27
+ """AlbertForZeroShotClassification using a `ModelForSequenceClassification` trained on NLI (natural language
28
+ inference) tasks. Equivalent of `DistilBertForSequenceClassification` models, but these models don't require a hardcoded
29
+ number of potential classes, they can be chosen at runtime. It usually means it's slower but it is much more
30
+ flexible.
31
+
32
+ Note that the model will loop through all provided labels. So the more labels you have, the
33
+ longer this process will take.
34
+
35
+ Any combination of sequences and labels can be passed and each combination will be posed as a premise/hypothesis
36
+ pair and passed to the pretrained model.
37
+
38
+ Pretrained models can be loaded with :meth:`.pretrained` of the companion
39
+ object:
40
+
41
+ >>> sequenceClassifier = AlbertForZeroShotClassification.pretrained() \\
42
+ ... .setInputCols(["token", "document"]) \\
43
+ ... .setOutputCol("label")
44
+
45
+ The default model is ``"albert_base_zero_shot_classifier_onnx"``, if no name is
46
+ provided.
47
+
48
+ For available pretrained models please see the `Models Hub
49
+ <https://sparknlp.orgtask=Text+Classification>`__.
50
+
51
+ To see which models are compatible and how to import them see
52
+ `Import Transformers into Spark NLP 🚀
53
+ <https://github.com/JohnSnowLabs/spark-nlp/discussions/5669>`_.
54
+
55
+ ====================== ======================
56
+ Input Annotation types Output Annotation type
57
+ ====================== ======================
58
+ ``DOCUMENT, TOKEN`` ``CATEGORY``
59
+ ====================== ======================
60
+
61
+ Parameters
62
+ ----------
63
+ batchSize
64
+ Batch size. Large values allows faster processing but requires more
65
+ memory, by default 8
66
+ caseSensitive
67
+ Whether to ignore case in tokens for embeddings matching, by default
68
+ True
69
+ configProtoBytes
70
+ ConfigProto from tensorflow, serialized into byte array.
71
+ maxSentenceLength
72
+ Max sentence length to process, by default 128
73
+ coalesceSentences
74
+ Instead of 1 class per sentence (if inputCols is `sentence`) output 1
75
+ class per document by averaging probabilities in all sentences, by
76
+ default False
77
+ activation
78
+ Whether to calculate logits via Softmax or Sigmoid, by default
79
+ `"softmax"`.
80
+
81
+ Examples
82
+ --------
83
+ >>> import sparknlp
84
+ >>> from sparknlp.base import *
85
+ >>> from sparknlp.annotator import *
86
+ >>> from pyspark.ml import Pipeline
87
+ >>> documentAssembler = DocumentAssembler() \\
88
+ ... .setInputCol("text") \\
89
+ ... .setOutputCol("document")
90
+ >>> tokenizer = Tokenizer() \\
91
+ ... .setInputCols(["document"]) \\
92
+ ... .setOutputCol("token")
93
+ >>> sequenceClassifier = AlbertForZeroShotClassification.pretrained() \\
94
+ ... .setInputCols(["token", "document"]) \\
95
+ ... .setOutputCol("label") \\
96
+ ... .setCaseSensitive(True)
97
+ >>> pipeline = Pipeline().setStages([
98
+ ... documentAssembler,
99
+ ... tokenizer,
100
+ ... sequenceClassifier
101
+ ... ])
102
+ >>> data = spark.createDataFrame([["I have a problem with my iphone that needs to be resolved asap!!"]]).toDF("text")
103
+ >>> result = pipeline.fit(data).transform(data)
104
+ >>> result.select("label.result").show(truncate=False)
105
+ +---------+
106
+ |result |
107
+ +---------+
108
+ |[urgent] |
109
+ +---------+
110
+ """
111
+ name = "AlbertForZeroShotClassification"
112
+
113
+ inputAnnotatorTypes = [AnnotatorType.DOCUMENT, AnnotatorType.TOKEN]
114
+
115
+ outputAnnotatorType = AnnotatorType.CATEGORY
116
+
117
+ configProtoBytes = Param(Params._dummy(),
118
+ "configProtoBytes",
119
+ "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
120
+ TypeConverters.toListInt)
121
+
122
+ coalesceSentences = Param(Params._dummy(), "coalesceSentences",
123
+ "Instead of 1 class per sentence (if inputCols is '''sentence''') output 1 class per document by averaging probabilities in all sentences.",
124
+ TypeConverters.toBoolean)
125
+
126
+ def getClasses(self):
127
+ """
128
+ Returns labels used to train this model
129
+ """
130
+ return self._call_java("getClasses")
131
+
132
+ def setConfigProtoBytes(self, b):
133
+ """Sets configProto from tensorflow, serialized into byte array.
134
+
135
+ Parameters
136
+ ----------
137
+ b : List[int]
138
+ ConfigProto from tensorflow, serialized into byte array
139
+ """
140
+ return self._set(configProtoBytes=b)
141
+
142
+ def setCoalesceSentences(self, value):
143
+ """Instead of 1 class per sentence (if inputCols is '''sentence''') output 1 class per document by averaging
144
+ probabilities in all sentences. Due to max sequence length limit in almost all transformer models such as Bart
145
+ (512 tokens), this parameter helps to feed all the sentences into the model and averaging all the probabilities
146
+ for the entire document instead of probabilities per sentence. (Default: true)
147
+
148
+ Parameters
149
+ ----------
150
+ value : bool
151
+ If the output of all sentences will be averaged to one output
152
+ """
153
+ return self._set(coalesceSentences=value)
154
+
155
+ @keyword_only
156
+ def __init__(self, classname="com.johnsnowlabs.nlp.annotators.classifier.dl.AlbertForZeroShotClassification",
157
+ java_model=None):
158
+ super(AlbertForZeroShotClassification, self).__init__(
159
+ classname=classname,
160
+ java_model=java_model
161
+ )
162
+ self._setDefault(
163
+ batchSize=8,
164
+ maxSentenceLength=128,
165
+ caseSensitive=True,
166
+ coalesceSentences=False,
167
+ activation="softmax"
168
+ )
169
+
170
+ @staticmethod
171
+ def loadSavedModel(folder, spark_session):
172
+ """Loads a locally saved model.
173
+
174
+ Parameters
175
+ ----------
176
+ folder : str
177
+ Folder of the saved model
178
+ spark_session : pyspark.sql.SparkSession
179
+ The current SparkSession
180
+
181
+ Returns
182
+ -------
183
+ AlbertForZeroShotClassification
184
+ The restored model
185
+ """
186
+ from sparknlp.internal import _AlbertForZeroShotClassificationLoader
187
+ jModel = _AlbertForZeroShotClassificationLoader(folder, spark_session._jsparkSession)._java_obj
188
+ return AlbertForZeroShotClassification(java_model=jModel)
189
+
190
+ @staticmethod
191
+ def pretrained(name="albert_zero_shot_classifier_onnx", lang="en", remote_loc=None):
192
+ """Downloads and loads a pretrained model.
193
+
194
+ Parameters
195
+ ----------
196
+ name : str, optional
197
+ Name of the pretrained model, by default
198
+ "albert_zero_shot_classifier_onnx"
199
+ lang : str, optional
200
+ Language of the pretrained model, by default "en"
201
+ remote_loc : str, optional
202
+ Optional remote address of the resource, by default None. Will use
203
+ Spark NLPs repositories otherwise.
204
+
205
+ Returns
206
+ -------
207
+ BartForZeroShotClassification
208
+ The restored model
209
+ """
210
+ from sparknlp.pretrained import ResourceDownloader
211
+ return ResourceDownloader.downloadModel(AlbertForZeroShotClassification, name, lang, remote_loc)
@@ -0,0 +1,202 @@
1
+ # Copyright 2017-2024 John Snow Labs
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """Contains classes for CamemBertForSequenceClassification."""
15
+
16
+ from sparknlp.common import *
17
+
18
+
19
+ class CamemBertForZeroShotClassification(AnnotatorModel,
20
+ HasCaseSensitiveProperties,
21
+ HasBatchedAnnotate,
22
+ HasClassifierActivationProperties,
23
+ HasCandidateLabelsProperties,
24
+ HasEngine,
25
+ HasMaxSentenceLengthLimit):
26
+ """CamemBertForZeroShotClassification using a `ModelForSequenceClassification` trained on NLI (natural language
27
+ inference) tasks. Equivalent of `DeBertaForSequenceClassification` models, but these models don't require a hardcoded
28
+ number of potential classes, they can be chosen at runtime. It usually means it's slower but it is much more
29
+ flexible.
30
+ Any combination of sequences and labels can be passed and each combination will be posed as a premise/hypothesis
31
+ pair and passed to the pretrained model.
32
+ Pretrained models can be loaded with :meth:`.pretrained` of the companion
33
+ object:
34
+ >>> sequenceClassifier = CamemBertForZeroShotClassification.pretrained() \\
35
+ ... .setInputCols(["token", "document"]) \\
36
+ ... .setOutputCol("label")
37
+ The default model is ``"camembert_zero_shot_classifier_xnli_onnx"``, if no name is
38
+ provided.
39
+ For available pretrained models please see the `Models Hub
40
+ <https://sparknlp.orgtask=Text+Classification>`__.
41
+ To see which models are compatible and how to import them see
42
+ `Import Transformers into Spark NLP 🚀
43
+ <https://github.com/JohnSnowLabs/spark-nlp/discussions/5669>`_.
44
+ ====================== ======================
45
+ Input Annotation types Output Annotation type
46
+ ====================== ======================
47
+ ``DOCUMENT, TOKEN`` ``CATEGORY``
48
+ ====================== ======================
49
+ Parameters
50
+ ----------
51
+ batchSize
52
+ Batch size. Large values allows faster processing but requires more
53
+ memory, by default 8
54
+ caseSensitive
55
+ Whether to ignore case in tokens for embeddings matching, by default
56
+ True
57
+ configProtoBytes
58
+ ConfigProto from tensorflow, serialized into byte array.
59
+ maxSentenceLength
60
+ Max sentence length to process, by default 128
61
+ coalesceSentences
62
+ Instead of 1 class per sentence (if inputCols is `sentence`) output 1
63
+ class per document by averaging probabilities in all sentences, by
64
+ default False
65
+ activation
66
+ Whether to calculate logits via Softmax or Sigmoid, by default
67
+ `"softmax"`.
68
+ Examples
69
+ --------
70
+ >>> import sparknlp
71
+ >>> from sparknlp.base import *
72
+ >>> from sparknlp.annotator import *
73
+ >>> from pyspark.ml import Pipeline
74
+ >>> documentAssembler = DocumentAssembler() \\
75
+ ... .setInputCol("text") \\
76
+ ... .setOutputCol("document")
77
+ >>> tokenizer = Tokenizer() \\
78
+ ... .setInputCols(["document"]) \\
79
+ ... .setOutputCol("token")
80
+ >>> sequenceClassifier = CamemBertForZeroShotClassification.pretrained() \\
81
+ ... .setInputCols(["token", "document"]) \\
82
+ ... .setOutputCol("multi_class") \\
83
+ ... .setCaseSensitive(True)
84
+ ... .setCandidateLabels(["sport", "politique", "science"])
85
+ >>> pipeline = Pipeline().setStages([
86
+ ... documentAssembler,
87
+ ... tokenizer,
88
+ ... sequenceClassifier
89
+ ... ])
90
+ >>> data = spark.createDataFrame([["L'équipe de France joue aujourd'hui au Parc des Princes"]]).toDF("text")
91
+ >>> result = pipeline.fit(data).transform(data)
92
+ >>> result.select("class.result").show(truncate=False)
93
+ +------+
94
+ |result|
95
+ +------+
96
+ |[sport]|
97
+ +------+
98
+ """
99
+ name = "CamemBertForZeroShotClassification"
100
+
101
+ inputAnnotatorTypes = [AnnotatorType.DOCUMENT, AnnotatorType.TOKEN]
102
+
103
+ outputAnnotatorType = AnnotatorType.CATEGORY
104
+
105
+ configProtoBytes = Param(Params._dummy(),
106
+ "configProtoBytes",
107
+ "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
108
+ TypeConverters.toListInt)
109
+
110
+ coalesceSentences = Param(Params._dummy(), "coalesceSentences",
111
+ "Instead of 1 class per sentence (if inputCols is '''sentence''') output 1 class per document by averaging probabilities in all sentences.",
112
+ TypeConverters.toBoolean)
113
+
114
+ def getClasses(self):
115
+ """
116
+ Returns labels used to train this model
117
+ """
118
+ return self._call_java("getClasses")
119
+
120
+ def setConfigProtoBytes(self, b):
121
+ """Sets configProto from tensorflow, serialized into byte array.
122
+
123
+ Parameters
124
+ ----------
125
+ b : List[int]
126
+ ConfigProto from tensorflow, serialized into byte array
127
+ """
128
+ return self._set(configProtoBytes=b)
129
+
130
+ def setCoalesceSentences(self, value):
131
+ """Instead of 1 class per sentence (if inputCols is '''sentence''') output 1
132
+ class per document by averaging probabilities in all sentences, by default True.
133
+
134
+ Due to max sequence length limit in almost all transformer models such as BERT
135
+ (512 tokens), this parameter helps feeding all the sentences into the model and
136
+ averaging all the probabilities for the entire document instead of probabilities
137
+ per sentence.
138
+
139
+ Parameters
140
+ ----------
141
+ value : bool
142
+ If the output of all sentences will be averaged to one output
143
+ """
144
+ return self._set(coalesceSentences=value)
145
+
146
+ @keyword_only
147
+ def __init__(self, classname="com.johnsnowlabs.nlp.annotators.classifier.dl.CamemBertForZeroShotClassification",
148
+ java_model=None):
149
+ super(CamemBertForZeroShotClassification, self).__init__(
150
+ classname=classname,
151
+ java_model=java_model
152
+ )
153
+ self._setDefault(
154
+ batchSize=8,
155
+ maxSentenceLength=128,
156
+ caseSensitive=True,
157
+ coalesceSentences=False,
158
+ activation="softmax"
159
+ )
160
+
161
+ @staticmethod
162
+ def loadSavedModel(folder, spark_session):
163
+ """Loads a locally saved model.
164
+
165
+ Parameters
166
+ ----------
167
+ folder : str
168
+ Folder of the saved model
169
+ spark_session : pyspark.sql.SparkSession
170
+ The current SparkSession
171
+
172
+ Returns
173
+ -------
174
+ CamemBertForZeroShotClassification
175
+ The restored model
176
+ """
177
+ from sparknlp.internal import _CamemBertForZeroShotClassificationLoader
178
+ jModel = _CamemBertForZeroShotClassificationLoader(folder, spark_session._jsparkSession)._java_obj
179
+ return CamemBertForZeroShotClassification(java_model=jModel)
180
+
181
+ @staticmethod
182
+ def pretrained(name="camembert_zero_shot_classifier_xnli_onnx", lang="fr", remote_loc=None):
183
+ """Downloads and loads a pretrained model.
184
+
185
+ Parameters
186
+ ----------
187
+ name : str, optional
188
+ Name of the pretrained model, by default
189
+ "camembert_zero_shot_classifier_xnli_onnx"
190
+ lang : str, optional
191
+ Language of the pretrained model, by default "fr"
192
+ remote_loc : str, optional
193
+ Optional remote address of the resource, by default None. Will use
194
+ Spark NLPs repositories otherwise.
195
+
196
+ Returns
197
+ -------
198
+ CamemBertForSequenceClassification
199
+ The restored model
200
+ """
201
+ from sparknlp.pretrained import ResourceDownloader
202
+ return ResourceDownloader.downloadModel(CamemBertForZeroShotClassification, name, lang, remote_loc)
@@ -21,7 +21,8 @@ class DeBertaForZeroShotClassification(AnnotatorModel,
21
21
  HasBatchedAnnotate,
22
22
  HasClassifierActivationProperties,
23
23
  HasCandidateLabelsProperties,
24
- HasEngine):
24
+ HasEngine,
25
+ HasMaxSentenceLengthLimit):
25
26
  """DeBertaForZeroShotClassification using a `ModelForSequenceClassification` trained on NLI (natural language
26
27
  inference) tasks. Equivalent of `DeBertaForSequenceClassification` models, but these models don't require a hardcoded
27
28
  number of potential classes, they can be chosen at runtime. It usually means it's slower but it is much more
@@ -101,11 +102,6 @@ class DeBertaForZeroShotClassification(AnnotatorModel,
101
102
 
102
103
  outputAnnotatorType = AnnotatorType.CATEGORY
103
104
 
104
- maxSentenceLength = Param(Params._dummy(),
105
- "maxSentenceLength",
106
- "Max sentence length to process",
107
- typeConverter=TypeConverters.toInt)
108
-
109
105
  configProtoBytes = Param(Params._dummy(),
110
106
  "configProtoBytes",
111
107
  "ConfigProto from tensorflow, serialized into byte array. Get with config_proto.SerializeToString()",
@@ -130,15 +126,6 @@ class DeBertaForZeroShotClassification(AnnotatorModel,
130
126
  """
131
127
  return self._set(configProtoBytes=b)
132
128
 
133
- def setMaxSentenceLength(self, value):
134
- """Sets max sentence length to process, by default 128.
135
- Parameters
136
- ----------
137
- value : int
138
- Max sentence length to process
139
- """
140
- return self._set(maxSentenceLength=value)
141
-
142
129
  def setCoalesceSentences(self, value):
143
130
  """Instead of 1 class per sentence (if inputCols is '''sentence''') output 1 class per document by averaging
144
131
  probabilities in all sentences. Due to max sequence length limit in almost all transformer models such as DeBerta
@@ -37,3 +37,6 @@ from sparknlp.annotator.embeddings.xlm_roberta_sentence_embeddings import *
37
37
  from sparknlp.annotator.embeddings.xlnet_embeddings import *
38
38
  from sparknlp.annotator.embeddings.bge_embeddings import *
39
39
  from sparknlp.annotator.embeddings.uae_embeddings import *
40
+ from sparknlp.annotator.embeddings.mxbai_embeddings import *
41
+ from sparknlp.annotator.embeddings.snowflake_embeddings import *
42
+ from sparknlp.annotator.embeddings.nomic_embeddings import *
@@ -0,0 +1,184 @@
1
+ # Copyright 2017-2022 John Snow Labs
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """Contains classes for MxbaiEmbeddings."""
15
+
16
+ from sparknlp.common import *
17
+
18
+
19
+ class MxbaiEmbeddings(AnnotatorModel,
20
+ HasEmbeddingsProperties,
21
+ HasCaseSensitiveProperties,
22
+ HasStorageRef,
23
+ HasBatchedAnnotate,
24
+ HasMaxSentenceLengthLimit):
25
+ """Sentence embeddings using Mxbai Embeddings.
26
+
27
+
28
+ Pretrained models can be loaded with :meth:`.pretrained` of the companion
29
+ object:
30
+
31
+ >>> embeddings = MxbaiEmbeddings.pretrained() \\
32
+ ... .setInputCols(["document"]) \\
33
+ ... .setOutputCol("Mxbai_embeddings")
34
+
35
+
36
+ The default model is ``"mxbai_large_v1"``, if no name is provided.
37
+
38
+ For available pretrained models please see the
39
+ `Models Hub <https://sparknlp.org/models?q=Mxbai>`__.
40
+
41
+
42
+ ====================== ======================
43
+ Input Annotation types Output Annotation type
44
+ ====================== ======================
45
+ ``DOCUMENT`` ``SENTENCE_EMBEDDINGS``
46
+ ====================== ======================
47
+
48
+ Parameters
49
+ ----------
50
+ batchSize
51
+ Size of every batch , by default 8
52
+ dimension
53
+ Number of embedding dimensions, by default 768
54
+ caseSensitive
55
+ Whether to ignore case in tokens for embeddings matching, by default False
56
+ maxSentenceLength
57
+ Max sentence length to process, by default 512
58
+ configProtoBytes
59
+ ConfigProto from tensorflow, serialized into byte array.
60
+
61
+
62
+
63
+ Examples
64
+ --------
65
+ >>> import sparknlp
66
+ >>> from sparknlp.base import *
67
+ >>> from sparknlp.annotator import *
68
+ >>> from pyspark.ml import Pipeline
69
+ >>> documentAssembler = DocumentAssembler() \\
70
+ ... .setInputCol("text") \\
71
+ ... .setOutputCol("document")
72
+ >>> embeddings = MxbaiEmbeddings.pretrained() \\
73
+ ... .setInputCols(["document"]) \\
74
+ ... .setOutputCol("embeddings")
75
+ >>> embeddingsFinisher = EmbeddingsFinisher() \\
76
+ ... .setInputCols("embeddings") \\
77
+ ... .setOutputCols("finished_embeddings") \\
78
+ ... .setOutputAsVector(True)
79
+ >>> pipeline = Pipeline().setStages([
80
+ ... documentAssembler,
81
+ ... embeddings,
82
+ ... embeddingsFinisher
83
+ ... ])
84
+ >>> data = spark.createDataFrame([["hello world", "hello moon"]]).toDF("text")
85
+ >>> result = pipeline.fit(data).transform(data)
86
+ >>> result.selectExpr("explode(finished_embeddings) as result").show(5, 80)
87
+ +--------------------------------------------------------------------------------+
88
+ | result|
89
+ +--------------------------------------------------------------------------------+
90
+ |[0.50387806, 0.5861606, 0.35129607, -0.76046336, -0.32446072, -0.117674336, 0...|
91
+ |[0.6660665, 0.961762, 0.24854276, -0.1018044, -0.6569202, 0.027635604, 0.1915...|
92
+ +--------------------------------------------------------------------------------+
93
+ """
94
+
95
+ name = "MxbaiEmbeddings"
96
+
97
+ inputAnnotatorTypes = [AnnotatorType.DOCUMENT]
98
+
99
+ outputAnnotatorType = AnnotatorType.SENTENCE_EMBEDDINGS
100
+ poolingStrategy = Param(Params._dummy(),
101
+ "poolingStrategy",
102
+ "Pooling strategy to use for sentence embeddings",
103
+ TypeConverters.toString)
104
+
105
+ def setPoolingStrategy(self, value):
106
+ """Pooling strategy to use for sentence embeddings.
107
+
108
+ Available pooling strategies for sentence embeddings are:
109
+ - `"cls"`: leading `[CLS]` token
110
+ - `"cls_avg"`: leading `[CLS]` token + mean of all other tokens
111
+ - `"last"`: embeddings of the last token in the sequence
112
+ - `"avg"`: mean of all tokens
113
+ - `"max"`: max of all embedding features of the entire token sequence
114
+ - `"int"`: An integer number, which represents the index of the token to use as the
115
+ embedding
116
+
117
+ Parameters
118
+ ----------
119
+ value : str
120
+ Pooling strategy to use for sentence embeddings
121
+ """
122
+
123
+ valid_strategies = {"cls", "cls_avg", "last", "avg", "max"}
124
+ if value in valid_strategies or value.isdigit():
125
+ return self._set(poolingStrategy=value)
126
+ else:
127
+ raise ValueError(f"Invalid pooling strategy: {value}. "
128
+ f"Valid strategies are: {', '.join(self.valid_strategies)} or an integer.")
129
+
130
+ @keyword_only
131
+ def __init__(self, classname="com.johnsnowlabs.nlp.embeddings.MxbaiEmbeddings", java_model=None):
132
+ super(MxbaiEmbeddings, self).__init__(
133
+ classname=classname,
134
+ java_model=java_model
135
+ )
136
+ self._setDefault(
137
+ dimension=1024,
138
+ batchSize=8,
139
+ maxSentenceLength=512,
140
+ caseSensitive=False,
141
+ poolingStrategy="cls"
142
+ )
143
+
144
+ @staticmethod
145
+ def loadSavedModel(folder, spark_session):
146
+ """Loads a locally saved model.
147
+
148
+ Parameters
149
+ ----------
150
+ folder : str
151
+ Folder of the saved model
152
+ spark_session : pyspark.sql.SparkSession
153
+ The current SparkSession
154
+
155
+ Returns
156
+ -------
157
+ MxbaiEmbeddings
158
+ The restored model
159
+ """
160
+ from sparknlp.internal import _MxbaiEmbeddingsLoader
161
+ jModel = _MxbaiEmbeddingsLoader(folder, spark_session._jsparkSession)._java_obj
162
+ return MxbaiEmbeddings(java_model=jModel)
163
+
164
+ @staticmethod
165
+ def pretrained(name="mxbai_large_v1", lang="en", remote_loc=None):
166
+ """Downloads and loads a pretrained model.
167
+
168
+ Parameters
169
+ ----------
170
+ name : str, optional
171
+ Name of the pretrained model, by default "mxbai_large_v1"
172
+ lang : str, optional
173
+ Language of the pretrained model, by default "en"
174
+ remote_loc : str, optional
175
+ Optional remote address of the resource, by default None. Will use
176
+ Spark NLPs repositories otherwise.
177
+
178
+ Returns
179
+ -------
180
+ MxbaiEmbeddings
181
+ The restored model
182
+ """
183
+ from sparknlp.pretrained import ResourceDownloader
184
+ return ResourceDownloader.downloadModel(MxbaiEmbeddings, name, lang, remote_loc)