Update app.py

app.py

@@ -18,6 +18,18 @@ os.environ["MKL_NUM_THREADS"] = str(psutil.cpu_count(logical=False))
 # Set device globally
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+@st.cache_resource()
+def load_model(model_name, model_class, is_bc=False, device=None):
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = model_class.from_pretrained(model_name, num_labels=3 if not is_bc else 2)
+    model.eval()
+    
+    model.to(device)
+    return tokenizer, model
+
 @st.cache_data
 def preprocess_text(text):
     # Add any text cleaning or normalization here
@@ -309,6 +321,10 @@ with st.sidebar:
 # Main content
 tabs = st.tabs(["🔍 Kiểm chứng", "📊 Lịch sử", "ℹ️ Thông tin"])
 
+tokenizer_qatc, model_qatc = load_model(qatc_model_name, QATCForQuestionAnswering, device=DEVICE)
+tokenizer_bc, model_bc = load_model(bc_model_name, ClaimModelForClassification, is_bc=True, device=DEVICE)
+tokenizer_tc, model_tc = load_model(tc_model_name, ClaimModelForClassification, device=DEVICE)
+
 # --- Tab Verify ---
 with tabs[0]:
     col1, col2 = st.columns([2, 1])