Update app.py

app.py

@@ -2,33 +2,110 @@ import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer
 import pysrt
 import tempfile
+from tqdm import tqdm
 
-def translate_subtitle(input_file, model_name="Helsinki-NLP/opus-mt-en-vi"):
-    # Load model
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    
-    # Đọc file
-    subs = pysrt.open(input_file.name)
+# Danh sách các ngôn ngữ và model tương ứng
+LANGUAGE_MODELS = {
+    "Tiếng Việt": "Helsinki-NLP/opus-mt-en-vi",
+    "Tiếng Pháp": "Helsinki-NLP/opus-mt-en-fr",
+    "Tiếng Đức": "Helsinki-NLP/opus-mt-en-de",
+    "Tiếng Tây Ban Nha": "Helsinki-NLP/opus-mt-en-es",
+    "Tiếng Bồ Đào Nha": "Helsinki-NLP/opus-mt-en-pt",
+    "Tiếng Ý": "Helsinki-NLP/opus-mt-en-it",
+    "Tiếng Nga": "Helsinki-NLP/opus-mt-en-ru",
+    "Tiếng Hà Lan": "Helsinki-NLP/opus-mt-en-nl",
+    "Tiếng Thụy Điển": "Helsinki-NLP/opus-mt-en-sv",
+    "Tiếng Phần Lan": "Helsinki-NLP/opus-mt-en-fi",
+    "Tiếng Đan Mạch": "Helsinki-NLP/opus-mt-en-da",
+    "Tiếng Na Uy": "Helsinki-NLP/opus-mt-en-no",
+    "Tiếng Ba Lan": "Helsinki-NLP/opus-mt-en-pl",
+    "Tiếng Séc": "Helsinki-NLP/opus-mt-en-cs",
+    "Tiếng Hungary": "Helsinki-NLP/opus-mt-en-hu",
+    "Tiếng Romania": "Helsinki-NLP/opus-mt-en-ro",
+    "Tiếng Hy Lạp": "Helsinki-NLP/opus-mt-en-el",
+    "Tiếng Thổ Nhĩ Kỳ": "Helsinki-NLP/opus-mt-en-tr",
+    "Tiếng Hindi (Ấn Độ)": "Helsinki-NLP/opus-mt-en-hi",
+    "Tiếng Ả Rập": "Helsinki-NLP/opus-mt-en-ar",
+    "Tiếng Trung (Giản thể)": "Helsinki-NLP/opus-mt-en-zh",
+    "Tiếng Nhật": "Helsinki-NLP/opus-mt-en-ja",
+    "Tiếng Hàn": "Helsinki-NLP/opus-mt-en-ko"
+}
+
+# Cache models để tăng tốc độ
+model_cache = {}
+
+def get_model(language):
+    if language not in model_cache:
+        model_name = LANGUAGE_MODELS[language]
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        model_cache[language] = (model, tokenizer)
+    return model_cache[language]
+
+def translate_text(text, model, tokenizer):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    translated = model.generate(**inputs)
+    return tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
+
+def translate_subtitle(input_file, target_language):
+    if input_file is None:
+        raise gr.Error("Vui lòng upload file phụ đề!")
     
-    # Dịch từng dòng
-    for sub in subs:
-        if sub.text.strip():
-            inputs = tokenizer(sub.text, return_tensors="pt", truncation=True)
-            translated = model.generate(**inputs)
-            sub.text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
+    try:
+        model, tokenizer = get_model(target_language)
+        subs = pysrt.open(input_file.name)
+        
+        # Dịch từng dòng với thanh tiến trình
+        for sub in tqdm(subs, desc="Đang dịch"):
+            if sub.text.strip():
+                sub.text = translate_text(sub.text, model, tokenizer)
+        
+        # Lưu file tạm
+        output_path = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
+        subs.save(output_path, encoding='utf-8')
+        return output_path
     
-    # Lưu file tạm
-    output_path = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
-    subs.save(output_path)
-    return output_path
+    except Exception as e:
+        raise gr.Error(f"Có lỗi xảy ra: {str(e)}")
 
-iface = gr.Interface(
-    fn=translate_subtitle,
-    inputs=gr.File(label="Upload Subtitle File"),
-    outputs=gr.File(label="Download Translated Subtitle"),
-    title="Subtitle Translator",
-    description="Upload a subtitle file (.srt) to translate it using Hugging Face models"
-)
+# Giao diện Gradio
+with gr.Blocks(title="Subtitle Translator", theme="soft") as demo:
+    gr.Markdown("# 🎬 Subtitle Translator")
+    gr.Markdown("Dịch phụ đề (.srt) sang nhiều ngôn ngữ khác nhau sử dụng mô hình MarianMT từ Hugging Face")
+    
+    with gr.Row():
+        with gr.Column():
+            file_input = gr.File(label="Upload file phụ đề (.srt)", type="file")
+            language_dropdown = gr.Dropdown(
+                choices=list(LANGUAGE_MODELS.keys()),
+                value="Tiếng Việt",
+                label="Chọn ngôn ngữ đích"
+            )
+            translate_btn = gr.Button("Dịch phụ đề", variant="primary")
+        
+        with gr.Column():
+            file_output = gr.File(label="File phụ đề đã dịch", interactive=False)
+            gr.Examples(
+                examples=[["sample.srt", "Tiếng Việt"]],
+                inputs=[file_input, language_dropdown],
+                outputs=file_output,
+                fn=translate_subtitle,
+                cache_examples=True
+            )
+    
+    translate_btn.click(
+        fn=translate_subtitle,
+        inputs=[file_input, language_dropdown],
+        outputs=file_output
+    )
+    
+    gr.Markdown("### Thông tin")
+    gr.Markdown("""
+    - Hỗ trợ định dạng .srt
+    - Sử dụng mô hình MarianMT từ Hugging Face
+    - Dịch chính xác từng dòng phụ đề
+    - Hỗ trợ nhiều ngôn ngữ khác nhau
+    """)
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()