Update app.py

app.py

@@ -1,31 +1,68 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
 
-model_name = "Helsinki-NLP/opus-mt-ru-vi"
+# RU -> VI
+ruvi_model_name = "Helsinki-NLP/opus-mt-ru-vi"
+ruvi_tokenizer = AutoTokenizer.from_pretrained(ruvi_model_name)
+ruvi_model = AutoModelForSeq2SeqLM.from_pretrained(ruvi_model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# VI -> RU
+viru_model_name = "Helsinki-NLP/opus-mt-vi-ru"
+viru_tokenizer = AutoTokenizer.from_pretrained(viru_model_name)
+viru_model = AutoModelForSeq2SeqLM.from_pretrained(viru_model_name)
 
-def translate(text):
-    inputs = tokenizer(text, return_tensors="pt")
+def ru_to_vi(text):
+    inputs = ruvi_tokenizer(text, return_tensors="pt")
+    outputs = ruvi_model.generate(**inputs)
 
-    outputs = model.generate(**inputs)
+    return ruvi_tokenizer.decode(
+        outputs[0],
+        skip_special_tokens=True
+    )
+
+def vi_to_ru(text):
+    inputs = viru_tokenizer(text, return_tensors="pt")
+    outputs = viru_model.generate(**inputs)
 
-    translated = tokenizer.decode(
+    return viru_tokenizer.decode(
         outputs[0],
         skip_special_tokens=True
     )
 
-    return translated
-
-iface = gr.Interface(
-    fn=translate,
-    inputs=gr.Textbox(
-        lines=2,
-        placeholder="Nhập tiếng Nga..."
-    ),
-    outputs="text",
-    title="Nga → Việt"
-)
+with gr.Blocks() as iface:
+
+    gr.Markdown("# 🇷🇺⇄🇻🇳 Russian ↔ Vietnamese")
+
+    with gr.Tab("🇷🇺 Nga → Việt"):
+        input_ru = gr.Textbox(
+            lines=2,
+            placeholder="Nhập tiếng Nga..."
+        )
+
+        output_vi = gr.Textbox()
+
+        btn1 = gr.Button("Translate")
+
+        btn1.click(
+            ru_to_vi,
+            inputs=input_ru,
+            outputs=output_vi
+        )
+
+    with gr.Tab("🇻🇳 Việt → Nga"):
+        input_vi = gr.Textbox(
+            lines=2,
+            placeholder="Nhập tiếng Việt..."
+        )
+
+        output_ru = gr.Textbox()
+
+        btn2 = gr.Button("Translate")
+
+        btn2.click(
+            vi_to_ru,
+            inputs=input_vi,
+            outputs=output_ru
+        )
 
 iface.launch(server_name="0.0.0.0")