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)