Update app.py

app.py

@@ -8,14 +8,76 @@ from semviqa.tvc.tvc_eval import classify_claim
 import time
 import pandas as pd
 
-# Load models with caching
+# Load models with caching and optimization
 @st.cache_resource()
-def load_model(model_name, model_class, is_bc=False):
+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)
+    
+    # Enable CUDA optimizations if available
+    if device == "cuda":
+        if hasattr(model, 'half') and not model_name.startswith("SemViQA/bc-erniem") and not model_name.startswith("SemViQA/tc-erniem"):
+            model = model.half()  # Use FP16 for most models (except ERNIE which might not support it)
+    
     return tokenizer, model
 
+# Set device globally
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Pre-process text function to avoid doing it multiple times
+@st.cache_data
+def preprocess_text(text):
+    # Add any text cleaning or normalization here
+    return text.strip()
+
+# Optimized function for evidence extraction and classification
+def perform_verification(claim, context, model_qatc, tokenizer_qatc, model_tc, tokenizer_tc, 
+                          model_bc, tokenizer_bc, tfidf_threshold, length_ratio_threshold):
+    with torch.no_grad():
+        # Extract evidence
+        evidence_start_time = time.time()
+        evidence = extract_evidence_tfidf_qatc(
+            claim, context, model_qatc, tokenizer_qatc,
+            DEVICE,
+            confidence_threshold=tfidf_threshold,
+            length_ratio_threshold=length_ratio_threshold
+        )
+        evidence_time = time.time() - evidence_start_time
+
+        # Classify the claim
+        verdict_start_time = time.time()
+        prob3class, pred_tc = classify_claim(
+            claim, evidence, model_tc, tokenizer_tc, DEVICE
+        )
+        
+        # Only run binary classifier if needed
+        prob2class, pred_bc = 0, "Not used"
+        if pred_tc != 0:
+            prob2class, pred_bc = classify_claim(
+                claim, evidence, model_bc, tokenizer_bc, DEVICE
+            )
+            verdict = "SUPPORTED" if pred_bc == 0 else "REFUTED" if prob2class > prob3class else ["NEI", "SUPPORTED", "REFUTED"][pred_tc]
+        else:
+            verdict = "NEI"
+            
+        verdict_time = time.time() - verdict_start_time
+
+    return {
+        "evidence": evidence,
+        "verdict": verdict,
+        "evidence_time": evidence_time,
+        "verdict_time": verdict_time,
+        "prob3class": prob3class,
+        "pred_tc": pred_tc,
+        "prob2class": prob2class,
+        "pred_bc": pred_bc
+    }
+
 # Set up page configuration
 st.set_page_config(page_title="SemViQA Demo", layout="wide")
 
@@ -166,11 +228,31 @@ with st.container():
         st.session_state.history = []
     if 'latest_result' not in st.session_state:
         st.session_state.latest_result = None
+    if 'models_loaded' not in st.session_state:
+        st.session_state.models_loaded = False
 
