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