Update app.py

app.py

@@ -1,520 +1,521 @@
-import os
-import json
-import glob
-from pathlib import Path
-import torch
-import streamlit as st
-from dotenv import load_dotenv
-from langchain_groq import ChatGroq
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-from langchain_core.documents import Document
-from langchain_core.prompts import ChatPromptTemplate
-from langchain.chains import create_retrieval_chain
-from langchain.chains.combine_documents import create_stuff_documents_chain
-import numpy as np  
-from sentence_transformers import util
-import time
-
-# Set device for model (CUDA if available)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load environment variables - works for both local and Hugging Face Spaces
-load_dotenv()
-
-# Set up the clinical assistant LLM
-# Try to get API key from Hugging Face Spaces secrets first, then fall back to .env file
-try:
-    # For Hugging Face Spaces
-    from huggingface_hub.inference_api import InferenceApi
-    import os
-    groq_api_key = os.environ.get('GROQ_API_KEY')
-    
-    # If not found in environment, try to get from st.secrets (Streamlit Cloud/Spaces)
-    if not groq_api_key and hasattr(st, 'secrets') and 'GROQ_API_KEY' in st.secrets:
-        groq_api_key = st.secrets['GROQ_API_KEY']
-        
-    if not groq_api_key:
-        st.warning("API Key is not set in the secrets. Using a placeholder for UI demonstration.")
-        # For UI demonstration without API key
-        class MockLLM:
-            def invoke(self, prompt):
-                return {"answer": "This is a placeholder response. Please set up your GROQ_API_KEY to get real responses."}
-        llm = MockLLM()
-    else:    
-        llm = ChatGroq(groq_api_key=groq_api_key, model_name="llama-3.3-70b-versatile")
-        
-except Exception as e:
-    st.error(f"Error setting up LLM: {str(e)}")
-    class MockLLM:
-        def invoke(self, prompt):
-            return {"answer": f"Error setting up LLM: {str(e)}. Please check your API key configuration."}
-    llm = MockLLM()
-
-# Set up embeddings for clinical context (Bio_ClinicalBERT)
-embeddings = HuggingFaceEmbeddings(
-    model_name="emilyalsentzer/Bio_ClinicalBERT",
-    model_kwargs={"device": device}
-)
-
-def load_clinical_data():
-    """Load both flowcharts and patient cases"""
-    docs = []
-    
-    # Get the absolute path to the current script
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    
-    # Try to handle potential errors with file loading
-    try:
-        # Load diagnosis flowcharts
-        flowchart_dir = os.path.join(current_dir, "Diagnosis_flowchart")
-        if os.path.exists(flowchart_dir):
-            for fpath in glob.glob(os.path.join(flowchart_dir, "*.json")):
-                try:
-                    with open(fpath, 'r', encoding='utf-8') as f:
-                        data = json.load(f)
-                        content = f"""
-                        DIAGNOSTIC FLOWCHART: {Path(fpath).stem}
-                        Diagnostic Path: {data.get('diagnostic', 'N/A')}
-                        Key Criteria: {data.get('knowledge', 'N/A')}
-                        """
-                        docs.append(Document(
-                            page_content=content,
-                            metadata={"source": fpath, "type": "flowchart"}
-                        ))
-                except Exception as e:
-                    st.warning(f"Error loading flowchart file {fpath}: {str(e)}")
-        else:
-            st.warning(f"Flowchart directory not found at {flowchart_dir}")
-
-        # Load patient cases
-        finished_dir = os.path.join(current_dir, "Finished")
-        if os.path.exists(finished_dir):
-            for category_dir in glob.glob(os.path.join(finished_dir, "*")):
-                if os.path.isdir(category_dir):
-                    for case_file in glob.glob(os.path.join(category_dir, "*.json")):
-                        try:
-                            with open(case_file, 'r', encoding='utf-8') as f:
-                                case_data = json.load(f)
-                                notes = "\n".join(
-                                    f"{k}: {v}" for k, v in case_data.items() if k.startswith("input")
-                                )
-                                docs.append(Document(
-                                    page_content=f"""
-                                    PATIENT CASE: {Path(case_file).stem}
-                                    Category: {Path(category_dir).name}
-                                    Notes: {notes}
-                                    """,
-                                    metadata={"source": case_file, "type": "patient_case"}
-                                ))
-                        except Exception as e:
-                            st.warning(f"Error loading case file {case_file}: {str(e)}")
-        else:
-            st.warning(f"Finished directory not found at {finished_dir}")
-            
-        # If no documents were loaded, add a sample document for testing
-        if not docs:
-            st.warning("No clinical data files found. Using sample data for demonstration.")
-            docs.append(Document(
-                page_content="""SAMPLE CLINICAL DATA: This is sample data for demonstration purposes.
-                This application requires clinical data files to be present in the correct directories.
-                Please ensure the Diagnosis_flowchart and Finished directories exist with proper JSON files.""",
-                metadata={"source": "sample", "type": "sample"}
-            ))
-    except Exception as e:
-        st.error(f"Error loading clinical data: {str(e)}")
-        # Add a fallback document
-        docs.append(Document(
-            page_content="Error loading clinical data. This is a fallback document for demonstration purposes.",
-            metadata={"source": "error", "type": "error"}
-        ))
-    return docs
-
-def build_vectorstore():
-    """Build and return the vectorstore using FAISS"""
-    documents = load_clinical_data()
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    splits = splitter.split_documents(documents)
-    vectorstore = FAISS.from_documents(splits, embeddings)
-    return vectorstore
-
-# Path for saving/loading the vectorstore
-def get_vectorstore_path():
-    """Get the path for saving/loading the vectorstore"""
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    return os.path.join(current_dir, "vectorstore")
-
-# Initialize vectorstore with disk persistence
-@st.cache_resource(show_spinner="Loading clinical knowledge base...")
-def get_vectorstore():
-    """Get or create the vectorstore with disk persistence"""
-    vectorstore_path = get_vectorstore_path()
-    
-    # Try to load from disk first
-    try:
-        if os.path.exists(vectorstore_path):
-            st.info("Loading vectorstore from disk...")
-            # Set allow_dangerous_deserialization to True since we trust our own vectorstore files
-            return FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True)
-    except Exception as e:
-        st.warning(f"Could not load vectorstore from disk: {str(e)}. Building new vectorstore.")
-    
-    # If loading fails or doesn't exist, build a new one
-    st.info("Building new vectorstore...")
-    vectorstore = build_vectorstore()
-    
-    # Save to disk for future use
-    try:
-        os.makedirs(vectorstore_path, exist_ok=True)
-        vectorstore.save_local(vectorstore_path)
-        st.success("Vectorstore saved to disk for future use")
-    except Exception as e:
-        st.warning(f"Could not save vectorstore to disk: {str(e)}")
-    
-    return vectorstore
-
-def run_rag_chat(query, vectorstore):
-    """Run the Retrieval-Augmented Generation (RAG) for clinical questions"""
-    try:
-        retriever = vectorstore.as_retriever()
-
-        prompt_template = ChatPromptTemplate.from_template("""
-        You are a clinical assistant AI. Based on the following clinical context, provide a reasoned and medically sound answer to the question.
-
-        <context>
-        {context}
-        </context>
-
-        Question: {input}
-
-        Answer:
-        """)
-
-        retrieved_docs = retriever.invoke(query, k=3)
-        retrieved_context = "\n".join([doc.page_content for doc in retrieved_docs])
-
-        # Create document chain first
-        document_chain = create_stuff_documents_chain(llm, prompt_template)
-        
-        # Then create retrieval chain
-        chain = create_retrieval_chain(retriever, document_chain)
-
-        # Invoke the chain
-        response = chain.invoke({"input": query})
-        
-        # Add retrieved documents to response for transparency
-        response["context"] = retrieved_docs
-        
-        return response
-    except Exception as e:
-        st.error(f"Error in RAG processing: {str(e)}")
-        # Return a fallback response
-        return {
-            "answer": f"I encountered an error processing your query: {str(e)}",
-            "context": [],
-            "input": query
-        }
-
-def calculate_hit_rate(retriever, query, expected_docs, k=3):
-    """Calculate the hit rate for top-k retrieved documents"""
-    retrieved_docs = retriever.get_relevant_documents(query, k=k)
-    retrieved_contents = [doc.page_content for doc in retrieved_docs]
-    
-    hits = 0
-    for expected in expected_docs:
-        if any(expected in retrieved for retrieved in retrieved_contents):
-            hits += 1
-    
-    return hits / len(expected_docs) if expected_docs else 0.0
-
-def evaluate_rag_response(response, embeddings):
-    """Evaluate the RAG response for faithfulness and hit rate"""
-    scores = {}
-
-    # Faithfulness: Answer-Context Similarity
-    answer_embed = embeddings.embed_query(response["answer"])
-    context_embeds = [embeddings.embed_query(doc.page_content) for doc in response["context"]]
-    similarities = [util.cos_sim(answer_embed, ctx_embed).item() for ctx_embed in context_embeds]
-    scores["faithfulness"] = float(np.mean(similarities)) if similarities else 0.0
-
-    # Custom Hit Rate Calculation
-    retriever = response["retriever"]
-    scores["hit_rate"] = calculate_hit_rate(
-        retriever,
-        query=response["input"],
-        expected_docs=[doc.page_content for doc in response["context"]],
-        k=3
-    )
-    
-    return scores
-
-def main():
-    """Main function to run the Streamlit app"""
-    # Set page configuration
-    st.set_page_config(
-        page_title="DiReCT - Clinical Diagnostic Assistant",
-        page_icon="🩺",
-        layout="wide",
-        initial_sidebar_state="expanded"
-    )
-    
-    # Load vectorstore only once using session state
-    if 'vectorstore' not in st.session_state:
-        with st.spinner("Loading clinical knowledge base... This may take a minute."):
-            try:
-                st.session_state.vectorstore = get_vectorstore()
-                # Use custom styled message without the success icon
-                st.markdown("<div style='padding:10px 15px;background-color:rgba(40,167,69,0.2);border-radius:5px;border-left:5px solid rgba(40,167,69,0.8);'>Clinical knowledge base loaded successfully!</div>", unsafe_allow_html=True)
-            except Exception as e:
-                st.error(f"Error loading knowledge base: {str(e)}")
-                st.session_state.vectorstore = None
-
-    # Custom CSS for modern look with dark theme compatibility
-    st.markdown("""
-    <style>
-    .stApp {max-width: 1200px; margin: 0 auto;}
-    .css-18e3th9 {padding-top: 2rem;}
-    .stButton>button {background-color: #3498db; color: white;}
-    .stButton>button:hover {background-color: #2980b9;}
-    .chat-message {border-radius: 10px; padding: 10px; margin-bottom: 10px;}
-    .chat-message-user {background-color: rgba(52, 152, 219, 0.2); color: inherit;}
-    .chat-message-assistant {background-color: rgba(240, 240, 240, 0.2); color: inherit;}
-    .source-box {background-color: rgba(255, 255, 255, 0.1); color: inherit; border-radius: 5px; padding: 15px; margin-bottom: 10px; border-left: 5px solid #3498db;}
-    .metrics-box {background-color: rgba(255, 255, 255, 0.1); color: inherit; border-radius: 5px; padding: 15px; margin-top: 20px;}
-    .features-container {display: flex; flex-wrap: wrap; gap: 20px; justify-content: center; margin-top: 30px;}
-    .feature-item {flex: 1 1 calc(50% - 20px); min-width: 300px; display: flex; align-items: center; padding: 20px; border-radius: 10px; background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); transition: transform 0.3s, box-shadow 0.3s; border: 1px solid rgba(255, 255, 255, 0.1);}
-    .feature-item:hover {transform: translateY(-5px); box-shadow: 0 10px 20px rgba(0, 0, 0, 0.1);}
-    .feature-icon {width: 60px; height: 60px; border-radius: 50%; background: linear-gradient(135deg, #3498db, #2980b9); display: flex; align-items: center; justify-content: center; margin-right: 20px; box-shadow: 0 5px 15px rgba(52, 152, 219, 0.3);}
-    .feature-icon i {font-size: 24px; color: white;}
-    .feature-content {flex: 1;}
-    .feature-content h3 {margin-top: 0; margin-bottom: 10px; color: inherit;}
-    .feature-content p {margin: 0; font-size: 0.9em; color: inherit; opacity: 0.8;}
-    .input-container {margin-bottom: 20px; padding: 15px; border-radius: 10px; background-color: rgba(255, 255, 255, 0.05); border: 1px solid rgba(255, 255, 255, 0.1);}
-    </style>
-    """, unsafe_allow_html=True)
-
-    # App states
-    if 'chat_history' not in st.session_state:
-        st.session_state.chat_history = []
-    if 'page' not in st.session_state:
-        st.session_state.page = 'cover'
-
-    # Sidebar
-    with st.sidebar:
-        st.image("https://img.icons8.com/color/96/000000/caduceus.png", width=80)
-        st.title("DiReCT")
-        st.markdown("### Diagnostic Reasoning for Clinical Text")
-        st.markdown("---")
-        
-        if st.button("Home", key="home_btn"):
-            st.session_state.page = 'cover'
-        if st.button("Diagnostic Assistant", key="assistant_btn"):
-            st.session_state.page = 'chat'
-        if st.button("About", key="about_btn"):
-            st.session_state.page = 'about'
-            
-        st.markdown("---")
-        st.markdown("### Model Information")
-        st.markdown("**Embedding Model:** Bio_ClinicalBERT")
-        st.markdown("**LLM:** Llama-3.3-70B")
-        st.markdown("**Vector Store:** FAISS")
-
-    # Cover page
-    if st.session_state.page == 'cover':
-        # Hero section with animation
-        col1, col2 = st.columns([2, 1])
-        with col1:
-            st.markdown("<h1 style='font-size:3.5em;'>DiReCT</h1>", unsafe_allow_html=True)
-            st.markdown("<h2 style='font-size:1.8em;color:#3498db;'>Diagnostic Reasoning for Clinical Text</h2>", unsafe_allow_html=True)
-            st.markdown("""<p style='font-size:1.2em;'>A powerful RAG-based clinical diagnostic assistant that leverages the MIMIC-IV-Ext dataset to provide accurate medical insights and diagnostic reasoning.</p>""", unsafe_allow_html=True)
-            
-            st.markdown("""<br>""", unsafe_allow_html=True)
-            if st.button("Get Started", key="get_started"):
-                st.session_state.page = 'chat'
-                st.rerun()
-        
-        with col2:
-            # Animated medical icon
-            st.markdown("""
-            <div style='display:flex;justify-content:center;align-items:center;height:100%;'>
-                <img src="https://img.icons8.com/color/240/000000/healthcare-and-medical.png" style='max-width:90%;'>
-            </div>
-            """, unsafe_allow_html=True)
-        
-        # Modern Features section with Streamlit native components
-        st.markdown("<br><br>", unsafe_allow_html=True)
-        st.markdown("<h2 style='text-align:center;'>Key Features</h2>", unsafe_allow_html=True)
-        
-        # Create a 2x2 grid for features using Streamlit columns
-        col1, col2 = st.columns(2)
-        
-        # Feature 1
-        with col1:
-            st.markdown("""
-            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
-                        padding: 20px; border-radius: 10px; height: 100%; 
-                        border: 1px solid rgba(255, 255, 255, 0.1); margin-bottom: 20px;">
-                <h3>🔍 Intelligent Retrieval</h3>
-                <p>Finds the most relevant clinical information from the MIMIC-IV-Ext dataset</p>
-            </div>
-            """, unsafe_allow_html=True)
-        
-        # Feature 2
-        with col2:
-            st.markdown("""
-            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
-                        padding: 20px; border-radius: 10px; height: 100%; 
-                        border: 1px solid rgba(255, 255, 255, 0.1); margin-bottom: 20px;">
-                <h3>🧠 Advanced Reasoning</h3>
-                <p>Applies clinical knowledge to generate accurate diagnostic insights</p>
-            </div>
-            """, unsafe_allow_html=True)
-        
-        # Feature 3
-        with col1:
-            st.markdown("""
-            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
-                        padding: 20px; border-radius: 10px; height: 100%; 
-                        border: 1px solid rgba(255, 255, 255, 0.1);">
-                <h3>📄 Source Transparency</h3>
-                <p>Provides references to all clinical sources used in generating responses</p>
-            </div>
-            """, unsafe_allow_html=True)
-        
-        # Feature 4
-        with col2:
-            st.markdown("""
-            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
-                        padding: 20px; border-radius: 10px; height: 100%; 
-                        border: 1px solid rgba(255, 255, 255, 0.1);">
-                <h3>🌓 Dark/Light Theme Compatible</h3>
-                <p>Optimized interface that works seamlessly in both dark and light themes</p>
-            </div>
-            """, unsafe_allow_html=True)
-
-    # Chat interface
-    elif st.session_state.page == 'chat':
-        # Initialize session state for input if not exists
-        if 'user_input' not in st.session_state:
-            st.session_state.user_input = ""
-        
-        # Header with clear button
-        col1, col2 = st.columns([3, 1])
-        with col1:
-            st.markdown("<h1>Clinical Diagnostic Assistant</h1>", unsafe_allow_html=True)
-        with col2:
-            # Add a clear button in the header
-            if st.button("🗑️ Clear Chat"):
-                st.session_state.chat_history = []
-                st.session_state.user_input = ""
-                st.rerun()
-        
-        st.markdown("Ask any clinical diagnostic question and get insights based on medical knowledge and patient cases.")
-        
-        # Fixed input area at the top
-        with st.container():
-            st.markdown("<div class='input-container'>", unsafe_allow_html=True)
-            user_input = st.text_area("Ask a clinical question:", st.session_state.user_input, height=100, key="question_input")
-            col1, col2 = st.columns([1, 5])
-            with col1:
-                submit_button = st.button("Submit")
-            st.markdown("</div>", unsafe_allow_html=True)
-        
-        # Create a container for chat history
-        chat_container = st.container()
-        
-        # Process query
-        if submit_button and user_input:
-            if st.session_state.vectorstore is None:
-                st.error("Knowledge base not loaded. Please refresh the page and try again.")
-            else:
-                with st.spinner("Analyzing clinical data..."):
-                    try:
-                        # Add a small delay for UX
-                        time.sleep(0.5)
-                        
-                        # Run RAG
-                        response = run_rag_chat(user_input, st.session_state.vectorstore)
-                        response["retriever"] = st.session_state.vectorstore.as_retriever()
-                        
-                        # Clear previous chat history and only keep the current response
-                        st.session_state.chat_history = [(user_input, response)]
-                        
-                        # Clear the input field
-                        st.session_state.user_input = ""
-                        
-                        # Rerun to update UI
-                        st.rerun()
-                    except Exception as e:
-                        st.error(f"Error processing query: {str(e)}")
-        
-        # Display chat history in the container
-        with chat_container:
-            for i, (query, response) in enumerate(st.session_state.chat_history):
-                st.markdown(f"<div class='chat-message chat-message-user'><b>🧑‍⚕️ You:</b> {query}</div>", unsafe_allow_html=True)
-                
-                st.markdown(f"<div class='chat-message chat-message-assistant'><b>🩺 DiReCT:</b> {response['answer']}</div>", unsafe_allow_html=True)
-                
-                with st.expander("View Sources"):
-                    for doc in response["context"]:
-                        st.markdown(f"<div class='source-box'>"
-                                  f"<b>Source:</b> {Path(doc.metadata['source']).stem}<br>"
-                                  f"<b>Type:</b> {doc.metadata['type']}<br>"
-                                  f"<b>Content:</b> {doc.page_content[:300]}...</div>", 
-                                  unsafe_allow_html=True)
-                
-                # Show evaluation metrics if available
-                try:
-                    eval_scores = evaluate_rag_response(response, embeddings)
-                    with st.expander("View Evaluation Metrics"):
-                        col1, col2 = st.columns(2)
-                        with col1:
-                            st.metric("Hit Rate (Top-3)", f"{eval_scores['hit_rate']:.2f}")
-                        with col2:
-                            st.metric("Faithfulness", f"{eval_scores['faithfulness']:.2f}")
-                except Exception as e:
-                    st.warning(f"Evaluation metrics unavailable: {str(e)}")
-    
-    # About page
-    elif st.session_state.page == 'about':
-        st.markdown("<h1>About DiReCT</h1>", unsafe_allow_html=True)
-        
-        st.markdown("""
-        ### Project Overview
-        
-        DiReCT (Diagnostic Reasoning for Clinical Text) is a Retrieval-Augmented Generation (RAG) system designed to assist medical professionals with diagnostic reasoning based on clinical notes and medical knowledge.
-        
-        ### Data Sources
-        
-        This application uses the MIMIC-IV-Ext dataset, which contains de-identified clinical notes and medical records. The system processes:
-        
-        - Diagnostic flowcharts
-        - Patient cases
-        - Clinical guidelines
-        
-        ### Technical Implementation
-        
-        - **Embedding Model**: Bio_ClinicalBERT for domain-specific text understanding
-        - **Vector Database**: FAISS for efficient similarity search
-        - **LLM**: Llama-3.3-70B for generating medically accurate responses
-        - **Framework**: Built with LangChain and Streamlit
-        
-        ### Evaluation Metrics
-        
-        The system evaluates responses using:
-        
-        - **Hit Rate**: Measures how many relevant documents were retrieved
-        - **Faithfulness**: Measures how well the response aligns with the retrieved context
-        
-        ### Ethical Considerations
-        
-        This system is designed as a clinical decision support tool and not as a replacement for professional medical judgment. All patient data used has been properly de-identified in compliance with healthcare privacy regulations.
-        """)
-        
-        st.markdown("<br>", unsafe_allow_html=True)
-        st.markdown("### Developers")
-        st.markdown("This project was developed as part of an academic assignment on RAG systems for clinical applications.")
-
-if __name__ == "__main__":
-    main()
+import os
+import json
+import glob
+from pathlib import Path
+import torch
+import streamlit as st
+from dotenv import load_dotenv
+from langchain_groq import ChatGroq
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_core.documents import Document
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.chains import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+import numpy as np  
+from sentence_transformers import util
+import time
+
+# Set device for model (CUDA if available)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load environment variables - works for both local and Hugging Face Spaces
+load_dotenv()
+
+# Set up the clinical assistant LLM
+# Try to get API key from Hugging Face Spaces secrets first, then fall back to .env file
+try:
+    # For Hugging Face Spaces
+    from huggingface_hub.inference_api import InferenceApi
+    import os
+    groq_api_key = os.environ.get('GROQ_API_KEY')
+    
+    # If not found in environment, try to get from st.secrets (Streamlit Cloud/Spaces)
+    if not groq_api_key and hasattr(st, 'secrets') and 'GROQ_API_KEY' in st.secrets:
+        groq_api_key = st.secrets['GROQ_API_KEY']
+        
+    if not groq_api_key:
+        st.warning("API Key is not set in the secrets. Using a placeholder for UI demonstration.")
+        # For UI demonstration without API key
+        class MockLLM:
+            def invoke(self, prompt):
+                return {"answer": "This is a placeholder response. Please set up your GROQ_API_KEY to get real responses."}
+        llm = MockLLM()
+    else:    
+        llm = ChatGroq(groq_api_key=groq_api_key, model_name="llama-3.3-70b-versatile")
+        
+except Exception as e:
+    st.error(f"Error setting up LLM: {str(e)}")
+    class MockLLM:
+        def invoke(self, prompt):
+            return {"answer": f"Error setting up LLM: {str(e)}. Please check your API key configuration."}
+    llm = MockLLM()
+
+# Set up embeddings for clinical context (Bio_ClinicalBERT)
+embeddings = HuggingFaceEmbeddings(
+    model_name="emilyalsentzer/Bio_ClinicalBERT",
+    model_kwargs={"device": device}
+)
+
+
+def load_clinical_data():
+    """Load both flowcharts and patient cases"""
+    docs = []
+    
+    # Get the absolute path to the current script
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    
+    # Try to handle potential errors with file loading
+    try:
+        # Load diagnosis flowcharts
+        flowchart_dir = os.path.join(current_dir, "Diagnosis_flowchart")
+        if os.path.exists(flowchart_dir):
+            for fpath in glob.glob(os.path.join(flowchart_dir, "*.json")):
+                try:
+                    with open(fpath, 'r', encoding='utf-8') as f:
+                        data = json.load(f)
+                        content = f"""
+                        DIAGNOSTIC FLOWCHART: {Path(fpath).stem}
+                        Diagnostic Path: {data.get('diagnostic', 'N/A')}
+                        Key Criteria: {data.get('knowledge', 'N/A')}
+                        """
+                        docs.append(Document(
+                            page_content=content,
+                            metadata={"source": fpath, "type": "flowchart"}
+                        ))
+                except Exception as e:
+                    st.warning(f"Error loading flowchart file {fpath}: {str(e)}")
+        else:
+            st.warning(f"Flowchart directory not found at {flowchart_dir}")
+
+        # Load patient cases
+        finished_dir = os.path.join(current_dir, "Finished")
+        if os.path.exists(finished_dir):
+            for category_dir in glob.glob(os.path.join(finished_dir, "*")):
+                if os.path.isdir(category_dir):
+                    for case_file in glob.glob(os.path.join(category_dir, "*.json")):
+                        try:
+                            with open(case_file, 'r', encoding='utf-8') as f:
+                                case_data = json.load(f)
+                                notes = "\n".join(
+                                    f"{k}: {v}" for k, v in case_data.items() if k.startswith("input")
+                                )
+                                docs.append(Document(
+                                    page_content=f"""
+                                    PATIENT CASE: {Path(case_file).stem}
+                                    Category: {Path(category_dir).name}
+                                    Notes: {notes}
+                                    """,
+                                    metadata={"source": case_file, "type": "patient_case"}
+                                ))
+                        except Exception as e:
+                            st.warning(f"Error loading case file {case_file}: {str(e)}")
+        else:
+            st.warning(f"Finished directory not found at {finished_dir}")
+            
+        # If no documents were loaded, add a sample document for testing
+        if not docs:
+            st.warning("No clinical data files found. Using sample data for demonstration.")
+            docs.append(Document(
+                page_content="""SAMPLE CLINICAL DATA: This is sample data for demonstration purposes.
+                This application requires clinical data files to be present in the correct directories.
+                Please ensure the Diagnosis_flowchart and Finished directories exist with proper JSON files.""",
+                metadata={"source": "sample", "type": "sample"}
+            ))
+    except Exception as e:
+        st.error(f"Error loading clinical data: {str(e)}")
+        # Add a fallback document
+        docs.append(Document(
+            page_content="Error loading clinical data. This is a fallback document for demonstration purposes.",
+            metadata={"source": "error", "type": "error"}
+        ))
+    return docs
+
+def build_vectorstore():
+    """Build and return the vectorstore using FAISS"""
+    documents = load_clinical_data()
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    splits = splitter.split_documents(documents)
+    vectorstore = FAISS.from_documents(splits, embeddings)
+    return vectorstore
+
+# Path for saving/loading the vectorstore
+def get_vectorstore_path():
+    """Get the path for saving/loading the vectorstore"""
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    return os.path.join(current_dir, "vectorstore")
+
+# Initialize vectorstore with disk persistence
+@st.cache_resource(show_spinner="Loading clinical knowledge base...")
+def get_vectorstore():
+    """Get or create the vectorstore with disk persistence"""
+    vectorstore_path = get_vectorstore_path()
+    
+    # Try to load from disk first
+    try:
+        if os.path.exists(vectorstore_path):
+            st.info("Loading vectorstore from disk...")
+            # Set allow_dangerous_deserialization to True since we trust our own vectorstore files
+            return FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True)
+    except Exception as e:
+        st.warning(f"Could not load vectorstore from disk: {str(e)}. Building new vectorstore.")
+    
+    # If loading fails or doesn't exist, build a new one
+    st.info("Building new vectorstore...")
+    vectorstore = build_vectorstore()
+    
+    # Save to disk for future use
+    try:
+        os.makedirs(vectorstore_path, exist_ok=True)
+        vectorstore.save_local(vectorstore_path)
+        st.success("Vectorstore saved to disk for future use")
+    except Exception as e:
+        st.warning(f"Could not save vectorstore to disk: {str(e)}")
+    
+    return vectorstore
+
+def run_rag_chat(query, vectorstore):
+    """Run the Retrieval-Augmented Generation (RAG) for clinical questions"""
+    try:
+        retriever = vectorstore.as_retriever()
+
+        prompt_template = ChatPromptTemplate.from_template("""
+        You are a clinical assistant AI. Based on the following clinical context, provide a reasoned and medically sound answer to the question.
+
+        <context>
+        {context}
+        </context>
+
+        Question: {input}
+
+        Answer:
+        """)
+
+        retrieved_docs = retriever.invoke(query, k=3)
+        retrieved_context = "\n".join([doc.page_content for doc in retrieved_docs])
+
+        # Create document chain first
+        document_chain = create_stuff_documents_chain(llm, prompt_template)
+        
+        # Then create retrieval chain
+        chain = create_retrieval_chain(retriever, document_chain)
+
+        # Invoke the chain
+        response = chain.invoke({"input": query})
+        
+        # Add retrieved documents to response for transparency
+        response["context"] = retrieved_docs
+        
+        return response
+    except Exception as e:
+        st.error(f"Error in RAG processing: {str(e)}")
+        # Return a fallback response
+        return {
+            "answer": f"I encountered an error processing your query: {str(e)}",
+            "context": [],
+            "input": query
+        }
+
+def calculate_hit_rate(retriever, query, expected_docs, k=3):
+    """Calculate the hit rate for top-k retrieved documents"""
+    retrieved_docs = retriever.get_relevant_documents(query, k=k)
+    retrieved_contents = [doc.page_content for doc in retrieved_docs]
+    
+    hits = 0
+    for expected in expected_docs:
+        if any(expected in retrieved for retrieved in retrieved_contents):
+            hits += 1
+    
+    return hits / len(expected_docs) if expected_docs else 0.0
+
+def evaluate_rag_response(response, embeddings):
+    """Evaluate the RAG response for faithfulness and hit rate"""
+    scores = {}
+
+    # Faithfulness: Answer-Context Similarity
+    answer_embed = embeddings.embed_query(response["answer"])
+    context_embeds = [embeddings.embed_query(doc.page_content) for doc in response["context"]]
+    similarities = [util.cos_sim(answer_embed, ctx_embed).item() for ctx_embed in context_embeds]
+    scores["faithfulness"] = float(np.mean(similarities)) if similarities else 0.0
+
+    # Custom Hit Rate Calculation
+    retriever = response["retriever"]
+    scores["hit_rate"] = calculate_hit_rate(
+        retriever,
+        query=response["input"],
+        expected_docs=[doc.page_content for doc in response["context"]],
+        k=3
+    )
+    
+    return scores
+
+def main():
+    """Main function to run the Streamlit app"""
+    # Set page configuration
+    st.set_page_config(
+        page_title="DiReCT - Clinical Diagnostic Assistant",
+        page_icon="🩺",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+    
+    # Load vectorstore only once using session state
+    if 'vectorstore' not in st.session_state:
+        with st.spinner("Loading clinical knowledge base... This may take a minute."):
+            try:
+                st.session_state.vectorstore = get_vectorstore()
+                # Use custom styled message without the success icon
+                st.markdown("<div style='padding:10px 15px;background-color:rgba(40,167,69,0.2);border-radius:5px;border-left:5px solid rgba(40,167,69,0.8);'>Clinical knowledge base loaded successfully!</div>", unsafe_allow_html=True)
+            except Exception as e:
+                st.error(f"Error loading knowledge base: {str(e)}")
+                st.session_state.vectorstore = None
+
+    # Custom CSS for modern look with dark theme compatibility
+    st.markdown("""
+    <style>
+    .stApp {max-width: 1200px; margin: 0 auto;}
+    .css-18e3th9 {padding-top: 2rem;}
+    .stButton>button {background-color: #3498db; color: white;}
+    .stButton>button:hover {background-color: #2980b9;}
+    .chat-message {border-radius: 10px; padding: 10px; margin-bottom: 10px;}
+    .chat-message-user {background-color: rgba(52, 152, 219, 0.2); color: inherit;}
+    .chat-message-assistant {background-color: rgba(240, 240, 240, 0.2); color: inherit;}
+    .source-box {background-color: rgba(255, 255, 255, 0.1); color: inherit; border-radius: 5px; padding: 15px; margin-bottom: 10px; border-left: 5px solid #3498db;}
+    .metrics-box {background-color: rgba(255, 255, 255, 0.1); color: inherit; border-radius: 5px; padding: 15px; margin-top: 20px;}
+    .features-container {display: flex; flex-wrap: wrap; gap: 20px; justify-content: center; margin-top: 30px;}
+    .feature-item {flex: 1 1 calc(50% - 20px); min-width: 300px; display: flex; align-items: center; padding: 20px; border-radius: 10px; background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); transition: transform 0.3s, box-shadow 0.3s; border: 1px solid rgba(255, 255, 255, 0.1);}
+    .feature-item:hover {transform: translateY(-5px); box-shadow: 0 10px 20px rgba(0, 0, 0, 0.1);}
+    .feature-icon {width: 60px; height: 60px; border-radius: 50%; background: linear-gradient(135deg, #3498db, #2980b9); display: flex; align-items: center; justify-content: center; margin-right: 20px; box-shadow: 0 5px 15px rgba(52, 152, 219, 0.3);}
+    .feature-icon i {font-size: 24px; color: white;}
+    .feature-content {flex: 1;}
+    .feature-content h3 {margin-top: 0; margin-bottom: 10px; color: inherit;}
+    .feature-content p {margin: 0; font-size: 0.9em; color: inherit; opacity: 0.8;}
+    .input-container {margin-bottom: 20px; padding: 15px; border-radius: 10px; background-color: rgba(255, 255, 255, 0.05); border: 1px solid rgba(255, 255, 255, 0.1);}
+    </style>
+    """, unsafe_allow_html=True)
+
+    # App states
+    if 'chat_history' not in st.session_state:
+        st.session_state.chat_history = []
+    if 'page' not in st.session_state:
+        st.session_state.page = 'cover'
+
+    # Sidebar
+    with st.sidebar:
+        st.image("https://img.icons8.com/color/96/000000/caduceus.png", width=80)
+        st.title("DiReCT")
+        st.markdown("### Diagnostic Reasoning for Clinical Text")
+        st.markdown("---")
+        
+        if st.button("Home", key="home_btn"):
+            st.session_state.page = 'cover'
+        if st.button("Diagnostic Assistant", key="assistant_btn"):
+            st.session_state.page = 'chat'
+        if st.button("About", key="about_btn"):
+            st.session_state.page = 'about'
+            
+        st.markdown("---")
+        st.markdown("### Model Information")
+        st.markdown("**Embedding Model:** Bio_ClinicalBERT")
+        st.markdown("**LLM:** Llama-3.3-70B")
+        st.markdown("**Vector Store:** FAISS")
+
+    # Cover page
+    if st.session_state.page == 'cover':
+        # Hero section with animation
+        col1, col2 = st.columns([2, 1])
+        with col1:
+            st.markdown("<h1 style='font-size:3.5em;'>DiReCT</h1>", unsafe_allow_html=True)
+            st.markdown("<h2 style='font-size:1.8em;color:#3498db;'>Diagnostic Reasoning for Clinical Text</h2>", unsafe_allow_html=True)
+            st.markdown("""<p style='font-size:1.2em;'>A powerful RAG-based clinical diagnostic assistant that leverages the MIMIC-IV-Ext dataset to provide accurate medical insights and diagnostic reasoning.</p>""", unsafe_allow_html=True)
+            
+            st.markdown("""<br>""", unsafe_allow_html=True)
+            if st.button("Get Started", key="get_started"):
+                st.session_state.page = 'chat'
+                st.rerun()
+        
+        with col2:
+            # Animated medical icon
+            st.markdown("""
+            <div style='display:flex;justify-content:center;align-items:center;height:100%;'>
+                <img src="https://img.icons8.com/color/240/000000/healthcare-and-medical.png" style='max-width:90%;'>
+            </div>
+            """, unsafe_allow_html=True)
+        
+        # Modern Features section with Streamlit native components
+        st.markdown("<br><br>", unsafe_allow_html=True)
+        st.markdown("<h2 style='text-align:center;'>Key Features</h2>", unsafe_allow_html=True)
+        
+        # Create a 2x2 grid for features using Streamlit columns
+        col1, col2 = st.columns(2)
+        
+        # Feature 1
+        with col1:
+            st.markdown("""
+            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
+                        padding: 20px; border-radius: 10px; height: 100%; 
+                        border: 1px solid rgba(255, 255, 255, 0.1); margin-bottom: 20px;">
+                <h3>🔍 Intelligent Retrieval</h3>
+                <p>Finds the most relevant clinical information from the MIMIC-IV-Ext dataset</p>
+            </div>
+            """, unsafe_allow_html=True)
+        
+        # Feature 2
+        with col2:
+            st.markdown("""
+            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
+                        padding: 20px; border-radius: 10px; height: 100%; 
+                        border: 1px solid rgba(255, 255, 255, 0.1); margin-bottom: 20px;">
+                <h3>🧠 Advanced Reasoning</h3>
+                <p>Applies clinical knowledge to generate accurate diagnostic insights</p>
+            </div>
+            """, unsafe_allow_html=True)
+        
+        # Feature 3
+        with col1:
+            st.markdown("""
+            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
+                        padding: 20px; border-radius: 10px; height: 100%; 
+                        border: 1px solid rgba(255, 255, 255, 0.1);">
+                <h3>📄 Source Transparency</h3>
+                <p>Provides references to all clinical sources used in generating responses</p>
+            </div>
+            """, unsafe_allow_html=True)
+        
+        # Feature 4
+        with col2:
+            st.markdown("""
+            <div style="background: linear-gradient(135deg, rgba(72, 126, 176, 0.1), rgba(72, 126, 176, 0.2)); 
+                        padding: 20px; border-radius: 10px; height: 100%; 
+                        border: 1px solid rgba(255, 255, 255, 0.1);">
+                <h3>🌓 Dark/Light Theme Compatible</h3>
+                <p>Optimized interface that works seamlessly in both dark and light themes</p>
+            </div>
+            """, unsafe_allow_html=True)
+
+    # Chat interface
+    elif st.session_state.page == 'chat':
+        # Initialize session state for input if not exists
+        if 'user_input' not in st.session_state:
+            st.session_state.user_input = ""
+        
+        # Header with clear button
+        col1, col2 = st.columns([3, 1])
+        with col1:
+            st.markdown("<h1>Clinical Diagnostic Assistant</h1>", unsafe_allow_html=True)
+        with col2:
+            # Add a clear button in the header
+            if st.button("🗑️ Clear Chat"):
+                st.session_state.chat_history = []
+                st.session_state.user_input = ""
+                st.rerun()
+        
+        st.markdown("Ask any clinical diagnostic question and get insights based on medical knowledge and patient cases.")
+        
+        # Fixed input area at the top
+        with st.container():
+            st.markdown("<div class='input-container'>", unsafe_allow_html=True)
+            user_input = st.text_area("Ask a clinical question:", st.session_state.user_input, height=100, key="question_input")
+            col1, col2 = st.columns([1, 5])
+            with col1:
+                submit_button = st.button("Submit")
+            st.markdown("</div>", unsafe_allow_html=True)
+        
+        # Create a container for chat history
+        chat_container = st.container()
+        
+        # Process query
+        if submit_button and user_input:
+            if st.session_state.vectorstore is None:
+                st.error("Knowledge base not loaded. Please refresh the page and try again.")
+            else:
+                with st.spinner("Analyzing clinical data..."):
+                    try:
+                        # Add a small delay for UX
+                        time.sleep(0.5)
+                        
+                        # Run RAG
+                        response = run_rag_chat(user_input, st.session_state.vectorstore)
+                        response["retriever"] = st.session_state.vectorstore.as_retriever()
+                        
+                        # Clear previous chat history and only keep the current response
+                        st.session_state.chat_history = [(user_input, response)]
+                        
+                        # Clear the input field
+                        st.session_state.user_input = ""
+                        
+                        # Rerun to update UI
+                        st.rerun()
+                    except Exception as e:
+                        st.error(f"Error processing query: {str(e)}")
+        
+        # Display chat history in the container
+        with chat_container:
+            for i, (query, response) in enumerate(st.session_state.chat_history):
+                st.markdown(f"<div class='chat-message chat-message-user'><b>🧑‍⚕️ You:</b> {query}</div>", unsafe_allow_html=True)
+                
+                st.markdown(f"<div class='chat-message chat-message-assistant'><b>🩺 DiReCT:</b> {response['answer']}</div>", unsafe_allow_html=True)
+                
+                with st.expander("View Sources"):
+                    for doc in response["context"]:
+                        st.markdown(f"<div class='source-box'>"
+                                  f"<b>Source:</b> {Path(doc.metadata['source']).stem}<br>"
+                                  f"<b>Type:</b> {doc.metadata['type']}<br>"
+                                  f"<b>Content:</b> {doc.page_content[:300]}...</div>", 
+                                  unsafe_allow_html=True)
+                
+                # Show evaluation metrics if available
+                try:
+                    eval_scores = evaluate_rag_response(response, embeddings)
+                    with st.expander("View Evaluation Metrics"):
+                        col1, col2 = st.columns(2)
+                        with col1:
+                            st.metric("Hit Rate (Top-3)", f"{eval_scores['hit_rate']:.2f}")
+                        with col2:
+                            st.metric("Faithfulness", f"{eval_scores['faithfulness']:.2f}")
+                except Exception as e:
+                    st.warning(f"Evaluation metrics unavailable: {str(e)}")
+    
+    # About page
+    elif st.session_state.page == 'about':
+        st.markdown("<h1>About DiReCT</h1>", unsafe_allow_html=True)
+        
+        st.markdown("""
+        ### Project Overview
+        
+        DiReCT (Diagnostic Reasoning for Clinical Text) is a Retrieval-Augmented Generation (RAG) system designed to assist medical professionals with diagnostic reasoning based on clinical notes and medical knowledge.
+        
+        ### Data Sources
+        
+        This application uses the MIMIC-IV-Ext dataset, which contains de-identified clinical notes and medical records. The system processes:
+        
+        - Diagnostic flowcharts
+        - Patient cases
+        - Clinical guidelines
+        
+        ### Technical Implementation
+        
+        - **Embedding Model**: Bio_ClinicalBERT for domain-specific text understanding
+        - **Vector Database**: FAISS for efficient similarity search
+        - **LLM**: Llama-3.3-70B for generating medically accurate responses
+        - **Framework**: Built with LangChain and Streamlit
+        
+        ### Evaluation Metrics
+        
+        The system evaluates responses using:
+        
+        - **Hit Rate**: Measures how many relevant documents were retrieved
+        - **Faithfulness**: Measures how well the response aligns with the retrieved context
+        
+        ### Ethical Considerations
+        
+        This system is designed as a clinical decision support tool and not as a replacement for professional medical judgment. All patient data used has been properly de-identified in compliance with healthcare privacy regulations.
+        """)
+        
+        st.markdown("<br>", unsafe_allow_html=True)
+        st.markdown("### Developers")
+        st.markdown("This project was developed as part of an academic assignment on RAG systems for clinical applications.")
+
+if __name__ == "__main__":
+    main()