Views
No views yet
1from transformers import AutoModelForSeq2SeqLM, NllbTokenizer
2
3model = AutoModelForSeq2SeqLM.from_pretrained("leks-forever/nllb-200-distilled-600M")
4tokenizer = NllbTokenizer.from_pretrained("leks-forever/nllb-200-distilled-600M")
5
6def predict(
7 text,
8 src_lang='lez_Cyrl',
9 tgt_lang='rus_Cyrl',
10 a=32, b=3,
11 max_input_length=1024,
12 num_beams=1,
13 **kwargs
14):
15 tokenizer.src_lang = src_lang
16 tokenizer.tgt_lang = tgt_lang
17 inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=max_input_length)
18 result = model.generate(
19 **inputs.to(model.device),
20 forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
21 max_new_tokens=int(a + b * inputs.input_ids.shape[1]),
22 num_beams=num_beams,
23 **kwargs
24 )
25 return tokenizer.batch_decode(result, skip_special_tokens=True)
26
27sentence: str = "Я люблю гулять по парку ранним утром, когда воздух свежий и тишина вокруг."
28
29translation = predict(sentence, src_lang='rus_Cyrl', tgt_lang='lez_Cyrl')
30
31print(translation)
32
33# ['Заз пакамахъ, хъсан гар алаз, сагъ-саламатдиз къекъвез кӀанзава.'
34