Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import gradio as gr
+import opencc
 
 # Load model và tokenizer
 model_name = "chi-vi/hirashiba-mt-jp-names"
@@ -9,6 +10,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
 
+# Khởi tạo OpenCC converter
+converter = opencc.OpenCC('t2s')  # Chuyển đổi từ Phồn thể sang Giản thể
+
 def translate_text(input_text):
     lines = input_text.split('\n')  # Tách từng dòng
     translated_lines = []
@@ -19,15 +23,18 @@ def translate_text(input_text):
             translated_lines.append('')  # Giữ dòng trống
             continue
         
+        # Chuyển đổi sang giản thể
+        simplified_text = converter.convert(raw_text)
+        
         # Tokenize input
-        inputs = tokenizer(raw_text, return_tensors="pt", padding=True, truncation=True).to(device)
+        inputs = tokenizer(simplified_text, return_tensors="pt", padding=True, truncation=True).to(device)
         
         # Dịch với mô hình (không cần tính gradient)
         with torch.no_grad():
             output_tokens = model.generate(**inputs, max_length=512)
         
         # Giải mã kết quả và viết hoa chữ đầu
-        translated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True).capwords()
+        translated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True).capitalize()
         translated_lines.append(translated_text)
     
     return '\n'.join(translated_lines)