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:
- 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]