Update rag_engine.py

rag_engine.py

@@ -11,6 +11,7 @@ import textwrap
 import unicodedata
 import streamlit as st
 from utils import setup_gcp_auth, setup_openai_auth
+import gc  # Added for explicit garbage collection
 
 # Force model to CPU for stability
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -22,6 +23,7 @@ def initialize_session_state():
         st.session_state.model = None
         st.session_state.tokenizer = None
         st.session_state.device = torch.device("cpu")
+        st.session_state.data_loaded = False
         print("Initialized session state variables")
 
 # Call the initialization function right away
@@ -38,7 +40,6 @@ def setup_gcp_client():
         return bucket
     except Exception as e:
         print(f"❌ GCP client initialization error: {str(e)}")
-        st.error(f"GCP client initialization error: {str(e)}")
         return None
 
 # Setup OpenAI authentication
@@ -49,7 +50,6 @@ def setup_openai_client():
         return True
     except Exception as e:
         print(f"❌ OpenAI client initialization error: {str(e)}")
-        st.error(f"OpenAI client initialization error: {str(e)}")
         return False
 
 # GCS Paths
@@ -66,32 +66,36 @@ local_metadata_file = "metadata.jsonl"
 
 def load_model():
     try:
-        # Initialize model if it doesn't exist
-        if 'model' not in st.session_state or 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 CPU explicitly
-            model = model.to('cpu')
-            model.eval()
-            torch.set_grad_enabled(False)
-            
-            # Store in session state
-            st.session_state.tokenizer = tokenizer
-            st.session_state.model = model
-            
-            print("✅ Model loaded successfully")
+        # Check if model is already loaded
+        if st.session_state.model is not None and st.session_state.tokenizer is not None:
+            print("Model already loaded, reusing existing instance")
+            return st.session_state.tokenizer, st.session_state.model
+
+        # 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 CPU explicitly
+        model = model.to('cpu')
+        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")
         
-        return st.session_state.tokenizer, st.session_state.model
+        return tokenizer, model
     except Exception as e:
         print(f"❌ Error loading model: {str(e)}")
         # Return None values instead of raising to avoid crashing
@@ -100,32 +104,41 @@ def load_model():
 def download_file_from_gcs(bucket, gcs_path, local_path):
     """Download a file from GCS to local storage."""
     try:
+        # Check if file already exists locally
+        if os.path.exists(local_path):
+            print(f"File already exists locally: {local_path}")
+            return True
+            
         blob = bucket.blob(gcs_path)
         blob.download_to_filename(local_path)
         print(f"✅ Downloaded {gcs_path} → {local_path}")
         return True
     except Exception as e:
         print(f"❌ Error downloading {gcs_path}: {str(e)}")
-        st.error(f"Error downloading {gcs_path}: {str(e)}")
         return False
 
 def load_data_files():
+    # Check if already loaded in session state
+    if hasattr(st.session_state, 'faiss_index') and st.session_state.faiss_index is not None:
+        print("Using cached data files from session state")
+        return st.session_state.faiss_index, st.session_state.text_chunks, st.session_state.metadata_dict
+    
     # Initialize GCP and OpenAI clients
     bucket = setup_gcp_client()
     openai_initialized = setup_openai_client()
     
     if not bucket or not openai_initialized:
-        st.error("Failed to initialize required services")
+        print("Failed to initialize required services")
         return None, None, None
     
-    # Download necessary files - remove the spinner from here
+    # Download necessary files
     success = True
     success &= download_file_from_gcs(bucket, faiss_index_file_gcs, local_faiss_index_file)
     success &= download_file_from_gcs(bucket, text_chunks_file_gcs, local_text_chunks_file)
     success &= download_file_from_gcs(bucket, metadata_file_gcs, local_metadata_file)
     
     if not success:
-        st.error("Failed to download required files")
+        print("Failed to download required files")
         return None, None, None
     
     # Load FAISS index
@@ -133,7 +146,6 @@ def load_data_files():
         faiss_index = faiss.read_index(local_faiss_index_file)
     except Exception as e:
         print(f"❌ Error loading FAISS index: {str(e)}")
-        st.error(f"Error loading FAISS index: {str(e)}")
         return None, None, None
     
     # Load text chunks
@@ -146,7 +158,6 @@ def load_data_files():
                     text_chunks[int(parts[0])] = (parts[1], parts[2], parts[3])
     except Exception as e:
         print(f"❌ Error loading text chunks: {str(e)}")
-        st.error(f"Error loading text chunks: {str(e)}")
         return None, None, None
     
     # Load metadata.jsonl for publisher information
@@ -158,10 +169,16 @@ def load_data_files():
                 metadata_dict[item["Title"]] = item  # Store for easy lookup
     except Exception as e:
         print(f"❌ Error loading metadata: {str(e)}")
-        st.error(f"Error loading metadata: {str(e)}")
         return None, None, None
     
     print(f"✅ FAISS index and text chunks loaded. {len(text_chunks)} passages available.")
