langtest.modelhandler.transformers_modelhandler.PretrainedModelForNER#

class PretrainedModelForNER(model, **kwargs)#

Bases: ModelAPI

Transformers pretrained model for NER tasks

Parameters:

model (transformers.pipeline.Pipeline) – Pretrained HuggingFace NER pipeline for predictions.

__init__(model, **kwargs)#

Constructor method

Parameters:

model (transformers.pipeline.Pipeline) – Pretrained HuggingFace NER pipeline for predictions.

Methods

__init__(model, **kwargs)	Constructor method
group_entities(entities)	Find and group together the adjacent tokens with the same entity predicted.
input_process(input_sen)	Process the input text to be used in the model
load_model(path, **kwargs)	Load the NER model into the model attribute.
predict(text, **kwargs)	Perform predictions on the input text.
predict_raw(text)	Predict a list of labels.

Attributes

model_registry

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]

input_process(input_sen)#

Process the input text to be used in the model

classmethod load_model(path: str | Any, **kwargs) → Pipeline#

Load the NER model into the model attribute.

Parameters:

path (str) – path to model or model name

Return type:

‘Pipeline’

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

Perform predictions on the input text.

Parameters:

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

• kwargs – Additional keyword arguments.

Keyword Arguments:

group_entities (bool) – Option to group entities.

Returns:

A list of named entities recognized in the input text.

Return type:

NEROutput

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

Predict a list of labels.

Parameters:

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

Returns:

A list of named entities recognized in the input text.

Return type:

List[str]