Update app.py

app.py

@@ -1,110 +1,108 @@
-import torch
-import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-
-# --- PHẦN 1: CẤU HÌNH & LOAD MODEL ---
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# 1. Load PhoWhisper (Nhận dạng giọng nói)
-print("Đang tải model PhoWhisper...")
-asr_pipeline = pipeline(
-    "automatic-speech-recognition",
-    model="vinai/PhoWhisper-small",
-    device=0 if torch.cuda.is_available() else -1
-)
-
-# 2. Load Model Dịch (ViT5)
-# LƯU Ý QUAN TRỌNG: 
-# Bạn phải upload folder chứa model (tên là 'trained_model') lên cùng chỗ với file app.py trên Hugging Face
-# Sau đó sửa đường dẫn dưới đây thành "./trained_model"
-TRANSLATION_MODEL_PATH = "./trained_model" 
-
-print(f"Đang tải model dịch...")
-try:
-    trans_tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_MODEL_PATH)
-    trans_model = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATION_MODEL_PATH).to(DEVICE)
-    print("Load model dịch thành công!")
-except Exception as e:
-    print(f"Lỗi load model dịch: {e}")
-    trans_model = None
-
-
-# --- PHẦN 2: CÁC HÀM XỬ LÝ (LOGIC) ---
-def speech_to_text(audio_path):
-    if audio_path is None: return ""
-    try:
-        output = asr_pipeline(audio_path)
-        return output['text']
-    except Exception as e:
-        return f"Lỗi: {str(e)}"
-
-def text_to_gloss(vietnamese_text):
-    if not vietnamese_text: return ""
-    if trans_model is None: return "Lỗi: Chưa load được model dịch."
-
-    input_text = f"vi: {vietnamese_text}"
-    inputs = trans_tokenizer(input_text, return_tensors="pt", max_length=128, truncation=True).to(DEVICE)
-
-    with torch.no_grad():
-        outputs = trans_model.generate(
-            inputs["input_ids"],
-            max_length=128,
-            num_beams=5,
-            early_stopping=True
-        )
-    
-    gloss_text = trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return gloss_text.replace("vsl: ", "") if gloss_text.startswith("vsl: ") else gloss_text
-
-def full_pipeline(audio, text_input, mode):
-    vietnamese_output = ""
-    if mode == "Giọng nói (Microphone/File)" and audio is not None:
-        vietnamese_output = speech_to_text(audio)
-    elif mode == "Văn bản (Nhập tay)" and text_input:
-        vietnamese_output = text_input
-    else:
-        return "Vui lòng nhập dữ liệu.", ""
-    
-    gloss_output = text_to_gloss(vietnamese_output)
-    return vietnamese_output, gloss_output
-
-
-# --- PHẦN 3: GIAO DIỆN (UI) - ĐOẠN CODE CỦA BẠN ---
-custom_css = """
-.container {max-width: 1200px; margin: auto; padding-top: 20px}
-.header-text {text-align: center; font-family: 'Arial', sans-serif;}
-.uni-name {font-size: 24px; font-weight: bold; color: #003366; margin-bottom: 5px;}
-.faculty-name {font-size: 18px; font-weight: normal; color: #cc0000; margin-bottom: 20px;}
-.project-title {font-size: 28px; font-weight: bold; color: #2c3e50; margin-bottom: 10px; border-bottom: 2px solid #eee; padding-bottom: 10px;}
-.note-text {font-size: 14px; font-style: italic; color: #7f8c8d; margin-top: 20px; border-top: 1px solid #eee; padding-top: 10px;}
-.output-box {border: 1px solid #e0e0e0; background-color: #f9f9f9; border-radius: 8px;}
-"""
-
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    # ... (Dán phần giao diện của bạn vào đây, y hệt đoạn bạn gửi) ...
-    # Để cho gọn, mình viết tắt, bạn giữ nguyên code UI của bạn nhé
-    with gr.Column(elem_classes="header-text"):
-        gr.Markdown("""<div class="project-title">HỆ THỐNG DỊCH VSL DEMO</div>""")
-    
-    with gr.Row():
-        with gr.Column():
-            with gr.Tabs():
-                with gr.TabItem("Giọng nói"):
-                    input_audio = gr.Audio(sources=["microphone", "upload"], type="filepath")
-                    mode_audio = gr.State(value="Giọng nói (Microphone/File)")
-                    btn_audio = gr.Button("Xử lý Giọng nói", variant="primary")
-                with gr.TabItem("Văn bản"):
-                    input_text = gr.Textbox(label="Nhập text")
-                    mode_text = gr.State(value="Văn bản (Nhập tay)")
-                    btn_text = gr.Button("Dịch Văn bản", variant="primary")
-
-        with gr.Column():
-            output_vi = gr.Textbox(label="Tiếng Việt")
-            output_gloss = gr.Textbox(label="VSL Gloss")
-
-    btn_audio.click(fn=full_pipeline, inputs=[input_audio, input_text, mode_audio], outputs=[output_vi, output_gloss])
-    btn_text.click(fn=full_pipeline, inputs=[input_audio, input_text, mode_text], outputs=[output_vi, output_gloss])
-
-# Chạy app
-if __name__ == "__main__":
+import torch
+import gradio as gr
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+
+# --- PHẦN 1: CẤU HÌNH & LOAD MODEL ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 1. Load PhoWhisper (Nhận dạng giọng nói)
+print("Đang tải model PhoWhisper...")
+asr_pipeline = pipeline(
+    "automatic-speech-recognition",
+    model="vinai/PhoWhisper-small",
+    device=0 if torch.cuda.is_available() else -1
+)
+
+# 2. Load Model Dịch (ViT5)
+# LƯU Ý QUAN TRỌNG: 
+TRANSLATION_MODEL_PATH = "." 
+
+print(f"Đang tải model dịch...")
+try:
+    trans_tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_MODEL_PATH)
+    trans_model = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATION_MODEL_PATH).to(DEVICE)
+    print("Load model dịch thành công!")
+except Exception as e:
+    print(f"Lỗi load model dịch: {e}")
+    trans_model = None
+
+
+# --- PHẦN 2: CÁC HÀM XỬ LÝ (LOGIC) ---
+def speech_to_text(audio_path):
+    if audio_path is None: return ""
+    try:
+        output = asr_pipeline(audio_path)
+        return output['text']
+    except Exception as e:
+        return f"Lỗi: {str(e)}"
+
+def text_to_gloss(vietnamese_text):
+    if not vietnamese_text: return ""
+    if trans_model is None: return "Lỗi: Chưa load được model dịch."
+
+    input_text = f"vi: {vietnamese_text}"
+    inputs = trans_tokenizer(input_text, return_tensors="pt", max_length=128, truncation=True).to(DEVICE)
+
+    with torch.no_grad():
+        outputs = trans_model.generate(
+            inputs["input_ids"],
+            max_length=128,
+            num_beams=5,
+            early_stopping=True
+        )
+    
+    gloss_text = trans_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return gloss_text.replace("vsl: ", "") if gloss_text.startswith("vsl: ") else gloss_text
+
+def full_pipeline(audio, text_input, mode):
+    vietnamese_output = ""
+    if mode == "Giọng nói (Microphone/File)" and audio is not None:
+        vietnamese_output = speech_to_text(audio)
+    elif mode == "Văn bản (Nhập tay)" and text_input:
+        vietnamese_output = text_input
+    else:
+        return "Vui lòng nhập dữ liệu.", ""
+    
+    gloss_output = text_to_gloss(vietnamese_output)
+    return vietnamese_output, gloss_output
+
+
+# --- PHẦN 3: GIAO DIỆN (UI) - ĐOẠN CODE CỦA BẠN ---
+custom_css = """
+.container {max-width: 1200px; margin: auto; padding-top: 20px}
+.header-text {text-align: center; font-family: 'Arial', sans-serif;}
+.uni-name {font-size: 24px; font-weight: bold; color: #003366; margin-bottom: 5px;}
+.faculty-name {font-size: 18px; font-weight: normal; color: #cc0000; margin-bottom: 20px;}
+.project-title {font-size: 28px; font-weight: bold; color: #2c3e50; margin-bottom: 10px; border-bottom: 2px solid #eee; padding-bottom: 10px;}
+.note-text {font-size: 14px; font-style: italic; color: #7f8c8d; margin-top: 20px; border-top: 1px solid #eee; padding-top: 10px;}
+.output-box {border: 1px solid #e0e0e0; background-color: #f9f9f9; border-radius: 8px;}
+"""
+
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    # ... (Dán phần giao diện của bạn vào đây, y hệt đoạn bạn gửi) ...
+    # Để cho gọn, mình viết tắt, bạn giữ nguyên code UI của bạn nhé
+    with gr.Column(elem_classes="header-text"):
+        gr.Markdown("""<div class="project-title">HỆ THỐNG DỊCH VSL DEMO</div>""")
+    
+    with gr.Row():
+        with gr.Column():
+            with gr.Tabs():
+                with gr.TabItem("Giọng nói"):
+                    input_audio = gr.Audio(sources=["microphone", "upload"], type="filepath")
+                    mode_audio = gr.State(value="Giọng nói (Microphone/File)")
+                    btn_audio = gr.Button("Xử lý Giọng nói", variant="primary")
+                with gr.TabItem("Văn bản"):
+                    input_text = gr.Textbox(label="Nhập text")
+                    mode_text = gr.State(value="Văn bản (Nhập tay)")
+                    btn_text = gr.Button("Dịch Văn bản", variant="primary")
+
+        with gr.Column():
+            output_vi = gr.Textbox(label="Tiếng Việt")
+            output_gloss = gr.Textbox(label="VSL Gloss")
+
+    btn_audio.click(fn=full_pipeline, inputs=[input_audio, input_text, mode_audio], outputs=[output_vi, output_gloss])
+    btn_text.click(fn=full_pipeline, inputs=[input_audio, input_text, mode_text], outputs=[output_vi, output_gloss])
+
+# Chạy app
+if __name__ == "__main__":
     demo.launch()