# 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.
import subprocess
import sys
import threading
from pyspark import __version__ as pyspark_version
from pyspark.conf import SparkConf
from pyspark.context import SparkContext
from pyspark.java_gateway import launch_gateway
from pyspark.sql import SparkSession
from sparknlp import annotator
from sparknlp._maven import resolve_spark_nlp_coordinate
# Must be declared here one by one or else PretrainedPipeline will fail with AttributeError
from sparknlp.base import DocumentAssembler, MultiDocumentAssembler, Finisher, EmbeddingsFinisher, TokenAssembler, \
Doc2Chunk, AudioAssembler, GraphFinisher, ImageAssembler, TableAssembler, MultiColumnAssembler
from sparknlp.reader import SparkNLPReader
sys.modules['com.johnsnowlabs.nlp.annotators'] = annotator
sys.modules['com.johnsnsowlabs.nlp.annotators.tokenizer'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.tokenizer.wordpiece'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ner'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ner.regex'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ner.crf'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ner.dl'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.pos'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.pos.perceptron'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sbd'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sbd.pragmatic'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sbd.deep'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sda'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sda.pragmatic'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sda.vivekn'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.spell'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.spell.norvig'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.spell.symmetric'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.parser'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.parser.dep'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.parser.typdep'] = annotator
sys.modules['com.johnsnowlabs.nlp.embeddings'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.classifier'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.classifier.dl'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.spell.context'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ld'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ld.dl'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sentence_detector_dl'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.sbd.sat'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.seq2seq'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.ws'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.er'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.matcher'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.coref'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.cv'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.audio'] = annotator
sys.modules['com.johnsnowlabs.nlp.annotators.similarity'] = annotator
sys.modules['com.johnsnowlabs.ml.ai'] = annotator
__version__ = "7.0.0"
[docs]def start(gpu=False,
apple_silicon=False,
aarch64=False,
memory="16G",
cache_folder="",
log_folder="",
cluster_tmp_dir="",
params=None,
real_time_output=False,
output_level=1,
skip_sparknlp_maven=False):
"""Starts a PySpark instance with default parameters for Spark NLP.
The default parameters would result in the equivalent of:
.. code-block:: python
SparkSession.builder \\
.appName("Spark NLP") \\
.master("local[*]") \\
.config("spark.driver.memory", "16G") \\
.config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \\
.config("spark.kryoserializer.buffer.max", "2000M") \\
.config("spark.driver.maxResultSize", "0") \\
.config("spark.jars.packages", "<resolved Spark NLP Maven coordinate>") \\
.getOrCreate()
Parameters
----------
gpu : bool, optional
Whether to enable GPU acceleration (must be set up correctly), by default False
apple_silicon : bool, optional
Whether to enable Apple Silicon support for macOS
aarch64 : bool, optional
Whether to enable Linux Aarch64 support
memory : str, optional
How much memory to allocate for the Spark driver, by default "16G"
cache_folder : str, optional
The location to download and extract pretrained Models and Pipelines. If not
set, it will be in the users home directory under `cache_pretrained`.
log_folder : str, optional
The location to use on a cluster for temporarily files such as unpacking indexes
for WordEmbeddings. By default, this locations is the location of
`hadoop.tmp.dir` set via Hadoop configuration for Apache Spark. NOTE: `S3` is
not supported and it must be local, HDFS, or DBFS.
params : dict, optional
Custom parameters to set for the Spark configuration, by default None.
Set ``skip_sparknlp_maven`` to ``True`` or ``"true"`` to avoid adding the
default Spark NLP Maven package when a custom Spark NLP jar is provided.
cluster_tmp_dir : str, optional
The location to save logs from annotators during training. If not set, it will
be in the users home directory under `annotator_logs`.
real_time_output : bool, optional
Whether to read and print JVM output in real time, by default False
output_level : int, optional
Output level for logs, by default 1
skip_sparknlp_maven : bool, optional
Whether to avoid adding the default Spark NLP Maven package. Use this
when providing a custom Spark NLP jar with ``spark.jars``.
Notes
-----
The Maven artifact is selected from the installed PySpark version. Spark 3.x
uses Scala 2.12. Spark 4.0.0 uses the dedicated ``spark400`` Scala 2.13
artifact. Later Spark 4.x releases use the default Scala 2.13 artifact.
Automatic artifact selection does not imply that every release has been
validated. Only release versions in major.minor.patch format are accepted.
Returns
-------
:class:`SparkSession`
The initiated Spark session.
"""
current_version = __version__
maven_version = current_version.split("-")[0].split("+")[0]
if params is None:
params = {}
else:
if not isinstance(params, dict):
raise TypeError('params must be a dictionary like {"spark.executor.memory": "8G"}')
if '_instantiatedSession' in dir(SparkSession) and SparkSession._instantiatedSession is not None:
print('Warning::Spark Session already created, some configs may not take.')
skip_sparknlp_maven_param = "skip_sparknlp_maven"
def is_skip_sparknlp_maven_enabled():
value = params.get(skip_sparknlp_maven_param, False)
if isinstance(value, str):
value = value.strip().lower() == "true"
else:
value = bool(value)
return bool(skip_sparknlp_maven) or value
skip_sparknlp_maven = is_skip_sparknlp_maven_enabled()
spark_jars_packages = ""
if not skip_sparknlp_maven:
spark_jars_packages = resolve_spark_nlp_coordinate(
pyspark_version,
maven_version,
gpu=gpu,
apple_silicon=apple_silicon,
aarch64=aarch64,
)
driver_cores = "*"
for key, value in params.items():
if key == "spark.driver.cores":
driver_cores = f"{value}"
else:
driver_cores = "*"
class SparkNLPConfig:
def __init__(self):
self.master, self.app_name = "local[{}]".format(driver_cores), "Spark NLP"
self.serializer, self.serializer_max_buffer = "org.apache.spark.serializer.KryoSerializer", "2000M"
self.driver_max_result_size = "0"
def start_without_realtime_output():
builder = SparkSession.builder \
.appName(spark_nlp_config.app_name) \
.master(spark_nlp_config.master) \
.config("spark.driver.memory", memory) \
.config("spark.serializer", spark_nlp_config.serializer) \
.config("spark.kryoserializer.buffer.max", spark_nlp_config.serializer_max_buffer) \
.config("spark.driver.maxResultSize", spark_nlp_config.driver_max_result_size)
if cache_folder != '':
builder.config("spark.jsl.settings.pretrained.cache_folder", cache_folder)
if log_folder != '':
builder.config("spark.jsl.settings.annotator.log_folder", log_folder)
if cluster_tmp_dir != '':
builder.config("spark.jsl.settings.storage.cluster_tmp_dir", cluster_tmp_dir)
if not skip_sparknlp_maven and params.get("spark.jars.packages") is None:
builder.config("spark.jars.packages", spark_jars_packages)
for key, value in params.items():
if key == skip_sparknlp_maven_param:
continue
if key == "spark.jars.packages" and not skip_sparknlp_maven:
packages = spark_jars_packages + "," + value
builder.config(key, packages)
else:
builder.config(key, value)
spark_session = builder.getOrCreate()
apply_hadoop_params(spark_session)
return spark_session
def start_with_realtime_output():
class SparkWithCustomGateway:
def __init__(self):
spark_conf = SparkConf()
spark_conf.setAppName(spark_nlp_config.app_name)
spark_conf.setMaster(spark_nlp_config.master)
spark_conf.set("spark.driver.memory", memory)
spark_conf.set("spark.serializer", spark_nlp_config.serializer)
spark_conf.set("spark.kryoserializer.buffer.max", spark_nlp_config.serializer_max_buffer)
spark_conf.set("spark.driver.maxResultSize", spark_nlp_config.driver_max_result_size)
if cache_folder != '':
spark_conf.set("spark.jsl.settings.pretrained.cache_folder", cache_folder)
if log_folder != '':
spark_conf.set("spark.jsl.settings.annotator.log_folder", log_folder)
if cluster_tmp_dir != '':
spark_conf.set("spark.jsl.settings.storage.cluster_tmp_dir", cluster_tmp_dir)
if not skip_sparknlp_maven and params.get("spark.jars.packages") is None:
spark_conf.set("spark.jars.packages", spark_jars_packages)
for key, value in params.items():
if key == skip_sparknlp_maven_param:
continue
if key == "spark.jars.packages" and not skip_sparknlp_maven:
packages = spark_jars_packages + "," + value
spark_conf.set(key, packages)
else:
spark_conf.set(key, value)
# Make the py4j JVM stdout and stderr available without buffering
popen_kwargs = {
'stdout': subprocess.PIPE,
'stderr': subprocess.PIPE,
'bufsize': 0
}
# Launch the gateway with our custom settings
self.gateway = launch_gateway(conf=spark_conf, popen_kwargs=popen_kwargs)
self.process = self.gateway.proc
# Use the gateway we launched
spark_context = SparkContext(gateway=self.gateway)
self.spark_session = SparkSession(spark_context)
apply_hadoop_params(self.spark_session)
self.out_thread = threading.Thread(target=self.output_reader)
self.error_thread = threading.Thread(target=self.error_reader)
self.std_background_listeners()
def std_background_listeners(self):
self.out_thread.start()
self.error_thread.start()
def output_reader(self):
for line in iter(self.process.stdout.readline, b''):
print('{0}'.format(line.decode('utf-8')), end='')
def error_reader(self):
RED = '\033[91m'
RESET = '\033[0m'
for line in iter(self.process.stderr.readline, b''):
if output_level == 0:
print(RED + '{0}'.format(line.decode('utf-8')) + RESET, end='')
else:
# output just info
pass
def shutdown(self):
self.spark_session.stop()
self.gateway.shutdown()
self.process.communicate()
self.out_thread.join()
self.error_thread.join()
return SparkWithCustomGateway()
def apply_hadoop_params(spark_session):
hadoop_prefix = "spark.hadoop."
hadoop_configuration = spark_session.sparkContext._jsc.hadoopConfiguration()
for key, value in params.items():
if key.startswith(hadoop_prefix):
hadoop_configuration.set(key[len(hadoop_prefix):], value)
spark_nlp_config = SparkNLPConfig()
if real_time_output:
# Available from Spark 3.0.x
class SparkRealTimeOutput:
def __init__(self):
self.__spark_with_custom_gateway = start_with_realtime_output()
self.spark_session = self.__spark_with_custom_gateway.spark_session
def shutdown(self):
self.__spark_with_custom_gateway.shutdown()
return SparkRealTimeOutput().spark_session
else:
spark_session = start_without_realtime_output()
return spark_session
[docs]def read(params=None):
spark_session = start()
return SparkNLPReader(spark_session, params)
[docs]def version():
"""Returns the current Spark NLP version.
Returns
-------
str
The current Spark NLP version.
"""
return __version__