|
from typing import Dict, List, Any |
|
from nemo.collections.nlp.models import PunctuationCapitalizationModel |
|
|
|
class PreTrainedPipeline(): |
|
def __init__(self, path=""): |
|
|
|
|
|
|
|
|
|
self.model = PunctuationCapitalizationModel.from_pretrained("dchaplinsky/punctuation_uk_bert") |
|
|
|
def __call__(self, inputs: str) -> List[Dict[str, Any]]: |
|
""" |
|
Args: |
|
inputs (:obj:`str`): |
|
a string containing some text |
|
Return: |
|
A :obj:`list`:. The object returned should be like [{"entity_group": "XXX", "word": "some word", "start": 3, "end": 6, "score": 0.82}] containing : |
|
- "entity_group": A string representing what the entity is. |
|
- "word": A substring of the original string that was detected as an entity. |
|
- "start": the offset within `input` leading to `answer`. context[start:stop] == word |
|
- "end": the ending offset within `input` leading to `answer`. context[start:stop] === word |
|
- "score": A score between 0 and 1 describing how confident the model is for this entity. |
|
""" |
|
inputs = inputs.strip() |
|
labels = self.model.add_punctuation_capitalization([inputs], return_labels=True)[0].split() |
|
|
|
tokens = inputs.split() |
|
|
|
res: List[Dict[str, Any]] = [] |
|
offset = 0 |
|
for tok, lab in zip(tokens, labels): |
|
if lab != "OO": |
|
res.append({ |
|
"entity_group": lab, |
|
"word": tok, |
|
"start": offset, |
|
"end": offset + len(tok), |
|
"score": 1 |
|
}) |
|
|
|
offset += len(tok) + 1 |
|
|
|
return res |
