Update rag_engine.py

rag_engine.py

@@ -175,18 +175,48 @@ query_embedding_cache = {}
 def get_embedding(text):
     if text in query_embedding_cache:
         return query_embedding_cache[text]
+
     try:
         # 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
+                return np.zeros((1, 384), dtype=np.float32)  # Fallback for e5-small-v2
         else:
             tokenizer, model = st.session_state.tokenizer, st.session_state.model
-        
+            
         input_text = f"query: {text}" if len(text) < 512 else f"passage: {text}"
         
-        # Rest of your code...
+        # 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()}
+        
+        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
+        
+        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
 
 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."""