+    
+    # Store in session state
+    st.session_state.faiss_index = faiss_index
+    st.session_state.text_chunks = text_chunks
+    st.session_state.metadata_dict = metadata_dict
+    st.session_state.data_loaded = True
+    
     return faiss_index, text_chunks, metadata_dict
 
 def average_pool(last_hidden_states, attention_mask):
@@ -177,16 +194,17 @@ def get_embedding(text):
 
     try:
         # Ensure model initialization
-        if 'model' not in st.session_state or st.session_state.model is None:
+        if not hasattr(st.session_state, 'model') or st.session_state.model is None:
             tokenizer, model = load_model()
             if model is None:
-                return np.zeros((1, 384), dtype=np.float32)  # Fallback for e5-small-v2
+                return np.zeros((1, 384), dtype=np.float32)
         else:
             tokenizer, model = st.session_state.tokenizer, st.session_state.model
             
+        # Prepare text
         input_text = f"query: {text}" if len(text) < 512 else f"passage: {text}"
         
-        # Explicitly specify truncation parameters to avoid warnings
+        # Explicitly specify truncation parameters
         inputs = tokenizer(
             input_text,
             padding=True,
@@ -196,26 +214,25 @@ def get_embedding(text):
             return_attention_mask=True
         )
         
-        # Move to CPU explicitly before processing
+        # Move to CPU explicitly
         inputs = {k: v.to('cpu') for k, v in inputs.items()}
         
         with torch.no_grad():
             outputs = model(**inputs)
             embeddings = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
             embeddings = nn.functional.normalize(embeddings, p=2, dim=1)
-            # Ensure we detach and move to numpy on CPU
             embeddings = embeddings.detach().cpu().numpy()
             
         # Explicitly clean up
-        del outputs
+        del outputs, inputs
+        gc.collect()
         torch.cuda.empty_cache() if torch.cuda.is_available() else None
         
         query_embedding_cache[text] = embeddings
         return embeddings
     except Exception as e:
         print(f"❌ Embedding error: {str(e)}")
-        st.error(f"Embedding error: {str(e)}")
-        return np.zeros((1, 384), dtype=np.float32)  # Changed from 1024 to 384 for e5-small-v2
+        return np.zeros((1, 384), dtype=np.float32)
 
 def retrieve_passages(query, faiss_index, text_chunks, metadata_dict, top_k=5, similarity_threshold=0.5):
     """Retrieve top-k most relevant passages using FAISS with metadata."""
@@ -242,6 +259,7 @@ def retrieve_passages(query, faiss_index, text_chunks, metadata_dict, top_k=5, s
                 if clean_title in cited_titles:
                     continue  
 
+                # Get metadata safely
                 metadata_entry = metadata_dict.get(clean_title, {})
                 author = metadata_entry.get("Author", "Unknown")
                 publisher = metadata_entry.get("Publisher", "Unknown")
@@ -258,7 +276,6 @@ def retrieve_passages(query, faiss_index, text_chunks, metadata_dict, top_k=5, s
         return retrieved_passages, retrieved_sources
     except Exception as e:
         print(f"❌ Error in retrieve_passages: {str(e)}")
-        st.error(f"Error in retrieve_passages: {str(e)}")
         return [], []
 
 def answer_with_llm(query, context=None, word_limit=100):
@@ -325,7 +342,6 @@ def answer_with_llm(query, context=None, word_limit=100):
 
     except Exception as e:
         print(f"❌ LLM API error: {str(e)}")
-        st.error(f"LLM API error: {str(e)}")
         return "I apologize, but I'm unable to answer at the moment."
 
 def format_citations(sources):
@@ -345,19 +361,21 @@ def process_query(query, top_k=5, word_limit=100):
     print(f"\n🔍 Processing query: {query}")
     
     # Load data files if not already loaded
-    if not hasattr(st.session_state, 'data_loaded') or not st.session_state.data_loaded:
-        st.session_state.faiss_index, st.session_state.text_chunks, st.session_state.metadata_dict = load_data_files()
-        st.session_state.data_loaded = True
+    faiss_index, text_chunks, metadata_dict = load_data_files()
     
     # Check if data loaded successfully
-    if not st.session_state.faiss_index or not st.session_state.text_chunks or not st.session_state.metadata_dict:
-        return {"query": query, "answer_with_rag": "⚠️ System error: Data files not loaded properly.", "citations": "No citations available."}
+    if faiss_index is None or text_chunks is None or metadata_dict is None:
+        return {
+            "query": query, 
+            "answer_with_rag": "⚠️ System error: Data files not loaded properly.", 
+            "citations": "No citations available."
+        }
 
     retrieved_context, retrieved_sources = retrieve_passages(
         query, 
-        st.session_state.faiss_index, 
-        st.session_state.text_chunks, 
-        st.session_state.metadata_dict,
+        faiss_index, 
+        text_chunks, 
+        metadata_dict,
         top_k=top_k
     )