PRE_ PYTHONFIY REFACTOR
@@ -1194,7 +1194,7 @@ class NLUPipeline(BasePipe):
print(f'Stored model in {path}')
# else : print('Please fit untrained pipeline first or predict on a String to save it')
def predict(self, data, output_level='', positions=False, keep_stranger_features=True, metadata=False,
- multithread=True, drop_irrelevant_cols=True, verbose=False):
+ multithread=True, drop_irrelevant_cols=True, verbose=False, return_spark_df = False):
'''
Annotates a Pandas Dataframe/Pandas Series/Numpy Array/Spark DataFrame/Python List strings /Python String
@@ -1205,6 +1205,7 @@ class NLUPipeline(BasePipe):
:param metadata: wether to keep additonal metadata in final df or not like confidiences of every possible class for preidctions.
:param multithread: Whether to use multithreading based lightpipeline. In some cases, this may cause errors.
:param drop_irellevant_cols: Wether to drop cols of different output levels, i.e. when predicting token level and dro_irrelevant_cols = True then chunk, sentence and Doc will be dropped
+ :param return_spark_df: Prediction results will be returned right after transforming with the Spark NLP pipeline
:return:
'''
@@ -1373,6 +1374,7 @@ class NLUPipeline(BasePipe):
print(
"If you use Modin, make sure you have installed 'pip install modin[ray]' or 'pip install modin[dask]' backend for Modin ")
+ if return_spark_df : return sdf # Returns RAW result of pipe prediction
return self.pythonify_spark_dataframe(sdf, self.output_different_levels,
keep_stranger_features=keep_stranger_features,
stranger_features=stranger_features, output_metadata=metadata,
