| from typing import Dict, List, Any | |
| from nemo.collections.nlp.models import PunctuationCapitalizationModel | |
| class PreTrainedPipeline: | |
| def __init__(self, path=""): | |
| # IMPLEMENT_THIS | |
| # Preload all the elements you are going to need at inference. | |
| # For instance your model, processors, tokenizer that might be needed. | |
| # This function is only called once, so do all the heavy processing I/O here""" | |
| self.model = PunctuationCapitalizationModel.from_pretrained("dchaplinsky/punctuation_uk_bert") | |
| def __call__(self, inputs: str) -> List[Dict]: | |
| """ | |
| Args: | |
| inputs (:obj:`str`): | |
| a string containing some text | |
| Return: | |
| A :obj:`str` | |
| """ | |
| inputs = inputs.strip() | |
| return [{"generated_text": self.model.add_punctuation_capitalization([inputs])[0]}] | |