| import torch | |
| import transformers | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| from transformers.generation import LogitsProcessor | |
| class RepetitionPenaltyLogitsProcessor(LogitsProcessor): | |
| def __init__(self, penalty: float, model): | |
| last_bias = model.classifier.nonlinearity[-1].bias.data | |
| last_bias = torch.nn.functional.log_softmax(last_bias) | |
| self.penalty = penalty * (last_bias - last_bias.max()) | |
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: | |
| penalized_score = torch.gather(scores + self.penalty.unsqueeze(0).to(input_ids.device), 1, input_ids).to(scores.dtype) | |
| scores.scatter_(1, input_ids, penalized_score) | |
| return scores | |
| class Translator: | |
| def __init__(self, model_path="ltg/nort5-base-en-no-translation", device="cpu"): | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| self.cls_index = self.tokenizer.convert_tokens_to_ids("[CLS]") | |
| self.sep_index = self.tokenizer.convert_tokens_to_ids("[SEP]") | |
| self.eos_index = self.tokenizer.convert_tokens_to_ids("[EOS]") | |
| self.pad_index = self.tokenizer.convert_tokens_to_ids("[PAD]") | |
| self.eng_index = self.tokenizer.convert_tokens_to_ids(">>eng<<") | |
| self.nob_index = self.tokenizer.convert_tokens_to_ids(">>nob<<") | |
| self.nno_index = self.tokenizer.convert_tokens_to_ids(">>nno<<") | |
| self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path, trust_remote_code=True) | |
| self.device = device | |
| print(f"SYSTEM: Running on {self.device}", flush=True) | |
| self.model = self.model.to(device) | |
| self.model.eval() | |
| print(f"Sucessfully loaded the model to the memory") | |
| self.LANGUAGE_IDS = { | |
| "en": self.eng_index, | |
| "nb": self.nob_index, | |
| "nn": self.nno_index | |
| } | |
| def __call__(self, source, source_language, target_language): | |
| source = [s.strip() for s in source.split('\n')] | |
| source_subwords = self.tokenizer(source).input_ids | |
| source_subwords = [[self.cls_index, self.LANGUAGE_IDS[target_language], self.LANGUAGE_IDS[source_language]] + s + [self.sep_index] for s in source_subwords] | |
| source_subwords = [torch.tensor(s) for s in source_subwords] | |
| source_subwords = torch.nn.utils.rnn.pad_sequence(source_subwords, batch_first=True, padding_value=self.pad_index) | |
| source_subwords = source_subwords[:, :512].to(self.device) | |
| def generate(model, **kwargs): | |
| with torch.inference_mode(): | |
| with torch.autocast(enabled=self.device != "cpu", device_type="cuda", dtype=torch.bfloat16): | |
| return model.generate(**kwargs) | |
| generate_kwargs = dict( | |
| input_ids=source_subwords, | |
| attention_mask=(source_subwords != self.pad_index).long(), | |
| max_new_tokens = 512-1, | |
| num_beams=8, | |
| length_penalty=1.6, | |
| early_stopping=True, | |
| do_sample=False, | |
| use_cache=True, | |
| logits_processor=[RepetitionPenaltyLogitsProcessor(0.5, self.model), transformers.LogitNormalization()] | |
| ) | |
| output = generate(self.model, **generate_kwargs).tolist() | |
| paragraphs = [self.tokenizer.decode(c, skip_special_tokens=True).strip() for c in output] | |
| translation = '\n'.join(paragraphs) | |
| return translation | |
| if __name__ == "__main__": | |
| translator = Translator() | |
| en_text = "How are you feeling right now? Better?" | |
| no_text = translator(en_text, "en", "nb") | |
| print(en_text) | |
| print(no_text) |