langtest.modelhandler.jsl_modelhandler.PretrainedModelForNER#

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

Bases: PretrainedJSLModel, ModelAPI

Pretrained model for NER tasks.

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

Constructor method

Parameters:

model (LightPipeline) – Loaded SparkNLP LightPipeline for inference.

Methods

__init__(model)

Constructor method

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 langtest

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.

Attributes

hub

model_registry

task

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 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, *args, **kwargs) NEROutput#

Perform predictions with SparkNLP LightPipeline on the input text.

Parameters:

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

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.

Parameters:

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

Returns:

Predicted labels.

Return type:

List[str]