Views
No views yet
1from transformers import MBartForConditionalGeneration, MBart50Tokenizer
2model = MBartForConditionalGeneration.from_pretrained("SaranaAbidueva/mbart50_ru_bua")
3tokenizer = MBart50Tokenizer.from_pretrained("SaranaAbidueva/mbart50_ru_bua")
4
5
6def translate(text, max_length=200, num_beams=5, repetition_penalty=5.0, **kwargs):
7 encoded = tokenizer(text, return_tensors="pt")
8 generated_tokens = model.generate(
9 **encoded.to(model.device),
10 max_length=max_length,
11 num_beams=num_beams,
12 repetition_penalty=repetition_penalty
13 )
14 return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
15
16translate('Евгений Онегин интересная книга')