Views
No views yet
1import torch
2device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
3
4from transformers import T5Tokenizer, MT5ForConditionalGeneration
5
6tokenizer = T5Tokenizer.from_pretrained('werent4/mt5TranslatorLT')
7model = MT5ForConditionalGeneration.from_pretrained("werent4/mt5TranslatorLT")
8model.to(device)
9def translate(text, model, tokenizer, device, translation_way = "en-lt"):
10 translations_ways = {
11 "en-lt": "<EN2LT>",
12 "lt-en": "<LT2EN>"
13 }
14 if translation_way not in translations_ways:
15 raise ValueError(f"Invalid translation way. Supported ways: {list(translations_ways.keys())}")
16 input_text = f"{translations_ways[translation_way]} {text}"
17 encoded_input = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
18 with torch.no_grad():
19 output_tokens = model.generate(
20 **encoded_input,
21 max_length=128,
22 num_beams=5,
23 no_repeat_ngram_size=2,
24 early_stopping=True
25 )
26
27 translated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
28 return translated_text
29
30text = "How are you?"
31translate(text, model, tokenizer, device)
32`Kaip esate?`
33
34text = "I live in Kaunas"
35translate(text, model, tokenizer, device)
36`Aš gyvenu Kaunas`
37
38text = "Mano vardas yra Karolis"
39translate(text, model, tokenizer, device, translation_way= "lt-en")
40`My name is Karolis`