Update app.py

app.py

@@ -1,29 +1,31 @@
 import math
 import gradio as gr
-from transformers import MarianTokenizer, MarianMTModel
+from transformers import MT5Tokenizer, MT5ForConditionalGeneration
 import re
 import functools
 
 ###################################
-# 1) โหลดโมเดล MarianMT (Thai->En) ครั้งเดียวเมื่อเริ่มแอป
+# 1) โหลดโมเดล mT5 (Thai->En) ครั้งเดียวเมื่อเริ่มแอป
 ###################################
-model_name = "Helsinki-NLP/opus-mt-th-en"
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
+model_name = "google/mt5-small"
+tokenizer = MT5Tokenizer.from_pretrained(model_name)
+model = MT5ForConditionalGeneration.from_pretrained(model_name)
 
 # เพิ่มการแคชผลลัพธ์การแปล
 @functools.lru_cache(maxsize=1024)
 def translate_th_to_en(text_th: str) -> str:
     """
-    แปลไทย -> อังกฤษ ด้วย MarianMT บน CPU
+    แปลไทย -> อังกฤษ ด้วย mT5 บน CPU
     * เรียกเฉพาะส่วนที่ผู้ใช้พิมพ์ เช่น ชื่อดาว, ชื่อสิ่งมีชีวิต
     """
     text_th = text_th.strip()
     if not text_th:
         return ""
-    inputs = tokenizer(text_th, return_tensors="pt", max_length=512, truncation=True)
-    translation_tokens = model.generate(**inputs, max_length=512)
-    en_text = tokenizer.decode(translation_tokens[0], skip_special_tokens=True)
+    # เพิ่ม prefix เพื่อบอกว่าเป็นการแปลภาษาไทยเป็นอังกฤษ
+    input_text = f"translate Thai to English: {text_th}"
+    inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
+    outputs = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
+    en_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return en_text
 
 ###################################