nlptest.modelhandler.jsl_modelhandler.PretrainedModelForNER#

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

Bases: _ModelHandler

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

model#

Loaded SparkNLP LightPipeline for inference.

Type:

LightPipeline

Methods

__init__(model)	model#
group_entities(entities)	Find and group together the adjacent tokens with the same entity predicted.
is_ner_annotator(model_instance)	Check ner model instance is supported by nlptest
load_model(path)	Load the NER model into the model attribute.
predict(text, *args, **kwargs)	Perform predictions with SparkNLP LightPipeline on the input text.
predict_raw(text)	Perform predictions with SparkNLP LightPipeline on the input text.

group_entities(entities: List[Dict]) → List[Dict]#

Find and group together the adjacent tokens with the same entity predicted. Inspired and adapted from: huggingface/transformers

Parameters:

entities (List[Dict]) – The entities predicted by the pipeline.

Returns:

grouped entities

Return type:

List[Dict]

static is_ner_annotator(model_instance) → bool#

Check ner model instance is supported by nlptest

classmethod load_model(path: str) → NLUPipeline#

Load the NER model into the model attribute. :param path: Path to pretrained local or NLP Models Hub SparkNLP model :type path: str

predict(text: str, *args, **kwargs) → NEROutput#

Perform predictions with SparkNLP LightPipeline on the input text. :param text: Input text to perform NER on. :type text: str

Returns:

A list of named entities recognized in the input text.

Return type:

NEROutput

predict_raw(text: str) → List[str]#

Perform predictions with SparkNLP LightPipeline on the input text. :param text: Input text to perform NER on. :type text: str

Returns:

Predicted labels.

Return type:

List[str]