Update app.py

app.py

@@ -10,117 +10,77 @@ from transformers import (
 )
 import re
 
-# --- THAY ĐỔI CÁC THÔNG TIN SAU CHO ĐÚNG VỚI REPO CỦA BẠN ---
-BART_MODEL_REPO = "Tin113/bart_model"
-VIT5_MODEL_REPO = "Tin113/vit5_model"
-# -------------------------------------------------------------
+# --- THAY ĐỔI USERNAME CỦA BẠN VÀO ĐÂY ---
+HF_USERNAME = "Tin113"
+# -----------------------------------------
 
-# Chọn thiết bị (GPU nếu có)
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-print(f"Bắt đầu tải model lên thiết bị: {DEVICE}...")
+BART_MODEL_REPO = f"{HF_USERNAME}/bart_model"
+VIT5_MODEL_REPO = f"{HF_USERNAME}/vit5_model"
 
-# Hàm tải model, với cơ chế thử lại để tránh lỗi tải tạm thời
-def load_model(model_class, tokenizer_class, repo_id):
-    try:
-        tokenizer = tokenizer_class.from_pretrained(repo_id)
-        model = model_class.from_pretrained(repo_id).to(DEVICE)
-        model.eval()
-        print(f"Tải thành công model: {repo_id}")
-        return model, tokenizer
-    except Exception as e:
-        print(f"Lỗi khi tải model {repo_id}: {e}")
-        # Trả về None nếu có lỗi để xử lý ở giao diện
-        return None, None
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Thiết bị sử dụng: {DEVICE}")
 
 # Tải các model
-model_bart, tokenizer_bart = load_model(AutoModelForSeq2SeqLM, AutoTokenizer, BART_MODEL_REPO)
-model_vit5, tokenizer_vit5 = load_model(T5ForConditionalGeneration, T5Tokenizer, VIT5_MODEL_REPO)
+try:
+    print(f"Đang tải model BART từ {BART_MODEL_REPO}...")
+    tokenizer_bart = AutoTokenizer.from_pretrained(BART_MODEL_REPO)
+    model_bart = AutoModelForSeq2SeqLM.from_pretrained(BART_MODEL_REPO).to(DEVICE)
+    model_bart.eval()
+    print("Tải model BART thành công.")
+except Exception as e:
+    print(f"Lỗi khi tải model BART: {e}")
+    model_bart, tokenizer_bart = None, None
+
+try:
+    print(f"Đang tải model ViT5 từ {VIT5_MODEL_REPO}...")
+    tokenizer_vit5 = T5Tokenizer.from_pretrained(VIT5_MODEL_REPO)
+    model_vit5 = T5ForConditionalGeneration.from_pretrained(VIT5_MODEL_REPO).to(DEVICE)
+    model_vit5.eval()
+    print("Tải model ViT5 thành công.")
+except Exception as e:
+    print(f"Lỗi khi tải model ViT5: {e}")
+    model_vit5, tokenizer_vit5 = None, None
 
-# Hàm clean text, lấy từ notebook của bạn
 def clean_text(text):
-    if not isinstance(text, str):
-        return ""
+    if not isinstance(text, str): return ""
     return re.sub(r'\s+', ' ', text).strip()
 
-# Hàm xử lý việc sửa lỗi
 def correct_grammar(sentence, model_choice):
-    if not sentence.strip():
-        return "Vui lòng nhập một câu."
-
-    model = None
-    tokenizer = None
-    prefix = ""
+    if not sentence.strip(): return "Vui lòng nhập một câu."
 
+    model, tokenizer, prefix = None, None, ""
     if model_choice == "BARTpho-syllable":
-        if model_bart and tokenizer_bart:
-            model = model_bart
-            tokenizer = tokenizer_bart
-            prefix = "Fix: "
+        if model_bart:
+            model, tokenizer, prefix = model_bart, tokenizer_bart, "Fix: "
         else:
-            return "Lỗi: Model BART không khả dụng. Vui lòng thử lại sau."
-
+            return "Lỗi: Model BART không khả dụng. Vui lòng kiểm tra lại Space."
     elif model_choice == "ViT5-base":
