File size: 3,624 Bytes
a7616e8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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-large-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) |