| from typing import List | |
| from transformers import BartTokenizer | |
| from idiomify.builders import SourcesBuilder | |
| from idiomify.models import Idiomifier | |
| class Pipeline: | |
| def __init__(self, model: Idiomifier, tokenizer: BartTokenizer): | |
| self.model = model | |
| self.builder = SourcesBuilder(tokenizer) | |
| def __call__(self, sents: List[str], max_length=100) -> List[str]: | |
| srcs = self.builder(literal2idiomatic=[(sent, "") for sent in sents]) | |
| pred_ids = self.model.bart.generate( | |
| inputs=srcs[:, 0], # (N, 2, L) -> (N, L) | |
| attention_mask=srcs[:, 1], # (N, 2, L) -> (N, L) | |
| decoder_start_token_id=self.model.hparams['bos_token_id'], | |
| max_length=max_length, | |
| ) # -> (N, L_t) | |
| tgts = self.builder.tokenizer.batch_decode(pred_ids, skip_special_tokens=True) | |
| return tgts | |