langtest.modelhandler.jsl_modelhandler.PretrainedModelForTextClassification#

class PretrainedModelForTextClassification(model: NLUPipeline | PretrainedPipeline | LightPipeline | PipelineModel)#

Bases: PretrainedJSLModel, ModelAPI

Pretrained model for text classification tasks

__init__(model: NLUPipeline | PretrainedPipeline | LightPipeline | PipelineModel)#

Constructor class

Parameters:

model (LightPipeline) – Loaded SparkNLP LightPipeline for inference.

Methods

__init__(model)

Constructor class

is_classifier(model_instance)

Check classifier model instance is supported by langtest

load_model(path)

Load the NER model into the model attribute.

predict(text[, return_all_scores])

Perform predictions with SparkNLP LightPipeline on the input text.

predict_raw(text)

Perform predictions on the input text.

Attributes

hub

model_registry

task

static is_classifier(model_instance) bool#

Check classifier model instance is supported by langtest

classmethod load_model(path) NLUPipeline#

Load the NER model into the model attribute.

Parameters:

path (str) – Path to pretrained local or NLP Models Hub SparkNLP model

predict(text: str, return_all_scores: bool = False, *args, **kwargs) SequenceClassificationOutput#

Perform predictions with SparkNLP LightPipeline on the input text.

Parameters:
  • text (str) – Input text to perform NER on.

  • return_all_scores (bool) – Option to return score for all labels.

Returns:

Classification output from SparkNLP LightPipeline.

Return type:

SequenceClassificationOutput

predict_raw(text: str) List[str]#

Perform predictions on the input text.

Parameters:

text (str) – Input text to perform text classification on.

Returns:

Predictions as a list of strings.

Return type:

List[str]