Update app.py

app.py

@@ -12,12 +12,14 @@ import psutil
 import gc
 import threading
 from queue import Queue
-from concurrent.futures import ThreadPoolExecutor
 
 # Set environment variables to optimize CPU performance
 os.environ["OMP_NUM_THREADS"] = str(psutil.cpu_count(logical=False))
 os.environ["MKL_NUM_THREADS"] = str(psutil.cpu_count(logical=False))
 
+# Set device globally
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Load models with caching and CPU optimization
 @st.cache_resource()
 def load_model(model_name, model_class, is_bc=False, device=None):
@@ -43,9 +45,6 @@ def load_model(model_name, model_class, is_bc=False, device=None):
     model.to(device)
     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):
@@ -271,41 +270,19 @@ with st.container():
         os.environ["OMP_NUM_THREADS"] = str(num_threads)
         os.environ["MKL_NUM_THREADS"] = str(num_threads)
 
-    # Store verification history
-    if 'history' not in st.session_state:
-        st.session_state.history = []
-    if 'latest_result' not in st.session_state:
-        st.session_state.latest_result = None
+    # Load models once and keep them in memory
     if 'models_loaded' not in st.session_state:
-        st.session_state.models_loaded = False
-
-    # 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..."):
-            # Clear memory before loading new models
-            gc.collect()
-            if DEVICE == "cpu":
-                torch.set_num_threads(num_threads)
-                
             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)
+
+    # Store verification history
+    if 'history' not in st.session_state:
+        st.session_state.history = []
+    if 'latest_result' not in st.session_state:
+        st.session_state.latest_result = None
 
     # Icons for results
     verdict_icons = {