Views
No views yet

translate.py file:1import torch
2import transformers
3from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
4from transformers.generation import LogitsProcessor
5
6
7class RepetitionPenaltyLogitsProcessor(LogitsProcessor):
8 def __init__(self, penalty: float, model):
9 last_bias = model.classifier.nonlinearity[-1].bias.data
10 last_bias = torch.nn.functional.log_softmax(last_bias)
11 self.penalty = penalty * (last_bias - last_bias.max())
12
13 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
14 penalized_score = torch.gather(scores + self.penalty.unsqueeze(0).to(input_ids.device), 1, input_ids).to(scores.dtype)
15 scores.scatter_(1, input_ids, penalized_score)
16 return scores
17
18
19class Translator:
20 def __init__(self, model_path="ltg/nort5-base-en-no-translation", device="cpu"):
21 self.tokenizer = AutoTokenizer.from_pretrained(model_path)
22 self.cls_index = self.tokenizer.convert_tokens_to_ids("[CLS]")
23 self.sep_index = self.tokenizer.convert_tokens_to_ids("[SEP]")
24 self.eos_index = self.tokenizer.convert_tokens_to_ids("[EOS]")
25 self.pad_index = self.tokenizer.convert_tokens_to_ids("[PAD]")
26 self.eng_index = self.tokenizer.convert_tokens_to_ids(">>eng<<")
27 self.nob_index = self.tokenizer.convert_tokens_to_ids(">>nob<<")
28 self.nno_index = self.tokenizer.convert_tokens_to_ids(">>nno<<")
29
30 self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path, trust_remote_code=True)
31
32 self.device = device
33 print(f"SYSTEM: Running on {self.device}", flush=True)
34
35 self.model = self.model.to(device)
36 self.model.eval()
37
38 print(f"Sucessfully loaded the model to the memory")
39
40 self.LANGUAGE_IDS = {
41 "en": self.eng_index,
42 "nb": self.nob_index,
43 "nn": self.nno_index
44 }
45
46 def __call__(self, source, source_language, target_language):
47 source = [s.strip() for s in source.split('\n')]
48 source_subwords = self.tokenizer(source).input_ids
49 source_subwords = [[self.cls_index, self.LANGUAGE_IDS[target_language], self.LANGUAGE_IDS[source_language]] + s + [self.sep_index] for s in source_subwords]
50 source_subwords = [torch.tensor(s) for s in source_subwords]
51 source_subwords = torch.nn.utils.rnn.pad_sequence(source_subwords, batch_first=True, padding_value=self.pad_index)
52 source_subwords = source_subwords[:, :512].to(self.device)
53
54 def generate(model, **kwargs):
55 with torch.inference_mode():
56 with torch.autocast(enabled=self.device != "cpu", device_type="cuda", dtype=torch.bfloat16):
57 return model.generate(**kwargs)
58
59 generate_kwargs = dict(
60 input_ids=source_subwords,
61 attention_mask=(source_subwords != self.pad_index).long(),
62 max_new_tokens = 512-1,
63 num_beams=8,
64 length_penalty=1.6,
65 early_stopping=True,
66 do_sample=False,
67 use_cache=True,
68 logits_processor=[RepetitionPenaltyLogitsProcessor(0.5, self.model), transformers.LogitNormalization()]
69 )
70 output = generate(self.model, **generate_kwargs).tolist()
71 paragraphs = [self.tokenizer.decode(c, skip_special_tokens=True).strip() for c in output]
72 translation = '\n'.join(paragraphs)
73
74 return translation
75
76
77if __name__ == "__main__":
78
79 translator = Translator()
80
81 en_text = "How are you feeling right now? Better?"
82 no_text = translator(en_text, "en", "nb")
83
84 print(en_text)
85 print(no_text)1@inproceedings{samuel-etal-2023-norbench,
2 title = "{N}or{B}ench {--} A Benchmark for {N}orwegian Language Models",
3 author = "Samuel, David and
4 Kutuzov, Andrey and
5 Touileb, Samia and
6 Velldal, Erik and
7 {\O}vrelid, Lilja and
8 R{\o}nningstad, Egil and
9 Sigdel, Elina and
10 Palatkina, Anna",
11 booktitle = "Proceedings of the 24th Nordic Conference on Computational Linguistics (NoDaLiDa)",
12 month = may,
13 year = "2023",
14 address = "T{\'o}rshavn, Faroe Islands",
15 publisher = "University of Tartu Library",
16 url = "https://aclanthology.org/2023.nodalida-1.61",
17 pages = "618--633",
18 abstract = "We present NorBench: a streamlined suite of NLP tasks and probes for evaluating Norwegian language models (LMs) on standardized data splits and evaluation metrics. We also introduce a range of new Norwegian language models (both encoder and encoder-decoder based). Finally, we compare and analyze their performance, along with other existing LMs, across the different benchmark tests of NorBench.",
19}
20