Use a newer model; add randomization; allow same lang

.gitignore

@@ -0,0 +1 @@
+.idea

app.py

@@ -1,13 +1,12 @@
 import gradio as gr
 import torch
 from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
-MODEL_URL = 'slone/nllb-rus-tyv-v1'
-# tokenizer = NllbTokenizer.from_pretrained(MODEL_URL, force_download=True)
+MODEL_URL = 'slone/nllb-rus-tyv-v2-extvoc'
 
 
 lang_to_code = {
-    'Русский | Russian': 'rus_Cyrl',
-    'Тувинский | Tyvan': 'tyv_Cyrl',
+    'Орус | Русский | Russian': 'rus_Cyrl',
+    'Тыва | Тувинский | Tyvan': 'tyv_Cyrl',
 }
 
 
@@ -60,18 +59,19 @@ def translate(
     return out
 
 
-def translate_wrapper(text, src, trg, correct=None):
+def translate_wrapper(text, src, trg, random=False):
     src_lang = lang_to_code.get(src)
     tgt_lang = lang_to_code.get(trg)
-    if src == trg:
-        return 'Please choose two different languages'
-    print(text, src, trg)
+    # if src == trg:
+    #     return 'Please choose two different languages'
     result = translate(
         text=text,
         model=model,
         tokenizer=tokenizer,
         src_lang=src_lang,
         tgt_lang=tgt_lang,
+        do_sample=random,
+        num_beams=1 if random else 4,
     )
     return result
 
@@ -92,6 +92,7 @@ interface = gr.Interface(
         gr.Textbox(label="Text", lines=2, placeholder='text to translate '),
         gr.Dropdown(list(lang_to_code.keys()), type="value", label='source language', value=list(lang_to_code.keys())[0]),
         gr.Dropdown(list(lang_to_code.keys()), type="value", label='target language', value=list(lang_to_code.keys())[1]),
+        gr.Checkbox(label="random", value=False),
     ],
     "text",
     title='Tyvan-Russian translaton',