-    # Load the selected models
-    tokenizer_qatc, model_qatc = load_model(qatc_model_name, QATCForQuestionAnswering)
-    tokenizer_bc, model_bc = load_model(bc_model_name, ClaimModelForClassification, is_bc=True)
-    tokenizer_tc, model_tc = load_model(tc_model_name, ClaimModelForClassification)
+    # Load the selected models - only reload if model selection changes
+    if not st.session_state.models_loaded or 'prev_models' not in st.session_state or (
+            st.session_state.prev_models['qatc'] != qatc_model_name or
+            st.session_state.prev_models['bc'] != bc_model_name or
+            st.session_state.prev_models['tc'] != tc_model_name):
+        
+        with st.spinner("Loading models..."):
+            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)
+            
+            st.session_state.prev_models = {
+                'qatc': qatc_model_name,
+                'bc': bc_model_name,
+                'tc': tc_model_name
+            }
+            st.session_state.models_loaded = True
+    else:
+        # Reuse already loaded models
+        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)
 
     # Icons for results
     verdict_icons = {
@@ -196,52 +278,41 @@ with st.container():
         with col_result:
             st.markdown("<h3>Verification Result</h3>", unsafe_allow_html=True)
             if verify_button:
+                # Preprocess texts to improve tokenization speed
+                preprocessed_claim = preprocess_text(claim)
+                preprocessed_context = preprocess_text(context)
+                
                 # Placeholder for displaying result/loading
                 with st.spinner("Verifying..."): 
-                    start_time = time.time() 
-    
-                    # Extract evidence
-                    evidence_start_time = time.time()
-                    with torch.no_grad():
-                        evidence = extract_evidence_tfidf_qatc(
-                            claim, context, model_qatc, tokenizer_qatc,
-                            "cuda" if torch.cuda.is_available() else "cpu",
-                            confidence_threshold=tfidf_threshold,
-                            length_ratio_threshold=length_ratio_threshold
-                        )
-                    evidence_time = time.time() - evidence_start_time
-
-                    # Classify the claim
-                    verdict = "NEI"
+                    start_time = time.time()
+                    
+                    # Use the optimized verification function
+                    result = perform_verification(
+                        preprocessed_claim, preprocessed_context, 
+                        model_qatc, tokenizer_qatc,
+                        model_tc, tokenizer_tc,
+                        model_bc, tokenizer_bc,
+                        tfidf_threshold, length_ratio_threshold
+                    )
+                    
+                    total_time = time.time() - start_time
+                    
+                    # Format details if needed
                     details = ""
-                    verdict_start_time = time.time()
-                    with torch.no_grad():
-                        prob2class, pred_bc = 0, "Not used"
-                        prob3class, pred_tc = classify_claim(
-                            claim, evidence, model_tc, tokenizer_tc,
-                            "cuda" if torch.cuda.is_available() else "cpu"
-                        )
-                        if pred_tc != 0:
-                            prob2class, pred_bc = classify_claim(
-                                claim, evidence, model_bc, tokenizer_bc,
-                                "cuda" if torch.cuda.is_available() else "cpu"
-                            )
-                            verdict = "SUPPORTED" if pred_bc == 0 else "REFUTED" if prob2class > prob3class else ["NEI", "SUPPORTED", "REFUTED"][pred_tc]
-                    verdict_time = time.time() - verdict_start_time
                     if show_details:
                         details = f"""
-                            3-Class Probability: {prob3class.item():.2f}
-                            3-Class Predicted Label: {['NEI', 'SUPPORTED', 'REFUTED'][pred_tc]}
-                            2-Class Probability: {prob2class.item():.2f}
-                            2-Class Predicted Label: {['SUPPORTED', 'REFUTED'][pred_bc]}
+                            3-Class Probability: {result['prob3class'].item():.2f}
+                            3-Class Predicted Label: {['NEI', 'SUPPORTED', 'REFUTED'][result['pred_tc']]}
+                            2-Class Probability: {result['prob2class'].item():.2f}
+                            2-Class Predicted Label: {['SUPPORTED', 'REFUTED'][result['pred_bc']] if result['pred_tc'] != 0 else 'Not used'}
                         """
-                    total_time = time.time() - start_time
+                    
                     st.session_state.latest_result = {
                         "claim": claim,
-                        "evidence": evidence,
-                        "verdict": verdict,
-                        "evidence_time": evidence_time,
-                        "verdict_time": verdict_time,
+                        "evidence": result['evidence'],
+                        "verdict": result['verdict'],
+                        "evidence_time": result['evidence_time'],
+                        "verdict_time": result['verdict_time'],
                         "total_time": total_time,
                         "details": details,
                         "qatc_model": qatc_model_name,
@@ -249,9 +320,10 @@ with st.container():
                         "tc_model": tc_model_name
                     }
 
-                    # Thêm kết quả mới vào lịch sử
+                    # Add new result to history
                     st.session_state.history.append(st.session_state.latest_result)
 
+                    # Clear GPU cache to free memory
                     if torch.cuda.is_available():
                         torch.cuda.empty_cache()
 
@@ -264,8 +336,10 @@ with st.container():
                             <p class='verdict'><span class='verdict-icon'>{verdict_icons.get(res['verdict'], '')}</span>{res['verdict']}</p>
                             <p><strong>Evidence Inference Time:</strong> {res['evidence_time']:.2f} seconds</p>
                             <p><strong>Verdict Inference Time:</strong> {res['verdict_time']:.2f} seconds</p>
+                            <p><strong>Total Execution Time:</strong> {res['total_time']:.2f} seconds</p>
                         </div>
                     """, unsafe_allow_html=True)
+                    
                     # Download Verification Result Feature
                     result_text = f"Claim: {res['claim']}\nEvidence: {res['evidence']}\nVerdict: {res['verdict']}\nDetails: {res['details']}"
                     st.download_button("Download Result", data=result_text, file_name="verification_result.txt", mime="text/plain")