| from typing import Dict, Any | |
| from transformers import AutoModelForSeq2SeqLM, NllbTokenizer | |
| from translation import fix_tokenizer, TextPreprocessor, sentenize_with_fillers | |
| from sentence_splitter import SentenceSplitter | |
| import torch | |
| class PreTrainedPipeline(): | |
| def __init__(self, path=""): | |
| self.model = AutoModelForSeq2SeqLM.from_pretrained(path) | |
| if torch.cuda.is_available(): | |
| self.model = self.model.cuda() | |
| self.tokenizer = NllbTokenizer.from_pretrained(path) | |
| fix_tokenizer(self.tokenizer) | |
| self.splitter = SentenceSplitter(language='es') | |
| self.preprocessor = TextPreprocessor() | |
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: | |
| inputs = data.get("text", "") | |
| src_lang = data.get("src_lang", "spa_Latn") | |
| tgt_lang = data.get("tgt_lang", "agr_Latn") | |
| preprocess = data.get("preprocess", True) | |
| sentences, fillers = sentenize_with_fillers(inputs, self.splitter) | |
| if preprocess: | |
| sentences = [self.preprocessor(sent) for sent in sentences] | |
| translated_sentences = [] | |
| for sentence in sentences: | |
| self.tokenizer.src_lang = src_lang | |
| encoded = self.tokenizer(sentence, return_tensors="pt") | |
| generated_tokens = self.model.generate( | |
| **encoded.to(self.model.device), | |
| forced_bos_token_id=self.tokenizer.lang_code_to_id[tgt_lang] | |
| ) | |
| translated_sentences.append( | |
| self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0] | |
| ) | |
| output = "".join( | |
| filler + sentence for filler, sentence in zip(fillers, translated_sentences) | |
| ) + fillers[-1] | |
| return {"translation": output} | |