Update rag_engine.py

rag_engine.py

@@ -66,33 +66,32 @@ local_metadata_file = "metadata.jsonl"
 
 def load_model():
     try:
-        # Check if model is already loaded
-        if not st.session_state.model_initialized:
-            with st.spinner("Loading tokenizer and model... This may take a minute."):
-                print("Loading tokenizer...")
-                tokenizer = AutoTokenizer.from_pretrained("intfloat/e5-small-v2")
-                
-                print("Loading model...")
-                model = AutoModel.from_pretrained(
-                    "intfloat/e5-small-v2",
-                    torch_dtype=torch.float16,  # Use half precision
-                    low_cpu_mem_usage=True,
-                )
-                
-                model.eval()
-                torch.set_grad_enabled(False)
-                
-                # Store in session state
-                st.session_state.tokenizer = tokenizer
-                st.session_state.model = model
-                st.session_state.model_initialized = True
-                
-                print("✅ Model loaded successfully")
+        if st.session_state.model is None:
+            # Force model to CPU - more stable than GPU for this use case
+            os.environ["CUDA_VISIBLE_DEVICES"] = ""
+            
+            print("Loading tokenizer...")
+            tokenizer = AutoTokenizer.from_pretrained("intfloat/e5-small-v2")
+            
+            print("Loading model...")
+            model = AutoModel.from_pretrained(
+                "intfloat/e5-small-v2",
+                torch_dtype=torch.float16  # Use half precision
+            )
+            
+            # Move model to the designated device
+            model = model.to(st.session_state.device)
+            model.eval()
+            torch.set_grad_enabled(False)
+            
+            st.session_state.tokenizer = tokenizer
+            st.session_state.model = model
+            
+            print("✅ Model loaded successfully")
             
         return st.session_state.tokenizer, st.session_state.model
     except Exception as e:
         print(f"❌ Error loading model: {str(e)}")
-        st.error(f"Error loading model: {str(e)}")
         raise
 
 def download_file_from_gcs(bucket, gcs_path, local_path):