Update app.py

app.py

@@ -1,142 +1,210 @@
-import os
+# Optimized RAG System with E5-Mistral Embeddings and Llama3-70B Generation
 import json
-import pandas as pd
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from langchain_community.vectorstores import FAISS
-from langchain_core.prompts import PromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from operator import itemgetter
+import logging
+import re
+from typing import List, Tuple
 import gradio as gr
-from langchain_community.embeddings import HuggingFaceEmbeddings
-
-# Configuration
-USE_HF = True
-MODEL_NAME = "stanford-crfm/BioMedLM"
-BATCH_SIZE = 8  # Adjusted batch size for memory optimization
-
-# Load data
-with open('AskNatureNet_data.json', 'r', encoding='utf-8') as f:
-    data = json.load(f)
-df = pd.DataFrame(data)
-documents = [
-    f"Source: {item['Source']}\nApplication: {item['Application']}\nFunction1: {item['Function1']}\nStrategy: {item['Strategy']}"
-    for item in data
-]
-
-if USE_HF:
-    print("Using Hugging Face model...")
-
-    huggingface_token = os.environ.get("AskNature_RAG")
-
-    # Quantization configuration for 4-bit precision
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4"
-    )
-
-    # Load tokenizer and model with offloading and quantization
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=huggingface_token)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map="cpu",
-        offload_folder="offload",  # Specify the offload folder
-        quantization_config=bnb_config,
-        use_auth_token=huggingface_token
-    )
-    embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
-    lang_model = model
-else:
-    print("Using local model...")
-    # Local model loading logic here
-
-# Generate embeddings in batches
-batched_embeddings = [
-    embeddings.embed_documents(documents[i:i + BATCH_SIZE])
-    for i in range(0, len(documents), BATCH_SIZE)
-]
-batched_embeddings = [embed for batch in batched_embeddings for embed in batch]
-
-# FAISS index handling
-index_path = "faiss_index"
-if os.path.exists(index_path):
-    vectorstore = FAISS.load_local(index_path, embeddings)
-else:
-    vectorstore = FAISS.from_texts(documents, embeddings)
-    vectorstore.save_local(index_path)
-
-retriever = vectorstore.as_retriever()
-
-# Prompt template
-template = """
-Answer the question based on the context below. If you can't 
-answer the question, reply "I don't know".
-Context: {context}
-Question: {question}
-"""
-prompt = PromptTemplate.from_template(template)
-
-# Chain definition
-chain = {
-    "context": itemgetter("question") | retriever,
-    "question": itemgetter("question"),
-} | prompt | lang_model | StrOutputParser()
-
-# Question-answering function
-def rag_qa(question):
+from openai import OpenAI
+from functools import lru_cache
+from tenacity import retry, stop_after_attempt, wait_exponential
+from langchain_community.retrievers import BM25Retriever
+from langchain_community.vectorstores import FAISS
+from langchain_core.embeddings import Embeddings
+from langchain_core.documents import Document
+from collections import defaultdict
+
+embedding_model = "e5-mistral-7b-instruct"
+generation_model = "meta-llama-3-70b-instruct"
+# --- Configuration ---
+API_CONFIG = {
+    "api_key": "d9960fad1d2aaa16167902b0d26e369f",
+    "base_url": "https://chat-ai.academiccloud.de/v1"
+}
+CHUNK_SIZE = 800
+OVERLAP = 200
+
+# Initialize clients
+client = OpenAI(**API_CONFIG)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# --- Custom Embedding Handler ---
+class MistralEmbeddings(Embeddings):
+    """E5-Mistral-7B embedding adapter with error handling"""
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        try:
+            response = client.embeddings.create(
+                input=texts,
+                model=embedding_model,
+                encoding_format="float"
+            )
+            return [e.embedding for e in response.data]
+        except Exception as e:
+            logger.error(f"Embedding Error: {str(e)}")
+            return [[] for _ in texts]  # Return empty embeddings on failure
+
+    def embed_query(self, text: str) -> List[float]:
+        return self.embed_documents([text])[0]
+
+# --- Data Processing ---
+def load_and_chunk_data(file_path: str) -> List[Document]:
+    """Enhanced chunking with metadata preservation"""
+    with open(file_path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+    
+    documents = []
+    for item in data:
+        base_content = f"""Source: {item['Source']}
+Application: {item['Application']}
+Functions: {', '.join(filter(None, [item.get('Function1'), item.get('Function2')]))}
+Technical Concepts: {', '.join(item['technical_concepts'])}
+Biological Mechanisms: {', '.join(item['biological_mechanisms'])}"""
+        
+        strategy = item['Strategy']
+        for i in range(0, len(strategy), CHUNK_SIZE - OVERLAP):
+            chunk = strategy[i:i + CHUNK_SIZE]
+            documents.append(Document(
+                page_content=f"{base_content}\nStrategy Excerpt:\n{chunk}",
+                metadata={
+                    "source": item["Source"],
+                    "application": item["Application"],
+                    "technical_concepts": item["technical_concepts"],
+                    "sustainability_impacts": item["sustainability_impacts"],
+                    "hyperlink": item["Hyperlink"],
+                    "chunk_id": f"{item['Source']}-{len(documents)+1}"
+                }
+            ))
+    return documents
+
+# --- Hybrid Retrieval System ---
+class EnhancedRetriever:
+    """BM25 + E5-Mistral embeddings with fusion"""
+    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.vector_retriever = self.vector_store.as_retriever(search_kwargs={"k": 3})
+
+    @lru_cache(maxsize=200)
+    def retrieve(self, query: str) -> str:
+        try:
+            processed_query = self._preprocess_query(query)
+            expanded_query = self._hyde_expansion(processed_query)
+            
+            bm25_results = self.bm25.invoke(processed_query)
+            vector_results = self.vector_retriever.invoke(processed_query)
+            expanded_results = self.bm25.invoke(expanded_query)
+            
+            fused_results = self._fuse_results([bm25_results, vector_results, expanded_results])
+            return self._format_context(fused_results[:5])
+        except Exception as e:
+            logger.error(f"Retrieval Error: {str(e)}")
+            return ""
+
+    def _preprocess_query(self, query: str) -> str:
+        return query.lower().strip()
+
+    def _hyde_expansion(self, query: str) -> str:
+        try:
+            response = client.chat.completions.create(
+                model=generation_model,
+                messages=[{
+                    "role": "user",
+                    "content": f"Generate a technical draft about biomimicry for: {query}\nInclude domain-specific terms."
+                }],
+                temperature=0.5,
+                max_tokens=200
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            logger.error(f"HyDE Error: {str(e)}")
+            return query
+
+    def _fuse_results(self, result_sets: List[List[Document]]) -> List[Document]:
+        fused_scores = defaultdict(float)
+        for docs in result_sets:
+            for rank, doc in enumerate(docs, 1):
+                fused_scores[doc.metadata["chunk_id"]] += 1 / (rank + 60)
+        
+        seen = set()
+        return [
+            doc for doc in sorted(
+                (doc for docs in result_sets for doc in docs),
+                key=lambda x: fused_scores[x.metadata["chunk_id"]],
+                reverse=True
+            ) if not (doc.metadata["chunk_id"] in seen or seen.add(doc.metadata["chunk_id"]))
+        ]
+
+    def _format_context(self, docs: List[Document]) -> str:
+        context = []
+        for doc in docs:
+            context_str = f"""**Source**: {doc.metadata['source']}
+**Application**: {doc.metadata['application']}
+**Concepts**: {', '.join(doc.metadata['technical_concepts'])}
+**Excerpt**: {doc.page_content.split('Strategy Excerpt:')[-1].strip()}
+**Reference**: {doc.metadata['hyperlink']}"""
+            context.append(context_str)
+        return "\n\n---\n\n".join(context)
+
+# --- Generation System ---
+SYSTEM_PROMPT = """**Biomimicry Expert Guidelines**
+1. Base answers strictly on context
+2. Cite sources as [Source]
+3. **Bold** technical terms
+4. Include reference links
+
+Context: {context}"""
+
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=20))
+def get_ai_response(query: str, context: str) -> str:
+    try:
+        response = client.chat.completions.create(
+            model=generation_model,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT.format(context=context)},
+                {"role": "user", "content": f"Question: {query}\nProvide a detailed technical answer:"}
+            ],
+            temperature=0.4,
+            max_tokens=600
+        )
+        return _postprocess_response(response.choices[0].message.content)
+    except Exception as e:
+        logger.error(f"Generation Error: {str(e)}")
+        return "I'm unable to generate a response right now. Please try again later."
+
+def _postprocess_response(response: str) -> str:
+    response = re.sub(r"\[(.*?)\]", r"[\1](#)", response)
+    response = re.sub(r"\*\*([\w-]+)\*\*", r"**\1**", response)
+    return response
+
+# --- Pipeline Integration ---
+documents = load_and_chunk_data("mini_data_enhanced.json")
+retriever = EnhancedRetriever(documents)
+
+def generate_response(question: str) -> str:
     try:
-        return chain.invoke({'question': question})
+        context = retriever.retrieve(question)
+        return get_ai_response(question, context) if context else "No relevant information found."
     except Exception as e:
-        return f"Error: {str(e)}"
-
-# Gradio chatbot interface
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
-
-# Gradio interface setup
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+        logger.error(f"Pipeline Error: {str(e)}")
+        return "An error occurred processing your request."
+
+# --- Gradio Interface ---
+def chat_interface(question: str, history: List[Tuple[str, str]]):
+    response = generate_response(question)
+    return "", history + [(question, response)]
+
+with gr.Blocks(title="BioRAG Expert", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌿 BioRAG: Biomimicry Technical Assistant")
+    with gr.Row():
+        chatbot = gr.Chatbot(label="Dialogue History", height=500)
+    with gr.Row():
+        question = gr.Textbox(placeholder="Ask about nature-inspired innovations...", 
+                            label="Technical Inquiry", scale=4)
+        clear_btn = gr.Button("Clear History", variant="secondary")
+    
+    question.submit(chat_interface, [question, chatbot], [question, chatbot])
+    clear_btn.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(show_error=True)