space updated

app.py

@@ -3,23 +3,29 @@ import torch
 from transformers import T5Tokenizer, MT5ForConditionalGeneration
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-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, mode, max_length, num_beams):
-    text = f"translate English to Lithuanian: {text}"
-    encoded_input = TOKENIZER(text, return_tensors="pt", padding=True, truncation=True, max_length=max_length).to(DEVICE)
+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())}")
+    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(
+        output_tokens = model.generate(
           **encoded_input,
-          max_length=max_length,
-          num_beams=num_beams,
+          max_length=128,
+          num_beams=5,
           no_repeat_ngram_size=2,
           early_stopping=True
       )
+    return  TOKENIZER.decode(output_tokens[0], skip_special_tokens=True)
 
-    return TOKENIZER.decode(output_tokens[0], skip_special_tokens=True)
 
 
 with gr.Blocks() as interface:
@@ -30,18 +36,26 @@ with gr.Blocks() as interface:
   with gr.Row():
     input_text = gr.Textbox(label="Text input", placeholder="Enter your text here")
     with gr.Column():
-      mode = gr.Dropdown(label="Mode", choices=["Currenly wroking only in en-lt", "Currenly wroking only in en-lt"])
+      mode = gr.Dropdown(label="Mode", choices=["en-lt", "lt-en"])
       translate_button = gr.Button("Translate")
   output_text = gr.Textbox(label="Translated text")
   with gr.Accordion("How to run the model locally:", open=False):
     gr.Code("""import torch
-from transformers import T5Tokenizer, MT5ForConditionalGeneration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer = T5Tokenizer.from_pretrained('google/mt5-small')
+
+from transformers import T5Tokenizer, MT5ForConditionalGeneration
+
+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(
@@ -51,6 +65,7 @@ def translate(text, model, tokenizer, device):
           no_repeat_ngram_size=2,
           early_stopping=True
       )
+
     translated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
     return translated_text
 text = "I live in Kaunas"