langtest.modelhandler.jsl_modelhandler.PretrainedModelForNER#
- class PretrainedModelForNER(model: NLUPipeline | PretrainedPipeline | LightPipeline | PipelineModel)#
Bases:
PretrainedJSLModel,ModelAPIPretrained 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
hubmodel_registrytask- 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]