Update app.py

app.py

@@ -2,7 +2,9 @@
 import json
 import logging
 import re
-from typing import List, Tuple
+import os
+import pickle
+from typing import List, Tuple, Optional
 import gradio as gr
 from openai import OpenAI
 from functools import lru_cache
@@ -12,51 +14,74 @@ from langchain_community.vectorstores import FAISS
 from langchain_core.embeddings import Embeddings
 from langchain_core.documents import Document
 from collections import defaultdict
-import os
-
-
+import hashlib
+from tqdm import tqdm  
 
+# --- Configuration ---
+FAISS_INDEX_PATH = "faiss_index"
+BM25_INDEX_PATH = "bm25_index.pkl"
+CACHE_VERSION = "v1"  # Increment when data format changes
 embedding_model = "e5-mistral-7b-instruct"
 generation_model = "meta-llama-3-70b-instruct"
-# --- Configuration ---
+data_file_name = "AskNatureNet_data_enhanced.json"
 API_CONFIG = {
     "api_key": os.getenv("API_KEY"),
     "base_url": "https://chat-ai.academiccloud.de/v1"
 }
 CHUNK_SIZE = 800
 OVERLAP = 200
+EMBEDDING_BATCH_SIZE = 32  # Batch size for embedding API calls
 
 # Initialize clients
 client = OpenAI(**API_CONFIG)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- Custom Embedding Handler ---
+# --- Helper Functions ---
+def get_data_hash(file_path: str) -> str:
+    """Generate hash of data file for cache validation"""
+    with open(file_path, "rb") as f:
+        return hashlib.md5(f.read()).hexdigest()
+
+# --- Custom Embedding Handler with Progress Tracking ---
 class MistralEmbeddings(Embeddings):
-    """E5-Mistral-7B embedding adapter with error handling"""
+    """E5-Mistral-7B embedding adapter with error handling and progress tracking"""
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        embeddings = []
         try:
-            response = client.embeddings.create(
-                input=texts,
-                model=embedding_model,
-                encoding_format="float"
-            )
-            return [e.embedding for e in response.data]
+            # Process in batches with progress tracking
+            for i in tqdm(range(0, len(texts), EMBEDDING_BATCH_SIZE), desc="Embedding Progress"):
+                batch = texts[i:i + EMBEDDING_BATCH_SIZE]
+                response = client.embeddings.create(
+                    input=batch,
+                    model=embedding_model,
+                    encoding_format="float"
+                )
+                embeddings.extend([e.embedding for e in response.data])
+            return embeddings
         except Exception as e:
             logger.error(f"Embedding Error: {str(e)}")
-            return [[] for _ in texts]  # Return empty embeddings on failure
+            return [[] for _ in texts]
 
     def embed_query(self, text: str) -> List[float]:
         return self.embed_documents([text])[0]
 
-# --- Data Processing ---
+# --- Data Processing with Cache Validation ---
 def load_and_chunk_data(file_path: str) -> List[Document]:
     """Enhanced chunking with metadata preservation"""
+    current_hash = get_data_hash(file_path)
+    cache_file = f"documents_{CACHE_VERSION}_{current_hash}.pkl"
+    
+    if os.path.exists(cache_file):
+        logger.info("Loading cached documents")
+        with open(cache_file, "rb") as f:
+            return pickle.load(f)
+    
     with open(file_path, 'r', encoding='utf-8') as f:
         data = json.load(f)
     
     documents = []
-    for item in data:
+    for item in tqdm(data, desc="Chunking Progress"):
         base_content = f"""Source: {item['Source']}
 Application: {item['Application']}
 Functions: {', '.join(filter(None, [item.get('Function1'), item.get('Function2')]))}
@@ -77,18 +102,50 @@ Biological Mechanisms: {', '.join(item['biological_mechanisms'])}"""
                     "chunk_id": f"{item['Source']}-{len(documents)+1}"
                 }
             ))
+    
+    with open(cache_file, "wb") as f:
+        pickle.dump(documents, f)
     return documents
 
-# --- Hybrid Retrieval System ---
+# --- Optimized Retrieval System ---
 class EnhancedRetriever:
-    """BM25 + E5-Mistral embeddings with fusion"""
+    """Hybrid retriever with persistent caching"""
     def __init__(self, documents: List[Document]):
-        self.bm25 = BM25Retriever.from_documents(documents)
-        self.bm25.k = 5
-        self.vector_store = FAISS.from_documents(documents, MistralEmbeddings())
+        self.documents = documents
+        self.bm25 = self._init_bm25()
+        self.vector_store = self._init_faiss()
         self.vector_retriever = self.vector_store.as_retriever(search_kwargs={"k": 3})
 
-    @lru_cache(maxsize=200)
+    def _init_bm25(self) -> BM25Retriever:
+        cache_key = f"{BM25_INDEX_PATH}_{get_data_hash(data_file_name)}"
+        if os.path.exists(cache_key):
+            logger.info("Loading cached BM25 index")
+            with open(cache_key, "rb") as f:
+                return pickle.load(f)
+        
+        logger.info("Building new BM25 index")
+        retriever = BM25Retriever.from_documents(self.documents)
+        retriever.k = 5
+        with open(cache_key, "wb") as f:
+            pickle.dump(retriever, f)
+        return retriever
+
+    def _init_faiss(self) -> FAISS:
+        cache_key = f"{FAISS_INDEX_PATH}_{get_data_hash(data_file_name)}"
+        if os.path.exists(cache_key):
+            logger.info("Loading cached FAISS index")
+            return FAISS.load_local(
+                cache_key,
+                MistralEmbeddings(),
+                allow_dangerous_deserialization=True
+            )
+        
+        logger.info("Building new FAISS index")
+        vector_store = FAISS.from_documents(self.documents, MistralEmbeddings())
+        vector_store.save_local(cache_key)
+        return vector_store
+
+    @lru_cache(maxsize=500)
     def retrieve(self, query: str) -> str:
         try:
             processed_query = self._preprocess_query(query)
@@ -107,6 +164,7 @@ class EnhancedRetriever:
     def _preprocess_query(self, query: str) -> str:
         return query.lower().strip()
 
+    @lru_cache(maxsize=500)
     def _hyde_expansion(self, query: str) -> str:
         try:
             response = client.chat.completions.create(
@@ -167,8 +225,9 @@ def get_ai_response(query: str, context: str) -> str:
                 {"role": "user", "content": f"Question: {query}\nProvide a detailed technical answer:"}
             ],
             temperature=0.4,
-            max_tokens=600
+            max_tokens=2000  # Increased max_tokens
         )
+        logger.info(f"Raw Response: {response.choices[0].message.content}")  # Log raw response
         return _postprocess_response(response.choices[0].message.content)
     except Exception as e:
         logger.error(f"Generation Error: {str(e)}")
@@ -179,8 +238,8 @@ def _postprocess_response(response: str) -> str:
     response = re.sub(r"\*\*([\w-]+)\*\*", r"**\1**", response)
     return response
 
-# --- Pipeline Integration ---
-documents = load_and_chunk_data("mini_data_enhanced.json")
+# --- Optimized Pipeline ---
+documents = load_and_chunk_data(data_file_name)
 retriever = EnhancedRetriever(documents)
 
 def generate_response(question: str) -> str: