Update rag_engine.py

rag_engine.py

@@ -66,7 +66,8 @@ local_metadata_file = "metadata.jsonl"
 
 def load_model():
     try:
-        if st.session_state.model is None:
+        # 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"] = ""
             
@@ -79,20 +80,22 @@ def load_model():
                 torch_dtype=torch.float16  # Use half precision
             )
             
-            # Move model to the designated device
-            model = model.to(st.session_state.device)
+            # 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")
-            
+        
         return st.session_state.tokenizer, st.session_state.model
     except Exception as e:
         print(f"❌ Error loading model: {str(e)}")
-        raise
+        # Return None values instead of raising to avoid crashing
+        return None, None
 
 def download_file_from_gcs(bucket, gcs_path, local_path):
     """Download a file from GCS to local storage."""
@@ -172,41 +175,18 @@ query_embedding_cache = {}
 def get_embedding(text):
     if text in query_embedding_cache:
         return query_embedding_cache[text]
-
     try:
-        tokenizer, model = load_model()
-        input_text = f"query: {text}" if len(text) < 512 else f"passage: {text}"
-        
-        # Explicitly specify truncation parameters to avoid warnings
-        inputs = tokenizer(
-            input_text,
-            padding=True,
-            truncation=True,
-            return_tensors="pt",
-            max_length=512,
-            return_attention_mask=True
-        )
-        
-        # Move to CPU explicitly before processing
-        inputs = {k: v.to('cpu') for k, v in inputs.items()}
+        # Ensure model initialization
+        if 'model' not in st.session_state or st.session_state.model is None:
+            tokenizer, model = load_model()
+            if model is None:
+                return np.zeros((1, 384), dtype=np.float32)  # Fallback
+        else:
+            tokenizer, model = st.session_state.tokenizer, st.session_state.model
         
-        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
-        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        input_text = f"query: {text}" if len(text) < 512 else f"passage: {text}"
         
-        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
+        # Rest of your code...
 
 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."""