Update app.py

app.py

@@ -28,10 +28,8 @@ def fix_bartpho_output(text: str) -> str:
 
     
 # HuggingFace Transformers
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM
-)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel, PeftConfig
 #page config
 st.set_page_config(
     page_title="ViSum - Vietnamese News Summarization",
@@ -97,22 +95,29 @@ MODEL_ID = "OrdinaryAI/visum-qlora-5epochs"
 
 @st.cache_resource
 def load_model(model_id):
-
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-    # Load model
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-
-    # Kiểm tra có GPU không
+    # Đọc config PEFT để biết model gốc là gì
+    peft_config = PeftConfig.from_pretrained(model_id)
+    
+    # Load model gốc (vinai/bartpho-syllable)
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(
+        peft_config.base_model_name_or_path
+    )
+    
+    # Gắn trọng số QLoRA vào
+    model = PeftModel.from_pretrained(base_model, model_id)
+    
+    # Merge vào model gốc → inference nhanh hơn
+    model = model.merge_and_unload()
+    
+    # Load tokenizer từ model gốc
+    tokenizer = AutoTokenizer.from_pretrained(
+        peft_config.base_model_name_or_path
+    )
+    
     device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # Đưa model sang device tương ứng
     model = model.to(device)
-
-    # Chuyển sang chế độ inference
     model.eval()
-
+    
     return tokenizer, model