Update README.md

README.md

@@ -21,12 +21,17 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 from transformers import T5Tokenizer, MT5ForConditionalGeneration
 
-tokenizer = T5Tokenizer.from_pretrained('google/mt5-small')
+tokenizer = T5Tokenizer.from_pretrained('werent4/mt5TranslatorLT')
 model = MT5ForConditionalGeneration.from_pretrained("werent4/mt5TranslatorLT")
 model.to(device)
-
-def translate(text, model, tokenizer, device):
-    input_text = f"translate English to Lithuanian: {text}"
+def translate(text, model, tokenizer, device, translation_way = "en-lt"):
+    translations_ways = {
+        "en-lt": "<EN2LT>",
+        "lt-en": "<LT2EN>"
+    }
+    if translation_way not in translations_ways:
+        raise ValueError(f"Invalid translation way. Supported ways: {list(translations_ways.keys())}")
+    input_text = f"{translations_ways[translation_way]} {text}"
     encoded_input = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
     with torch.no_grad():
         output_tokens = model.generate(
@@ -40,10 +45,6 @@ def translate(text, model, tokenizer, device):
     translated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
     return translated_text
 
-text = "women"
-translate(text, model, tokenizer, device)
-`moteris`
-
 text = "How are you?"
 translate(text, model, tokenizer, device)
 `Kaip esate?`
@@ -51,6 +52,10 @@ translate(text, model, tokenizer, device)
 text = "I live in Kaunas"
 translate(text, model, tokenizer, device)
 `Aš gyvenu Kaunas`
+
+text = "Mano vardas yra Karolis"
+translate(text, model, tokenizer, device, translation_way= "lt-en")
+`My name is Karolis`
 ```