Views
No views yet
1!pip install sentencepiece transformers==4.33
2import torch
3from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
4
5def fix_tokenizer(tokenizer, new_lang='agr_Latn'):
6 old_len = len(tokenizer) - int(new_lang in tokenizer.added_tokens_encoder)
7 tokenizer.lang_code_to_id[new_lang] = old_len-1
8 tokenizer.id_to_lang_code[old_len-1] = new_lang
9 tokenizer.fairseq_tokens_to_ids["<mask>"] = len(tokenizer.sp_model) + len(tokenizer.lang_code_to_id) + tokenizer.fairseq_offset
10
11 tokenizer.fairseq_tokens_to_ids.update(tokenizer.lang_code_to_id)
12 tokenizer.fairseq_ids_to_tokens = {v: k for k, v in tokenizer.fairseq_tokens_to_ids.items()}
13 if new_lang not in tokenizer._additional_special_tokens:
14 tokenizer._additional_special_tokens.append(new_lang)
15 tokenizer.added_tokens_encoder = {}
16 tokenizer.added_tokens_decoder = {}
17
18MODEL_URL = "hectordiazgomez/nllb-spa-awa-v3"
19model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_URL)
20tokenizer = NllbTokenizer.from_pretrained(MODEL_URL)
21fix_tokenizer(tokenizer)
22
23def translate(
24 text,
25 model,
26 tokenizer,
27 src_lang='agr_Latn',
28 tgt_lang='spa_Latn',
29 max_length='auto',
30 num_beams=4,
31 n_out=None,
32 **kwargs
33):
34 tokenizer.src_lang = src_lang
35 encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
36 if max_length == 'auto':
37 max_length = int(32 + 2.0 * encoded.input_ids.shape[1])
38 model.eval()
39 generated_tokens = model.generate(
40 **encoded.to(model.device),
41 forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang],
42 max_length=max_length,
43 num_beams=num_beams,
44 num_return_sequences=n_out or 1,
45 **kwargs
46 )
47 out = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
48 if isinstance(text, str) and n_out is None:
49 return out[0]
50 return
51
52translate("Uchi piipichi buuke baejai.", model=model, tokenizer=tokenizer)
53# El niño se quedo con el pelo.