Update app.py

app.py

@@ -1,32 +1,97 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import requests
+from bs4 import BeautifulSoup
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM
+)
+
+# =========================
+# MODELS
+# =========================
 
 # 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)
+
+ruvi_tokenizer = AutoTokenizer.from_pretrained(
+    ruvi_model_name
+)
+
+ruvi_model = AutoModelForSeq2SeqLM.from_pretrained(
+    ruvi_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)
-
-# Simple synonym databases
-ru_synonyms = {
-    "хороший": ["отличный", "прекрасный", "добрый"],
-    "плохой": ["ужасный", "скверный"],
-    "большой": ["огромный", "крупный"]
-}
-
-vi_synonyms = {
-    "đẹp": ["xinh", "đẹp đẽ", "xinh đẹp"],
-    "lớn": ["to", "khổng lồ", "rộng"],
-    "tốt": ["hay", "ổn", "tuyệt"]
-}
+
+viru_tokenizer = AutoTokenizer.from_pretrained(
+    viru_model_name
+)
+
+viru_model = AutoModelForSeq2SeqLM.from_pretrained(
+    viru_model_name
+)
+
+# =========================
+# SYNONYM FUNCTIONS
+# =========================
+
+def get_ru_synonyms(word):
+
+    try:
+
+        url = f"https://sinonim.org/s/{word}"
+
+        headers = {
+            "User-Agent": "Mozilla/5.0"
+        }
+
+        response = requests.get(
+            url,
+            headers=headers,
+            timeout=10
+        )
+
+        soup = BeautifulSoup(
+            response.text,
+            "html.parser"
+        )
+
+        syns = []
+
+        for a in soup.select("a"):
+
+            text = a.text.strip()
+
+            if (
+                text
+                and text.lower() != word.lower()
+                and len(text) > 1
+            ):
+                syns.append(text)
+
+        syns = list(dict.fromkeys(syns))
+
+        if len(syns) == 0:
+            return "Không tìm thấy"
+
+        return ", ".join(syns[:10])
+
+    except Exception:
+        return "Lỗi tìm synonym"
+
+
+# =========================
+# TRANSLATION FUNCTIONS
+# =========================
 
 def ru_to_vi(word):
 
-    inputs = ruvi_tokenizer(word, return_tensors="pt")
+    inputs = ruvi_tokenizer(
+        word,
+        return_tensors="pt"
+    )
 
     outputs = ruvi_model.generate(**inputs)
 
@@ -35,16 +100,17 @@ def ru_to_vi(word):
         skip_special_tokens=True
     )
 
-    syns = ru_synonyms.get(
-        word.lower(),
-        ["Không tìm thấy"]
-    )
+    syns = get_ru_synonyms(word)
+
+    return translated, syns
 
-    return translated, ", ".join(syns)
 
 def vi_to_ru(word):
 
-    inputs = viru_tokenizer(word, return_tensors="pt")
+    inputs = viru_tokenizer(
+        word,
+        return_tensors="pt"
+    )
 
     outputs = viru_model.generate(**inputs)
 
@@ -53,16 +119,22 @@ def vi_to_ru(word):
         skip_special_tokens=True
     )
 
-    syns = vi_synonyms.get(
-        word.lower(),
-        ["Không tìm thấy"]
-    )
+    return translated
 
-    return translated, ", ".join(syns)
+
+# =========================
+# UI
+# =========================
 
 with gr.Blocks() as demo:
 
-    gr.Markdown("# 🇷🇺⇄🇻🇳 Russian ↔ Vietnamese Dictionary")
+    gr.Markdown(
+        "# 🇷🇺⇄🇻🇳 Russian ↔ Vietnamese Dictionary"
+    )
+
+    # -------------------------
+    # RU -> VI
+    # -------------------------
 
     with gr.Tab("🇷🇺 Nga → Việt"):
 
@@ -71,7 +143,7 @@ with gr.Blocks() as demo:
         )
 
         vi_output = gr.Textbox(
-            label="🇻🇳 Nghĩa"
+            label="🇻🇳 Nghĩa tiếng Việt"
         )
 
         ru_syn_output = gr.Textbox(
@@ -89,6 +161,10 @@ with gr.Blocks() as demo:
             ]
         )
 
+    # -------------------------
+    # VI -> RU
+    # -------------------------
+
     with gr.Tab("🇻🇳 Việt → Nga"):
 
         vi_input = gr.Textbox(
@@ -96,22 +172,15 @@ with gr.Blocks() as demo:
         )
 
         ru_output = gr.Textbox(
-            label="🇷🇺 Nghĩa"
+            label="🇷🇺 Nghĩa tiếng Nga"
         )
 
-        vi_syn_output = gr.Textbox(
-            label="📚 Từ đồng nghĩa"
-        )
-
-        vi_btn = gr.Button("Tra cứu")
+        vi_btn = gr.Button("Dịch")
 
         vi_btn.click(
             vi_to_ru,
             inputs=vi_input,
-            outputs=[
-                ru_output,
-                vi_syn_output
-            ]
+            outputs=ru_output
         )
 
 demo.launch(server_name="0.0.0.0")