Views
No views yet
1# the version of transformers is important!
2!pip install sentencepiece transformers==4.33
3import torch
4from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
5
6def fix_tokenizer(tokenizer, new_lang='ain_Latn'):
7 """ Add a new language token to the tokenizer vocabulary (this should be done each time after its initialization) """
8 old_len = len(tokenizer) - int(new_lang in tokenizer.added_tokens_encoder)
9 tokenizer.lang_code_to_id[new_lang] = old_len-1
10 tokenizer.id_to_lang_code[old_len-1] = new_lang
11 # always move "mask" to the last position
12 tokenizer.fairseq_tokens_to_ids["<mask>"] = len(tokenizer.sp_model) + len(tokenizer.lang_code_to_id) + tokenizer.fairseq_offset
13
14 tokenizer.fairseq_tokens_to_ids.update(tokenizer.lang_code_to_id)
15 tokenizer.fairseq_ids_to_tokens = {v: k for k, v in tokenizer.fairseq_tokens_to_ids.items()}
16 if new_lang not in tokenizer._additional_special_tokens:
17 tokenizer._additional_special_tokens.append(new_lang)
18 # clear the added token encoder; otherwise a new token may end up there by mistake
19 tokenizer.added_tokens_encoder = {}
20 tokenizer.added_tokens_decoder = {}
21
22MODEL_URL = "TwentyNine/nllb-jpn-ain-v1"
23model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_URL)
24tokenizer = NllbTokenizer.from_pretrained(MODEL_URL)
25fix_tokenizer(tokenizer)
26
27def translate(
28 text,
29 model,
30 tokenizer,
31 src_lang='jpn_Jpan',
32 tgt_lang='ain_Latn',
33 max_length='auto',
34 num_beams=4,
35 n_out=None,
36 **kwargs
37):
38 tokenizer.src_lang = src_lang
39 encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
40 if max_length == 'auto':
41 max_length = int(32 + 2.0 * encoded.input_ids.shape[1])
42 model.eval()
43 generated_tokens = model.generate(
44 **encoded.to(model.device),
45 forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
46 max_length=max_length,
47 num_beams=num_beams,
48 num_return_sequences=n_out or 1,
49 **kwargs
50 )
51 out = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
52 if isinstance(text, str) and n_out is None:
53 return out[0]
54 return
55
56translate("肉が食べたいな。", model=model, tokenizer=tokenizer)
57# 'kam c=e rusuy na.'