Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2
3def translate(text, src_lang='spa_Latn', tgt_lang='awa_Latn', max_input_length=1024, num_beams=4, **kwargs):
4 tokenizer.src_lang = src_lang
5 tokenizer.tgt_lang = src_lang
6 inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=max_input_length)
7 model.eval()
8 result = model.generate(
9 **inputs,
10 forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
11 num_beams=num_beams,
12 **kwargs
13 )
14 return tokenizer.batch_decode(result, skip_special_tokens=True)
15
16model_name = "hectordiazgomez/nllb-spa-awa-v2"
17tokenizer = AutoTokenizer.from_pretrained(model_name)
18model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
19
20text = 'Hola a todos mis amigos' #Hello to all my friends
21translated_text = translate(text, 'spa_Latn', 'awa_Latn')
22print(translated_text)