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

{spark_nlp-5.5.0.dist-info → spark_nlp-5.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: spark-nlp
-Version: 5.5.0
+Version: 5.5.1
 Summary: John Snow Labs Spark NLP is a natural language processing library built on top of Apache Spark ML. It provides simple, performant & accurate NLP annotations for machine learning pipelines, that scale easily in a distributed environment.
 Home-page: https://github.com/JohnSnowLabs/spark-nlp
 Author: John Snow Labs
@@ -95,7 +95,7 @@ $ java -version
 $ conda create -n sparknlp python=3.7 -y
 $ conda activate sparknlp
 # spark-nlp by default is based on pyspark 3.x
-$ pip install spark-nlp==5.5.0 pyspark==3.3.1
+$ pip install spark-nlp==5.5.1 pyspark==3.3.1
 ```
 
 In Python console or Jupyter `Python3` kernel:
@@ -161,7 +161,7 @@ For a quick example of using pipelines and models take a look at our official [d
 
 ### Apache Spark Support
 
-Spark NLP *5.5.0* has been built on top of Apache Spark 3.4 while fully supports Apache Spark 3.0.x, 3.1.x, 3.2.x, 3.3.x, 3.4.x, and 3.5.x
+Spark NLP *5.5.1* has been built on top of Apache Spark 3.4 while fully supports Apache Spark 3.0.x, 3.1.x, 3.2.x, 3.3.x, 3.4.x, and 3.5.x
 
 | Spark NLP | Apache Spark 3.5.x | Apache Spark 3.4.x | Apache Spark 3.3.x | Apache Spark 3.2.x | Apache Spark 3.1.x | Apache Spark 3.0.x | Apache Spark 2.4.x | Apache Spark 2.3.x |
 |-----------|--------------------|--------------------|--------------------|--------------------|--------------------|--------------------|--------------------|--------------------|
@@ -189,7 +189,7 @@ Find out more about 4.x `SparkNLP` versions in our official [documentation](http
 
 ### Databricks Support
 
-Spark NLP 5.5.0 has been tested and is compatible with the following runtimes:
+Spark NLP 5.5.1 has been tested and is compatible with the following runtimes:
 
 | **CPU**            | **GPU**            |
 |--------------------|--------------------|
@@ -206,7 +206,7 @@ We are compatible with older runtimes. For a full list check databricks support
 
 ### EMR Support
 
-Spark NLP 5.5.0 has been tested and is compatible with the following EMR releases:
+Spark NLP 5.5.1 has been tested and is compatible with the following EMR releases:
 
 | **EMR Release**    |
 |--------------------|
@@ -237,7 +237,7 @@ deployed to Maven central. To add any of our packages as a dependency in your ap
 from our official documentation.
 
 If you are interested, there is a simple SBT project for Spark NLP to guide you on how to use it in your
-projects [Spark NLP SBT S5.5.0r](https://github.com/maziyarpanahi/spark-nlp-starter)
+projects [Spark NLP SBT S5.5.1r](https://github.com/maziyarpanahi/spark-nlp-starter)
 
 ### Python
 
@@ -282,7 +282,7 @@ In Spark NLP we can define S3 locations to:
 
 Please check [these instructions](https://sparknlp.org/docs/en/install#s3-integration) from our official documentation.
 
-## Document5.5.0
+## Document5.5.1
 
 ### Examples
 
@@ -315,7 +315,7 @@ the Spark NLP library:
     keywords = {Spark, Natural language processing, Deep learning, Tensorflow, Cluster},
     abstract = {Spark NLP is a Natural Language Processing (NLP) library built on top of Apache Spark ML. It provides simple, performant & accurate NLP annotations for machine learning pipelines that can scale easily in a distributed environment. Spark NLP comes with 1100+ pretrained pipelines and models in more than 192+ languages. It supports nearly all the NLP tasks and modules that can be used seamlessly in a cluster. Downloaded more than 2.7 million times and experiencing 9x growth since January 2020, Spark NLP is used by 54% of healthcare organizations as the world’s most widely used NLP library in the enterprise.}
     }
-}5.5.0
+}5.5.1
 ```
 
 ## Community support

{spark_nlp-5.5.0.dist-info → spark_nlp-5.5.1.dist-info}/RECORD

@@ -3,7 +3,7 @@ com/johnsnowlabs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 com/johnsnowlabs/ml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 com/johnsnowlabs/ml/ai/__init__.py,sha256=YQiK2M7U4d8y5irPy_HB8ae0mSpqS9583MH44pnKJXc,295
 com/johnsnowlabs/nlp/__init__.py,sha256=DPIVXtONO5xXyOk-HB0-sNiHAcco17NN13zPS_6Uw8c,294
-sparknlp/__init__.py,sha256=AKjfuQ0s3OwAuC0XQj4FjGRISLl5RY8bO3WdLpZpQaA,13638
+sparknlp/__init__.py,sha256=26U34YGDYCBbYHr3rpxvy71snUYuoFhl5XfUvxkPv7M,13638
 sparknlp/annotation.py,sha256=I5zOxG5vV2RfPZfqN9enT1i4mo6oBcn3Lrzs37QiOiA,5635
 sparknlp/annotation_audio.py,sha256=iRV_InSVhgvAwSRe9NTbUH9v6OGvTM-FPCpSAKVu0mE,1917
 sparknlp/annotation_image.py,sha256=xhCe8Ko-77XqWVuuYHFrjKqF6zPd8Z-RY_rmZXNwCXU,2547
@@ -30,12 +30,13 @@ sparknlp/annotator/audio/__init__.py,sha256=dXjtvi5c0aTZFq1Q_JciUd1uFTBVSJoUdcq0
 sparknlp/annotator/audio/hubert_for_ctc.py,sha256=76PfwPZZvOHU5kfDqLueCFbmqa4W8pMNRGoCvOqjsEA,7859
 sparknlp/annotator/audio/wav2vec2_for_ctc.py,sha256=K78P1U6vA4O1UufsLYzy0H7arsKNmwPcIV7kzDFsA5Q,6210
 sparknlp/annotator/audio/whisper_for_ctc.py,sha256=uII51umuohqwnAW0Q7VdxEFyr_j5LMnfpcRlf8TbetA,9800
-sparknlp/annotator/classifier_dl/__init__.py,sha256=74WL0W2zBfx6v0tJpx1DcRfZENs86n9JxizDDBEE41A,3934
+sparknlp/annotator/classifier_dl/__init__.py,sha256=4v2_3kSWQFFBc_KzaJ0gEC6ANVJpy5tsHa6CJGc4nCw,4005
 sparknlp/annotator/classifier_dl/albert_for_question_answering.py,sha256=LG2dL6Fky1T35yXTUZBfIihIIGnkRFQ7ECQ3HRXXEG8,6517
 sparknlp/annotator/classifier_dl/albert_for_sequence_classification.py,sha256=kWx7f9pcKE2qw319gn8FN0Md5dX38gbmfeoY9gWCLNk,7842
 sparknlp/annotator/classifier_dl/albert_for_token_classification.py,sha256=5rdsjWnsAVmtP-idU7ATKJ8lkH2rtlKZLnpi4Mq27eI,6839
 sparknlp/annotator/classifier_dl/albert_for_zero_shot_classification.py,sha256=_TgV6EiIOiD_djA3fxfoz-o37mzMeKbn6iL2kZ6GzO0,8366
 sparknlp/annotator/classifier_dl/bart_for_zero_shot_classification.py,sha256=yqQeDdpLbNOKuSZejZjSAjT8ydYyxsTVf2aFDgSSDfc,8767
+sparknlp/annotator/classifier_dl/bert_for_multiple_choice.py,sha256=Ew_NGBj7F5ApgK3SyQh2HIfjD7ZTqTs0LZEQxjwoyto,5936
 sparknlp/annotator/classifier_dl/bert_for_question_answering.py,sha256=2euY_RAdMPA4IHJXZAd5MkQojFOtFNhB_hSc1iVQ5DQ,6433
 sparknlp/annotator/classifier_dl/bert_for_sequence_classification.py,sha256=AzD3RQcRuQc0DDTbL6vGiacTtHlZnbAqksNvRQq7EQE,7800
 sparknlp/annotator/classifier_dl/bert_for_token_classification.py,sha256=uJXoDLPfPWiRmKqtw_3lLBvneIirj87S2JWwfd33zq8,6668
@@ -173,7 +174,7 @@ sparknlp/annotator/token/regex_tokenizer.py,sha256=FG2HvFwMb1G_4grfyIQaeBpaAgKv_
 sparknlp/annotator/token/tokenizer.py,sha256=Me3P3wogUKUJ7O7_2wLdPzF00vKpp_sHuiztpGWRVpU,19939
 sparknlp/annotator/ws/__init__.py,sha256=-l8bnl8Z6lGXWOBdRIBZ6958fzTHt4o87QhhLHIFF8A,693
 sparknlp/annotator/ws/word_segmenter.py,sha256=rrbshwn5wzXIHpCCDji6ZcsmiARpuA82_p_6TgNHfRc,16365
-sparknlp/base/__init__.py,sha256=iC4b4NzTDsqEu3eE_f5QL8JD-uoig4Pn2h1ZMyPHR6Q,1266
+sparknlp/base/__init__.py,sha256=fCL-kReIavZceUa1OC99pSRH7MsXzqGB8BXgzVS_f7s,1311
 sparknlp/base/audio_assembler.py,sha256=HKa9mXvmuMUrjTihUZkppGj-WJjcUrm2BGapNuPifyI,3320
 sparknlp/base/doc2_chunk.py,sha256=TyvbdJNkVo9favHlOEoH5JwKbjpk5ZVJ75p8Cilp9jM,6551
 sparknlp/base/document_assembler.py,sha256=zl-SXWMTR3B0EZ8z6SWYchCwEo-61FhU6u7dHUKDIOg,6697
@@ -185,6 +186,7 @@ sparknlp/base/has_recursive_transform.py,sha256=UkGNgo4LMsjQC-Coeefg4bJcg7FoPcPi
 sparknlp/base/image_assembler.py,sha256=HytRoYJTLMqGtvScHoFnp6CasG9IVNYAHYiT2_rrmeE,3719
 sparknlp/base/light_pipeline.py,sha256=Jk2DLpT4PLHCANlOo_WetTdPba_5lYs3ywiyY3lM-PE,16577
 sparknlp/base/multi_document_assembler.py,sha256=4htET1fRAeOB6zhsNXsBq5rKZvn-LGD4vrFRjPZeqow,7070
+sparknlp/base/prompt_assembler.py,sha256=ysU4Vbmnuv2UBHK0JBkYrxiZiJ7_GTcVMip1-QRmheI,11570
 sparknlp/base/recursive_pipeline.py,sha256=V9rTnu8KMwgjoceykN9pF1mKGtOkkuiC_n9v8dE3LDk,4279
 sparknlp/base/table_assembler.py,sha256=Kxu3R2fY6JgCxEc07ibsMsjip6dgcPDHLiWAZ8gC_d8,5102
 sparknlp/base/token_assembler.py,sha256=qiHry07L7mVCqeHSH6hHxLygv1AsfZIE4jy1L75L3Do,5075
@@ -200,7 +202,7 @@ sparknlp/common/read_as.py,sha256=imxPGwV7jr4Li_acbo0OAHHRGCBbYv-akzEGaBWEfcY,12
 sparknlp/common/recursive_annotator_approach.py,sha256=vqugBw22cE3Ff7PIpRlnYFuOlchgL0nM26D8j-NdpqU,1449
 sparknlp/common/storage.py,sha256=D91H3p8EIjNspjqAYu6ephRpCUtdcAir4_PrAbkIQWE,4842
 sparknlp/common/utils.py,sha256=Yne6yYcwKxhOZC-U4qfYoDhWUP_6BIaAjI5X_P_df1E,1306
-sparknlp/internal/__init__.py,sha256=nK-9lncAVRXmyP8ATbiMwRnLJVe4IEd_r5Z3gEqDK3g,33672
+sparknlp/internal/__init__.py,sha256=ljEf4IUraCdKU7gKFxNwFxlX-FHcnkG6sqs1MxEhLSQ,33967
 sparknlp/internal/annotator_java_ml.py,sha256=UGPoThG0rGXUOXGSQnDzEDW81Mu1s5RPF29v7DFyE3c,1187
 sparknlp/internal/annotator_transformer.py,sha256=fXmc2IWXGybqZpbEU9obmbdBYPc798y42zvSB4tqV9U,1448
 sparknlp/internal/extended_java_wrapper.py,sha256=hwP0133-hDiDf5sBF-P3MtUsuuDj1PpQbtGZQIRwzfk,2240
@@ -242,8 +244,8 @@ sparknlp/training/_tf_graph_builders_1x/ner_dl/dataset_encoder.py,sha256=R4yHFN3
 sparknlp/training/_tf_graph_builders_1x/ner_dl/ner_model.py,sha256=EoCSdcIjqQ3wv13MAuuWrKV8wyVBP0SbOEW41omHlR0,23189
 sparknlp/training/_tf_graph_builders_1x/ner_dl/ner_model_saver.py,sha256=k5CQ7gKV6HZbZMB8cKLUJuZxoZWlP_DFWdZ--aIDwsc,2356
 sparknlp/training/_tf_graph_builders_1x/ner_dl/sentence_grouper.py,sha256=pAxjWhjazSX8Vg0MFqJiuRVw1IbnQNSs-8Xp26L4nko,870
-spark_nlp-5.5.0.dist-info/.uuid,sha256=1f6hF51aIuv9yCvh31NU9lOpS34NE-h3a0Et7R9yR6A,36
-spark_nlp-5.5.0.dist-info/METADATA,sha256=FccpjBJS2ERU0kJM5kPn_bUo4VyX4l8tHLImemv6czo,19156
-spark_nlp-5.5.0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-spark_nlp-5.5.0.dist-info/top_level.txt,sha256=uuytur4pyMRw2H_txNY2ZkaucZHUs22QF8-R03ch_-E,13
-spark_nlp-5.5.0.dist-info/RECORD,,
+spark_nlp-5.5.1.dist-info/.uuid,sha256=1f6hF51aIuv9yCvh31NU9lOpS34NE-h3a0Et7R9yR6A,36
+spark_nlp-5.5.1.dist-info/METADATA,sha256=Y7Y0nf18tO2RfHzagHWWZpn4QRrF50d5wP3hXG1eFyw,19156
+spark_nlp-5.5.1.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+spark_nlp-5.5.1.dist-info/top_level.txt,sha256=uuytur4pyMRw2H_txNY2ZkaucZHUs22QF8-R03ch_-E,13
+spark_nlp-5.5.1.dist-info/RECORD,,

sparknlp/__init__.py

@@ -129,7 +129,7 @@ def start(gpu=False,
         The initiated Spark session.
 
     """
-    current_version = "5.5.0"
+    current_version = "5.5.1"
 
     if params is None:
         params = {}
@@ -310,4 +310,4 @@ def version():
     str
         The current Spark NLP version.
     """
-    return '5.5.0'
+    return '5.5.1'

sparknlp/annotator/classifier_dl/__init__.py

@@ -54,4 +54,4 @@ from sparknlp.annotator.classifier_dl.mpnet_for_question_answering import *
 from sparknlp.annotator.classifier_dl.mpnet_for_token_classification import *
 from sparknlp.annotator.classifier_dl.albert_for_zero_shot_classification import *
 from sparknlp.annotator.classifier_dl.camembert_for_zero_shot_classification import *
-
+from sparknlp.annotator.classifier_dl.bert_for_multiple_choice import *

sparknlp/annotator/classifier_dl/bert_for_multiple_choice.py

@@ -0,0 +1,161 @@
+#  Copyright 2017-2024 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.
+
+from sparknlp.common import *
+
+class BertForMultipleChoice(AnnotatorModel,
+                            HasCaseSensitiveProperties,
+                            HasBatchedAnnotate,
+                            HasEngine,
+                            HasMaxSentenceLengthLimit):
+    """BertForMultipleChoice can load BERT Models with a multiple choice classification head on top
+    (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/SWAG tasks.
+
+    Pretrained models can be loaded with :meth:`.pretrained` of the companion
+    object:
+
+    >>> spanClassifier = BertForMultipleChoice.pretrained() \\
+    ...     .setInputCols(["document_question", "document_context"]) \\
+    ...     .setOutputCol("answer")
+
+    The default model is ``"bert_base_uncased_multiple_choice"``, if no name is
+    provided.
+
+    For available pretrained models please see the `Models Hub
+    <https://sparknlp.org/models?task=Multiple+Choice>`__.
+
+    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>`_.
+
+    ====================== ======================
+    Input Annotation types Output Annotation type
+    ====================== ======================
+    ``DOCUMENT, DOCUMENT``    ``CHUNK``
+    ====================== ======================
+
+    Parameters
+    ----------
+    batchSize
+        Batch size. Large values allows faster processing but requires more
+        memory, by default 8
+    caseSensitive
+        Whether to ignore case in tokens for embeddings matching, by default
+        False
+    maxSentenceLength
+        Max sentence length to process, by default 512
+
+    Examples
+    --------
+    >>> import sparknlp
+    >>> from sparknlp.base import *
+    >>> from sparknlp.annotator import *
+    >>> from pyspark.ml import Pipeline
+    >>> documentAssembler = MultiDocumentAssembler() \\
+    ...     .setInputCols(["question", "context"]) \\
+    ...     .setOutputCols(["document_question", "document_context"])
+    >>> questionAnswering = BertForMultipleChoice.pretrained() \\
+    ...     .setInputCols(["document_question", "document_context"]) \\
+    ...     .setOutputCol("answer") \\
+    ...     .setCaseSensitive(False)
+    >>> pipeline = Pipeline().setStages([
+    ...     documentAssembler,
+    ...     questionAnswering
+    ... ])
+    >>> data = spark.createDataFrame([["The Eiffel Tower is located in which country??", "Germany, France, Italy"]]).toDF("question", "context")
+    >>> result = pipeline.fit(data).transform(data)
+    >>> result.select("answer.result").show(truncate=False)
+    +--------------------+
+    |result              |
+    +--------------------+
+    |[France]             |
+    +--------------------+
+    """
+    name = "BertForMultipleChoice"
+
+    inputAnnotatorTypes = [AnnotatorType.DOCUMENT, AnnotatorType.DOCUMENT]
+
+    outputAnnotatorType = AnnotatorType.CHUNK
+
+    choicesDelimiter = Param(Params._dummy(),
+                         "choicesDelimiter",
+                         "Delimiter character use to split the choices",
+                         TypeConverters.toString)
+
+    def setChoicesDelimiter(self, value):
+        """Sets delimiter character use to split the choices
+
+        Parameters
+        ----------
+        value : string
+            Delimiter character use to split the choices
+        """
+        return self._set(caseSensitive=value)
+
+    @keyword_only
+    def __init__(self, classname="com.johnsnowlabs.nlp.annotators.classifier.dl.BertForMultipleChoice",
+                 java_model=None):
+        super(BertForMultipleChoice, self).__init__(
+            classname=classname,
+            java_model=java_model
+        )
+        self._setDefault(
+            batchSize=4,
+            maxSentenceLength=512,
+            caseSensitive=False,
+            choicesDelimiter = ","
+        )
+
+    @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
+        -------
+        BertForQuestionAnswering
+            The restored model
+        """
+        from sparknlp.internal import _BertMultipleChoiceLoader
+        jModel = _BertMultipleChoiceLoader(folder, spark_session._jsparkSession)._java_obj
+        return BertForMultipleChoice(java_model=jModel)
+
+    @staticmethod
+    def pretrained(name="bert_base_uncased_multiple_choice", lang="en", remote_loc=None):
+        """Downloads and loads a pretrained model.
+
+        Parameters
+        ----------
+        name : str, optional
+            Name of the pretrained model, by default
+            "bert_base_uncased_multiple_choice"
+        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
+        -------
+        BertForQuestionAnswering
+            The restored model
+        """
+        from sparknlp.pretrained import ResourceDownloader
+        return ResourceDownloader.downloadModel(BertForMultipleChoice, name, lang, remote_loc)

sparknlp/base/__init__.py

@@ -26,3 +26,4 @@ from sparknlp.base.token_assembler import *
 from sparknlp.base.image_assembler import *
 from sparknlp.base.audio_assembler import *
 from sparknlp.base.table_assembler import *
+from sparknlp.base.prompt_assembler import *

sparknlp/base/prompt_assembler.py

@@ -0,0 +1,207 @@
+#  Copyright 2017-2024 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 the PromptAssembler."""
+
+from pyspark import keyword_only
+from pyspark.ml.param import TypeConverters, Params, Param
+
+from sparknlp.common import AnnotatorType
+from sparknlp.internal import AnnotatorTransformer
+
+
+class PromptAssembler(AnnotatorTransformer):
+    """Assembles a sequence of messages into a single string using a template. These strings can then
+    be used as prompts for large language models.
+
+    This annotator expects an array of two-tuples as the type of the input column (one array of
+    tuples per row). The first element of the tuples should be the role and the second element is
+    the text of the message. Possible roles are "system", "user" and "assistant".
+
+    An assistant header can be added to the end of the generated string by using
+    ``setAddAssistant(True)``.
+
+    At the moment, this annotator uses llama.cpp as a backend to parse and apply the templates.
+    llama.cpp uses basic pattern matching to determine the type of the template, then applies a
+    basic version of the template to the messages. This means that more advanced templates are not
+    supported.
+
+    For an extended example see the
+    `example notebook <https://github.com/JohnSnowLabs/spark-nlp/blob/master/examples/python/llama.cpp/PromptAssember_with_AutoGGUFModel.ipynb>`__.
+
+    ====================== ======================
+    Input Annotation types Output Annotation type
+    ====================== ======================
+    ``NONE``               ``DOCUMENT``
+    ====================== ======================
+
+    Parameters
+    ----------
+    inputCol
+        Input column name
+    outputCol
+        Output column name
+    chatTemplate
+        Template used for the chat
+    addAssistant
+        Whether to add an assistant header to the end of the generated string
+
+    Examples
+    --------
+    >>> from sparknlp.base import *
+    >>> messages = [
+    ...     [
+    ...         ("system", "You are a helpful assistant."),
+    ...         ("assistant", "Hello there, how can I help you?"),
+    ...         ("user", "I need help with organizing my room."),
+    ...     ]
+    ... ]
+    >>> df = spark.createDataFrame([messages]).toDF("messages")
+    >>> template = (
+    ...     "{{- bos_token }} {%- if custom_tools is defined %} {%- set tools = custom_tools %} {%- "
+    ...     "endif %} {%- if not tools_in_user_message is defined %} {%- set tools_in_user_message = true %} {%- "
+    ...     'endif %} {%- if not date_string is defined %} {%- set date_string = "26 Jul 2024" %} {%- endif %} '
+    ...     "{%- if not tools is defined %} {%- set tools = none %} {%- endif %} {#- This block extracts the "
+    ...     "system message, so we can slot it into the right place. #} {%- if messages[0]['role'] == 'system' %}"
+    ...     " {%- set system_message = messages[0]['content']|trim %} {%- set messages = messages[1:] %} {%- else"
+    ...     ' %} {%- set system_message = "" %} {%- endif %} {#- System message + builtin tools #} {{- '
+    ...     '"<|start_header_id|>system<|end_header_id|>\\n\\n" }} {%- if builtin_tools is defined or tools is '
+    ...     'not none %} {{- "Environment: ipython\\n" }} {%- endif %} {%- if builtin_tools is defined %} {{- '
+    ...     '"Tools: " + builtin_tools | reject(\\'equalto\', \\'code_interpreter\\') | join(", ") + "\\n\\n"}} '
+    ...     '{%- endif %} {{- "Cutting Knowledge Date: December 2023\\n" }} {{- "Today Date: " + date_string '
+    ...     '+ "\\n\\n" }} {%- if tools is not none and not tools_in_user_message %} {{- "You have access to '
+    ...     'the following functions. To call a function, please respond with JSON for a function call." }} {{- '
+    ...     '\\'Respond in the format {"name": function name, "parameters": dictionary of argument name and its'
+    ...     ' value}.\\' }} {{- "Do not use variables.\\n\\n" }} {%- for t in tools %} {{- t | tojson(indent=4) '
+    ...     '}} {{- "\\n\\n" }} {%- endfor %} {%- endif %} {{- system_message }} {{- "<|eot_id|>" }} {#- '
+    ...     "Custom tools are passed in a user message with some extra guidance #} {%- if tools_in_user_message "
+    ...     "and not tools is none %} {#- Extract the first user message so we can plug it in here #} {%- if "
+    ...     "messages | length != 0 %} {%- set first_user_message = messages[0]['content']|trim %} {%- set "
+    ...     'messages = messages[1:] %} {%- else %} {{- raise_exception("Cannot put tools in the first user '
+    ...     "message when there's no first user message!\\") }} {%- endif %} {{- "
+    ...     "'<|start_header_id|>user<|end_header_id|>\\n\\n' -}} {{- \\"Given the following functions, please "
+    ...     'respond with a JSON for a function call " }} {{- "with its proper arguments that best answers the '
+    ...     'given prompt.\\n\\n" }} {{- \\'Respond in the format {"name": function name, "parameters": '
+    ...     'dictionary of argument name and its value}.\\' }} {{- "Do not use variables.\\n\\n" }} {%- for t in '
+    ...     'tools %} {{- t | tojson(indent=4) }} {{- "\\n\\n" }} {%- endfor %} {{- first_user_message + '
+    ...     "\\"<|eot_id|>\\"}} {%- endif %} {%- for message in messages %} {%- if not (message.role == 'ipython' "
+    ...     "or message.role == 'tool' or 'tool_calls' in message) %} {{- '<|start_header_id|>' + message['role']"
+    ...     " + '<|end_header_id|>\\n\\n'+ message['content'] | trim + '<|eot_id|>' }} {%- elif 'tool_calls' in "
+    ...     'message %} {%- if not message.tool_calls|length == 1 %} {{- raise_exception("This model only '
+    ...     'supports single tool-calls at once!") }} {%- endif %} {%- set tool_call = message.tool_calls[0]'
+    ...     ".function %} {%- if builtin_tools is defined and tool_call.name in builtin_tools %} {{- "
+    ...     "'<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}} {{- \\"<|python_tag|>\\" + tool_call.name + "
+    ...     '".call(" }} {%- for arg_name, arg_val in tool_call.arguments | items %} {{- arg_name + \\'="\\' + '
+    ...     'arg_val + \\'"\\' }} {%- if not loop.last %} {{- ", " }} {%- endif %} {%- endfor %} {{- ")" }} {%- '
+    ...     "else %} {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}} {{- '{\\"name\": \\"' + "
+    ...     'tool_call.name + \\'", \\' }} {{- \\'"parameters": \\' }} {{- tool_call.arguments | tojson }} {{- "}" '
+    ...     "}} {%- endif %} {%- if builtin_tools is defined %} {#- This means we're in ipython mode #} {{- "
+    ...     '"<|eom_id|>" }} {%- else %} {{- "<|eot_id|>" }} {%- endif %} {%- elif message.role == "tool" '
+    ...     'or message.role == "ipython" %} {{- "<|start_header_id|>ipython<|end_header_id|>\\n\\n" }} {%- '
+    ...     "if message.content is mapping or message.content is iterable %} {{- message.content | tojson }} {%- "
+    ...     'else %} {{- message.content }} {%- endif %} {{- "<|eot_id|>" }} {%- endif %} {%- endfor %} {%- if '
+    ...     "add_generation_prompt %} {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }} {%- endif %} "
+    ... )
+    >>> prompt_assembler = (
+    ...     PromptAssembler()
+    ...     .setInputCol("messages")
+    ...     .setOutputCol("prompt")
+    ...     .setChatTemplate(template)
+    ... )
+    >>> prompt_assembler.transform(df).select("prompt.result").show(truncate=False)
+    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+    |result                                                                                                                                                                                                                                                                                                                      |
+    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+    |[<|start_header_id|>system<|end_header_id|>\n\nYou are a helpful assistant.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nHello there, how can I help you?<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nI need help with organizing my room.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n]|
+    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+    """
+
+    outputAnnotatorType = AnnotatorType.DOCUMENT
+
+    inputCol = Param(
+        Params._dummy(),
+        "inputCol",
+        "input column name",
+        typeConverter=TypeConverters.toString,
+    )
+    outputCol = Param(
+        Params._dummy(),
+        "outputCol",
+        "output column name",
+        typeConverter=TypeConverters.toString,
+    )
+    chatTemplate = Param(
+        Params._dummy(),
+        "chatTemplate",
+        "Template used for the chat",
+        typeConverter=TypeConverters.toString,
+    )
+    addAssistant = Param(
+        Params._dummy(),
+        "addAssistant",
+        "Whether to add an assistant header to the end of the generated string",
+        typeConverter=TypeConverters.toBoolean,
+    )
+    name = "PromptAssembler"
+
+    @keyword_only
+    def __init__(self):
+        super(PromptAssembler, self).__init__(
+            classname="com.johnsnowlabs.nlp.PromptAssembler"
+        )
+        self._setDefault(outputCol="prompt", addAssistant=True)
+
+    @keyword_only
+    def setParams(self):
+        kwargs = self._input_kwargs
+        return self._set(**kwargs)
+
+    def setInputCol(self, value):
+        """Sets input column name.
+
+        Parameters
+        ----------
+        value : str
+            Name of the input column
+        """
+        return self._set(inputCol=value)
+
+    def setOutputCol(self, value):
+        """Sets output column name.
+
+        Parameters
+        ----------
+        value : str
+            Name of the Output Column
+        """
+        return self._set(outputCol=value)
+
+    def setChatTemplate(self, value):
+        """Sets the chat template.
+
+        Parameters
+        ----------
+        value : str
+            Template used for the chat
+        """
+        return self._set(chatTemplate=value)
+
+    def setAddAssistant(self, value):
+        """Sets whether to add an assistant header to the end of the generated string.
+
+        Parameters
+        ----------
+        value : bool
+            Whether to add an assistant header to the end of the generated string
+        """
+        return self._set(addAssistant=value)

sparknlp/internal/__init__.py

@@ -113,6 +113,13 @@ class _BertQuestionAnsweringLoader(ExtendedJavaWrapper):
             jspark,
         )
 
+class _BertMultipleChoiceLoader(ExtendedJavaWrapper):
+    def __init__(self, path, jspark):
+        super(_BertMultipleChoiceLoader, self).__init__(
+            "com.johnsnowlabs.nlp.annotators.classifier.dl.BertForMultipleChoice.loadSavedModel",
+            path,
+            jspark,
+        )
 
 class _DeBERTaLoader(ExtendedJavaWrapper):
     def __init__(self, path, jspark):