-        if model_vit5 and tokenizer_vit5:
-            model = model_vit5
-            tokenizer = tokenizer_vit5
-            prefix = "sửa lỗi: "
+        if model_vit5:
+            model, tokenizer, prefix = model_vit5, tokenizer_vit5, "sửa lỗi: "
         else:
-            return "Lỗi: Model ViT5 không khả dụng. Vui lòng thử lại sau."
-
+            return "Lỗi: Model ViT5 không khả dụng. Vui lòng kiểm tra lại Space."
+    
     input_text = prefix + sentence
-    input_ids = tokenizer(
-        input_text,
-        return_tensors="pt",
-        max_length=256,
-        truncation=True,
-        padding=True
-    ).input_ids.to(DEVICE)
+    input_ids = tokenizer(input_text, return_tensors="pt", max_length=256, truncation=True, padding=True).input_ids.to(DEVICE)
 
     with torch.no_grad():
-        outputs = model.generate(
-            input_ids,
-            max_length=256 + 20,
-            num_beams=2,
-            early_stopping=True,
-            repetition_penalty=1.05,
-            no_repeat_ngram_size=2
-        )
-
-    corrected_sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return clean_text(corrected_sentence)
+        outputs = model.generate(input_ids, max_length=276, num_beams=2, early_stopping=True, repetition_penalty=1.05, no_repeat_ngram_size=2)
+    
+    return clean_text(tokenizer.decode(outputs[0], skip_special_tokens=True))
 
-# Ví dụ cho giao diện Gradio
-examples = [
-    ["chương trỉnhnh được páht sóng vào lúc 19h", "ViT5-base"],
-    ["công nghề thônngg tin đáng phát chiển rất nhanh", "ViT5-base"],
-    ["Học hok tốt thì kho mak đc điểm cao.", "BARTpho-syllable"],
-    ["dù rất mệt nhưng anh ấy vẫn cố hoàn thành công việc", "BARTpho-syllable"],
-]
-
-# Mô tả cho ứng dụng
 description = """
-Đây là ứng dụng demo cho việc sửa lỗi ngữ pháp tiếng Việt (Vietnamese Grammatical Error Correction).
-Ứng dụng sử dụng hai model đã được fine-tune:
-1.  **BARTpho-syllable**: Dựa trên kiến trúc BART, được tối ưu cho tiếng Việt ở cấp độ âm tiết.
-2.  **ViT5-base**: Dựa trên kiến trúc T5, một model mạnh mẽ cho các tác vụ Text-to-Text.
-
-**Cách sử dụng:**
-1.  Nhập câu tiếng Việt có lỗi vào ô bên dưới.
-2.  Chọn một trong hai model để thực hiện sửa lỗi.
-3.  Nhấn "Submit" và xem kết quả.
+Demo sửa lỗi ngữ pháp tiếng Việt sử dụng hai model: BARTpho-syllable và ViT5-base.
+1. Nhập câu lỗi vào ô bên dưới.
+2. Chọn model bạn muốn dùng.
+3. Nhấn "Submit" để xem kết quả.
 """
+examples = [["chương trỉnhnh được páht sóng vào lúc 19h", "ViT5-base"], ["Học hok tốt thì kho mak đc điểm cao.", "BARTpho-syllable"]]
 
-# Tạo giao diện Gradio
 demo = gr.Interface(
     fn=correct_grammar,
     inputs=[
         gr.Textbox(lines=5, label="Nhập câu tiếng Việt bị lỗi"),
-        gr.Radio(
-            choices=["BARTpho-syllable", "ViT5-base"],
-            value="ViT5-base", # Model mặc định
-            label="Chọn Model"
-        )
+        gr.Radio(choices=["BARTpho-syllable", "ViT5-base"], value="ViT5-base", label="Chọn Model")
     ],
     outputs=gr.Textbox(label="Câu đã được sửa"),
     title="Sửa lỗi Ngữ pháp Tiếng Việt",
@@ -128,6 +88,5 @@ demo = gr.Interface(
     examples=examples
 )
 
-# Chạy ứng dụng
 if __name__ == "__main__":
     demo.launch()