Update medical_rag_system.py

medical_rag_system.py

@@ -1,13 +1,15 @@
 import faiss
 import pickle
-from sentence_transformers import SentenceTransformer
 import numpy as np
 import os
+import requests
+import json
+from typing import List, Dict, Any
+import base64
 
 class MedicalRAGSystem:
     def __init__(self, vector_store_path="medical_rag/vector_store"):
         self.vector_store_path = vector_store_path
-        self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
         self.load_vector_store()
     
     def load_vector_store(self):
@@ -18,7 +20,7 @@ class MedicalRAGSystem:
                 self.vector_store_path,
                 "medical_rag/vector_store",
                 "./medical_rag/vector_store",
-                "/app/medical_rag/vector_store",  # For Hugging Face Spaces
+                "/app/medical_rag/vector_store",
             ]
             
             loaded = False
@@ -31,6 +33,15 @@ class MedicalRAGSystem:
                         self.index = faiss.read_index(index_path)
                         with open(metadata_path, "rb") as f:
                             data = pickle.load(f)
+                        
+                        # Load pre-computed embeddings
+                        embeddings_path = f"{path}/embeddings.npy"
+                        if os.path.exists(embeddings_path):
+                            self.embeddings = np.load(embeddings_path)
+                        else:
+                            # If no embeddings file, we'll use the index directly
+                            self.embeddings = None
+                        
                         self.chunks = data['chunks']
                         self.metadata = data['metadata']
                         print(f"✅ Vector store loaded from: {path}")
@@ -41,23 +52,70 @@ class MedicalRAGSystem:
                     continue
             
             if not loaded:
-                raise Exception("Could not load vector store from any path. Check if vector store files exist.")
+                raise Exception("Could not load vector store from any path.")
             
         except Exception as e:
             raise Exception(f"Error loading vector store: {str(e)}")
     
-    def retrieve_similar_chunks(self, query, k=5):
+    def get_embedding_huggingface(self, text: str) -> np.ndarray:
+        """Get embedding using Hugging Face Inference API"""
+        try:
+            API_URL = "https://api-inference.huggingface.co/pipeline/feature-extraction/sentence-transformers/all-MiniLM-L6-v2"
+            headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN', '')}"}
+            
+            response = requests.post(
+                API_URL,
+                headers=headers,
+                json={"inputs": text, "options": {"wait_for_model": True}}
+            )
+            
+            if response.status_code == 200:
+                return np.array(response.json())
+            else:
+                # Fallback to simple TF-IDF like approach
+                print(f"Hugging Face API failed: {response.status_code}")
+                return self.simple_embedding(text)
+                
+        except Exception as e:
+            print(f"Embedding error: {e}")
+            return self.simple_embedding(text)
+    
+    def simple_embedding(self, text: str) -> np.ndarray:
+        """Simple fallback embedding (TF-IDF like)"""
+        # Convert text to lowercase and split
+        words = text.lower().split()
+        
+        # Create a simple bag-of-words vector (384 dimensions to match MiniLM)
+        embedding = np.zeros(384)
+        
+        # Simple hash-based word distribution
+        for word in words:
+            # Create a simple hash-based position
+            hash_val = hash(word) % 384
+            embedding[hash_val] += 1
+        
+        # Normalize
+        norm = np.linalg.norm(embedding)
+        if norm > 0:
+            embedding = embedding / norm
+        
+        return embedding
+    
+    def retrieve_similar_chunks(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
         """Retrieve similar medical chunks for a query"""
         try:
-            # Encode query
-            query_embedding = self.embedding_model.encode([query])
+            # Get query embedding
+            query_embedding = self.get_embedding_huggingface(query).astype('float32')
+            query_embedding = query_embedding.reshape(1, -1)
+            
+            # Normalize for cosine similarity
             faiss.normalize_L2(query_embedding)
             
-            # Search (get extra for potential filtering)
-            scores, indices = self.index.search(query_embedding, k*3)
+            # Search
+            scores, indices = self.index.search(query_embedding, min(k * 3, len(self.chunks)))
             
             results = []
-            seen_chunks = set()  # Avoid duplicates
+            seen_chunks = set()
             
             for score, idx in zip(scores[0], indices[0]):
                 if idx < len(self.chunks) and idx not in seen_chunks: