[spark] Format python

extensions/spark/setup/djl_spark/task/audio/__init__.py

@@ -10,7 +10,6 @@
 # or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
-
 """DJL Spark Tasks Audio API."""
 
 from . import (

extensions/spark/setup/djl_spark/task/audio/speech_recognizer.py

@@ -64,18 +64,18 @@ def recognize(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        recognizer = (
-            sc._jvm.ai.djl.spark.task.audio.SpeechRecognizer()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        recognizer = sc._jvm.ai.djl.spark.task.audio.SpeechRecognizer() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.engine is not None:
             recognizer = recognizer.setEngine(self.engine)
         if self.batch_size is not None:
             recognizer = recognizer.setBatchSize(self.batch_size)
         if self.translator_factory is not None:
-            recognizer = recognizer.setTranslatorFactory(self.translator_factory)
+            recognizer = recognizer.setTranslatorFactory(
+                self.translator_factory)
         if self.batchifier is not None:
             recognizer = recognizer.setBatchifier(self.batchifier)
-        return DataFrame(recognizer.recognize(dataset._jdf), dataset.sparkSession)
+        return DataFrame(recognizer.recognize(dataset._jdf),
+                         dataset.sparkSession)

extensions/spark/setup/djl_spark/task/audio/whisper_speech_recognizer.py

@@ -20,7 +20,6 @@
 from transformers import pipeline
 from ...util import files_util, dependency_util
 
-
 TASK = "automatic-speech-recognition"
 APPLICATION = "audio/automatic_speech_recognition"
 GROUP_ID = "ai/djl/huggingface/pytorch"
@@ -65,22 +64,31 @@ def recognize(self, dataset, generate_kwargs=None, **kwargs):
             raise ValueError("Only PyTorch engine is supported.")
 
         if self.model_url:
-            cache_dir = files_util.get_cache_dir(APPLICATION, GROUP_ID, self.model_url)
+            cache_dir = files_util.get_cache_dir(APPLICATION, GROUP_ID,
+                                                 self.model_url)
             files_util.download_and_extract(self.model_url, cache_dir)
             dependency_util.install(cache_dir)
             model_id_or_path = cache_dir
         elif self.hf_model_id:
             model_id_or_path = self.hf_model_id
         else:
-            raise ValueError("Either model_url or hf_model_id must be provided.")
+            raise ValueError(
+                "Either model_url or hf_model_id must be provided.")
 
         @pandas_udf(StringType())
         def predict_udf(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
-            pipe = pipeline(TASK, generate_kwargs=generate_kwargs, model=model_id_or_path,
-                            batch_size=self.batch_size, chunk_length_s=30, **kwargs)
+            pipe = pipeline(TASK,
+                            generate_kwargs=generate_kwargs,
+                            model=model_id_or_path,
+                            batch_size=self.batch_size,
+                            chunk_length_s=30,
+                            **kwargs)
             for s in iterator:
                 # Model expects single channel, 16000 sample rate audio
-                batch = [librosa.load(io.BytesIO(d), mono=True, sr=16000)[0] for d in s]
+                batch = [
+                    librosa.load(io.BytesIO(d), mono=True, sr=16000)[0]
+                    for d in s
+                ]
                 output = pipe(batch)
                 text = [o["text"] for o in output]
                 yield pd.Series(text)

extensions/spark/setup/djl_spark/task/binary/__init__.py

@@ -10,7 +10,6 @@
 # or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
-
 """DJL Spark Tasks Binary API."""
 
 from . import binary_predictor

extensions/spark/setup/djl_spark/task/binary/binary_predictor.py

@@ -63,12 +63,10 @@ def predict(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        predictor = (
-            sc._jvm.ai.djl.spark.task.binary.BinaryPredictor()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        predictor = sc._jvm.ai.djl.spark.task.binary.BinaryPredictor() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.engine is not None:
             predictor = predictor.setEngine(self.engine)
         if self.batch_size is not None:

extensions/spark/setup/djl_spark/task/text/__init__.py

@@ -10,7 +10,6 @@
 # or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
-
 """DJL Spark Tasks Text API."""
 
 from . import (

extensions/spark/setup/djl_spark/task/text/question_answerer.py

@@ -55,11 +55,9 @@ def answer(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        answerer = (
-            sc._jvm.ai.djl.spark.task.text.QuestionAnswerer()
-            .setOutputCol(self.output_col)
+        answerer = sc._jvm.ai.djl.spark.task.text.QuestionAnswerer() \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.input_cols is not None:
             # Convert the input_cols to Java array
             input_cols_arr = sc._gateway.new_array(sc._jvm.java.lang.String,

extensions/spark/setup/djl_spark/task/text/text2text_generator.py

@@ -60,18 +60,23 @@ def generate(self, dataset, **kwargs):
             raise ValueError("Only PyTorch engine is supported.")
 
         if self.model_url:
-            cache_dir = files_util.get_cache_dir(APPLICATION, GROUP_ID, self.model_url)
+            cache_dir = files_util.get_cache_dir(APPLICATION, GROUP_ID,
+                                                 self.model_url)
             files_util.download_and_extract(self.model_url, cache_dir)
             dependency_util.install(cache_dir)
             model_id_or_path = cache_dir
         elif self.hf_model_id:
             model_id_or_path = self.hf_model_id
         else:
-            raise ValueError("Either model_url or hf_model_id must be provided.")
+            raise ValueError(
+                "Either model_url or hf_model_id must be provided.")
 
         @pandas_udf(StringType())
         def predict_udf(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
-            pipe = pipeline(TASK, model=model_id_or_path, batch_size=self.batch_size, **kwargs)
+            pipe = pipeline(TASK,
+                            model=model_id_or_path,
+                            batch_size=self.batch_size,
+                            **kwargs)
             for s in iterator:
                 output = pipe(s.tolist())
                 text = [o["generated_text"] for o in output]

extensions/spark/setup/djl_spark/task/text/text_classifier.py

@@ -55,18 +55,18 @@ def classify(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        classifier = (
-            sc._jvm.ai.djl.spark.task.text.TextClassifier()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        classifier = sc._jvm.ai.djl.spark.task.text.TextClassifier() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.engine is not None:
             classifier = classifier.setEngine(self.engine)
         if self.batch_size is not None:
             classifier = classifier.setBatchSize(self.batch_size)
         if self.translator_factory is not None:
-            classifier = classifier.setTranslatorFactory(self.translator_factory)
+            classifier = classifier.setTranslatorFactory(
+                self.translator_factory)
         if self.batchifier is not None:
             classifier = classifier.setBatchifier(self.batchifier)
-        return DataFrame(classifier.classify(dataset._jdf), dataset.sparkSession)
+        return DataFrame(classifier.classify(dataset._jdf),
+                         dataset.sparkSession)

extensions/spark/setup/djl_spark/task/text/text_decoder.py

@@ -44,12 +44,10 @@ def decode(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        decoder = (
-            sc._jvm.ai.djl.spark.task.text.TextDecoder()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        decoder = sc._jvm.ai.djl.spark.task.text.TextDecoder() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setHfModelId(self.hf_model_id)
-        )
         if self.batch_size is not None:
             decoder = decoder.setBatchSize(self.batch_size)
         return DataFrame(decoder.decode(dataset._jdf), dataset.sparkSession)

extensions/spark/setup/djl_spark/task/text/text_embedder.py

@@ -55,12 +55,10 @@ def embed(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        embedder = (
-            sc._jvm.ai.djl.spark.task.text.TextEmbedder()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        embedder = sc._jvm.ai.djl.spark.task.text.TextEmbedder() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.engine is not None:
             embedder = embedder.setEngine(self.engine)
         if self.batch_size is not None:

extensions/spark/setup/djl_spark/task/text/text_encoder.py

@@ -44,12 +44,10 @@ def encode(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        encoder = (
-            sc._jvm.ai.djl.spark.task.text.TextEncoder()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        encoder = sc._jvm.ai.djl.spark.task.text.TextEncoder() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setHfModelId(self.hf_model_id)
-        )
         if self.batch_size is not None:
             encoder = encoder.setBatchSize(self.batch_size)
         return DataFrame(encoder.encode(dataset._jdf), dataset.sparkSession)

extensions/spark/setup/djl_spark/task/text/text_generator.py

@@ -18,7 +18,6 @@
 from transformers import pipeline
 from ...util import files_util, dependency_util
 
-
 TASK = "text-generation"
 APPLICATION = "nlp/text_generation"
 GROUP_ID = "ai/djl/huggingface/pytorch"
@@ -61,18 +60,23 @@ def generate(self, dataset, **kwargs):
             raise ValueError("Only PyTorch engine is supported.")
 
         if self.model_url:
-            cache_dir = files_util.get_cache_dir(APPLICATION, GROUP_ID, self.model_url)
+            cache_dir = files_util.get_cache_dir(APPLICATION, GROUP_ID,
+                                                 self.model_url)
             files_util.download_and_extract(self.model_url, cache_dir)
             dependency_util.install(cache_dir)
             model_id_or_path = cache_dir
         elif self.hf_model_id:
             model_id_or_path = self.hf_model_id
         else:
-            raise ValueError("Either model_url or hf_model_id must be provided.")
+            raise ValueError(
+                "Either model_url or hf_model_id must be provided.")
 
         @pandas_udf(StringType())
         def predict_udf(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
-            pipe = pipeline(TASK, model=model_id_or_path, batch_size=self.batch_size, **kwargs)
+            pipe = pipeline(TASK,
+                            model=model_id_or_path,
+                            batch_size=self.batch_size,
+                            **kwargs)
             for s in iterator:
                 output = pipe(s.tolist())
                 text = [o[0]["generated_text"] for o in output]

extensions/spark/setup/djl_spark/task/text/text_tokenizer.py

@@ -17,10 +17,7 @@
 
 class TextTokenizer:
 
-    def __init__(self,
-                 input_col: str,
-                 output_col: str,
-                 hf_model_id: str):
+    def __init__(self, input_col: str, output_col: str, hf_model_id: str):
         """
         Initializes the TextTokenizer.
 
@@ -40,10 +37,9 @@ def tokenize(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        tokenizer = (
-            sc._jvm.ai.djl.spark.task.text.TextTokenizer()
-            .setInputCol(self.input_col)
-            .setOutputCol(self.output_col)
+        tokenizer = sc._jvm.ai.djl.spark.task.text.TextTokenizer() \
+            .setInputCol(self.input_col) \
+            .setOutputCol(self.output_col) \
             .setHfModelId(self.hf_model_id)
-        )
-        return DataFrame(tokenizer.tokenize(dataset._jdf), dataset.sparkSession)
+        return DataFrame(tokenizer.tokenize(dataset._jdf),
+                         dataset.sparkSession)

extensions/spark/setup/djl_spark/task/vision/__init__.py

@@ -10,7 +10,6 @@
 # or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
-
 """DJL Spark Tasks Vision API."""
 
 from . import (

extensions/spark/setup/djl_spark/task/vision/image_classifier.py

@@ -63,11 +63,9 @@ def classify(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        classifier = (
-            sc._jvm.ai.djl.spark.task.vision.ImageClassifier()
-            .setOutputCol(self.output_col)
+        classifier = sc._jvm.ai.djl.spark.task.vision.ImageClassifier() \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.input_cols is not None:
             # Convert the input_cols to Java array
             input_cols_arr = sc._gateway.new_array(sc._jvm.java.lang.String,
@@ -79,11 +77,13 @@ def classify(self, dataset):
         if self.batch_size is not None:
             classifier = classifier.setBatchSize(self.batch_size)
         if self.translator_factory is not None:
-            classifier = classifier.setTranslatorFactory(self.translator_factory)
+            classifier = classifier.setTranslatorFactory(
+                self.translator_factory)
         if self.batchifier is not None:
             classifier = classifier.setBatchifier(self.batchifier)
         if self.apply_softmax is not None:
             classifier = classifier.setApplySoftmax(self.apply_softmax)
         if self.top_k is not None:
             classifier = classifier.setTopK(self.top_k)
-        return DataFrame(classifier.classify(dataset._jdf), dataset.sparkSession)
+        return DataFrame(classifier.classify(dataset._jdf),
+                         dataset.sparkSession)

extensions/spark/setup/djl_spark/task/vision/image_embedder.py

@@ -57,11 +57,9 @@ def embed(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        embedder = (
-            sc._jvm.ai.djl.spark.task.vision.ImageEmbedder()
-            .setOutputCol(self.output_col)
+        embedder = sc._jvm.ai.djl.spark.task.vision.ImageEmbedder() \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.input_cols is not None:
             # Convert the input_cols to Java array
             input_cols_arr = sc._gateway.new_array(sc._jvm.java.lang.String,

extensions/spark/setup/djl_spark/task/vision/instance_segmenter.py

@@ -59,11 +59,9 @@ def segment(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        segmenter = (
-            sc._jvm.ai.djl.spark.task.vision.InstanceSegmenter()
-            .setOutputCol(self.output_col)
+        segmenter = sc._jvm.ai.djl.spark.task.vision.InstanceSegmenter() \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.input_cols is not None:
             # Convert the input_cols to Java array
             input_cols_arr = sc._gateway.new_array(sc._jvm.java.lang.String,

extensions/spark/setup/djl_spark/task/vision/object_detector.py

@@ -57,11 +57,9 @@ def detect(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        detector = (
-            sc._jvm.ai.djl.spark.task.vision.ObjectDetector()
-            .setOutputCol(self.output_col)
+        detector = sc._jvm.ai.djl.spark.task.vision.ObjectDetector() \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.input_cols is not None:
             # Convert the input_cols to Java array
             input_cols_arr = sc._gateway.new_array(sc._jvm.java.lang.String,

extensions/spark/setup/djl_spark/task/vision/semantic_segmenter.py

@@ -59,11 +59,9 @@ def segment(self, dataset):
         :return: output dataset
         """
         sc = SparkContext._active_spark_context
-        segmenter = (
-            sc._jvm.ai.djl.spark.task.vision.SemanticSegmenter()
-            .setOutputCol(self.output_col)
+        segmenter = sc._jvm.ai.djl.spark.task.vision.SemanticSegmenter() \
+            .setOutputCol(self.output_col) \
             .setModelUrl(self.model_url)
-        )
         if self.input_cols is not None:
             # Convert the input_cols to Java array
             input_cols_arr = sc._gateway.new_array(sc._jvm.java.lang.String,