Update app.py

app.py

@@ -1,5 +1,4 @@
 # Step 1: Import Libraries
-# No changes needed here, these will be installed from your requirements.txt
 import gradio as gr
 from sentence_transformers import SentenceTransformer
 import faiss
@@ -9,9 +8,7 @@ import pickle
 import os
 
 # Step 2: Configure API Key from Hugging Face Secrets
-# This section is modified to securely access the API key from your Space's secrets.
 try:
-    # Make sure to set 'GOOGLE_API_KEY' in your Hugging Face Space's settings/secrets
     GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
     if GOOGLE_API_KEY is None:
         raise ValueError("GOOGLE_API_KEY not found in environment variables.")
@@ -21,18 +18,15 @@ except Exception as e:
     print(f"ERROR: Could not configure API key. Please ensure 'GOOGLE_API_KEY' is set in your Hugging Face Space secrets. Details: {e}")
 
 # Step 3: Define Data Path
-# This is updated to point to the 'rag_chatbot_data' folder you uploaded to the Space.
-# It no longer uses Google Drive.
 DATA_PATH = "data"
 vector_store_file = os.path.join(DATA_PATH, "vector_store.index")
 data_file = os.path.join(DATA_PATH, "data.pkl")
 
 # Step 4: Load Models and Pre-processed Data
-# This section is slightly restructured to load models only if the data is found.
 vector_store_data = None
 if os.path.exists(vector_store_file) and os.path.exists(data_file):
     try:
-        print("Loading pre-processed data from the repository...")
+        print(f"Loading pre-processed data from the '{DATA_PATH}' directory...")
         index = faiss.read_index(vector_store_file)
         with open(data_file, "rb") as f:
             stored_data = pickle.load(f)
@@ -40,30 +34,26 @@ if os.path.exists(vector_store_file) and os.path.exists(data_file):
             sources = stored_data["sources"]
         print("✅ Data loaded successfully.")
 
-        # Store everything in a state object for Gradio
         vector_store_data = (index, texts, sources)
 
-        # Load the embedding and LLM models
         print("Loading AI and embedding models...")
         embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5')
-        llm = genai.GenerativeModel('gemini-2.5-flash') # Using gemini-pro as it's a robust choice
+        llm = genai.GenerativeModel('gemini-1.5-flash-latest')
         print("✅ Models loaded successfully.")
 
     except Exception as e:
         print(f"❌ ERROR: An error occurred during data or model loading: {e}")
         vector_store_data = None
-
 else:
-    print("❌ ERROR: Pre-processed data not found in the repository.")
-    print(f"Please make sure '{vector_store_file}' and '{data_file}' exist in the '{DATA_PATH}' folder.")
+    print(f"❌ ERROR: Pre-processed data not found in the '{DATA_PATH}' directory.")
+    print(f"Please make sure '{vector_store_file}' and '{data_file}' exist.")
 
 
-# Step 5: RAG and Chat Functions (No changes needed here)
+# Step 5: RAG and Chat Functions (with Updated System Prompt)
 def get_relevant_context(query, index, top_k=5):
     query_embedding = embedding_model.encode([query])
     distances, indices = index.search(query_embedding, top_k)
     context = []
-    # Ensure that the indices are within the bounds of the texts list
     for i in indices[0]:
         if i < len(texts):
             context.append({"text": texts[i], "source": sources[i]})
@@ -74,14 +64,17 @@ def chat_with_rag(message, history, vector_store_data):
     relevant_context = get_relevant_context(message, index)
     context_str = "\n\n".join([f"Source: {c['source']}\nContent: {c['text']}" for c in relevant_context])
 
+    # --- NEW, ROBUST SYSTEM PROMPT ---
     prompt = f"""
-    You are an Expert scientist in the Halassa Lab at MIT, an expert in computational neuroscience.
-    Your primary goal is to answer questions as thoroughly and accurately as possible. You may intelligently synthesize information from the provided context, which consists of key papers from the lab.
+    You are a friendly and engaging science communicator for the Halassa Lab at MIT.
+    Your goal is to explain the lab's complex computational neuroscience research to a general audience that has little to no scientific background.
 
-    Follow these rules strictly if you are using the provided context:
-    1.  Do not simply copy-paste from the context. Read the relevant passages and formulate a comprehensive, well-written answer in your own words.
-    2.  When your answer uses information directly from a provided paper, you MUST cite the source at the end of the output in a list. Use the format [filename - Page X].
-    3.  If the provided papers offer relevant concepts but do not contain the full answer, use your broader knowledge of computational neuroscience to provide a more complete explanation.
+    Follow these rules strictly:
+    1.  **Simplify, Don't Dumb Down:** Break down complex topics into simple, easy-to-understand concepts. Use analogies and real-world examples (e.g., "think of the thalamus as a busy switchboard operator for the brain").
+    2.  **Be Engaging and Accessible:** Write in a clear, conversational, and friendly tone. Avoid jargon at all costs. If you must use a technical term, explain it immediately in simple terms.
+    3.  **Synthesize and Explain:** Do not just copy-paste from the provided research papers. Read the relevant context and formulate a comprehensive, well-written answer in your own words.
+    4.  **Cite Your Sources:** At the end of your response, if you used information from the provided papers, you MUST include a "Sources:" list. Use the format [filename - Page X]. This adds credibility.
+    5.  **Stay Focused:** Only answer questions related to the Halassa Lab's work based on the provided context. If the context doesn't contain the answer, state that the information isn't available in the provided documents.
 
     Context from the Halassa Lab's papers:
     ---
@@ -90,7 +83,7 @@ def chat_with_rag(message, history, vector_store_data):
 
     User Question: {message}
 
-    Expert Answer:
+    Your Friendly Explanation:
     """
     try:
         response = llm.generate_content(prompt)
@@ -98,27 +91,70 @@ def chat_with_rag(message, history, vector_store_data):
     except Exception as e:
         return f"An error occurred with the AI model: {str(e)}"
 
-# Step 6: Gradio User Interface
-with gr.Blocks(theme=gr.themes.Soft(), title="Halassa Literature Chatbot") as demo:
-    gr.Markdown("# Halassa Lab Onboarder")
-    if vector_store_data is None:
-        gr.Markdown("## ⚠️ Error: Could not load data or models. Please check the logs in the Hugging Face Space for details.")
-    else:
-        gr.Markdown("The documents have been pre-loaded. Ask your questions below.")
-        chatbot_ui = gr.Chatbot(label="Chat History", height=600, layout="panel")
-        message_box = gr.Textbox(label="Ask your question...", lines=3)
-        clear_button = gr.ClearButton(components=[chatbot_ui, message_box])
-
-        def respond(message, history):
-            # Pass the loaded vector_store_data to the chat function
-            response_text = chat_with_rag(message, history, vector_store_data)
-            history.append((message, response_text))
-            return "", history
-
-        message_box.submit(respond, inputs=[message_box, chatbot_ui], outputs=[message_box, chatbot_ui])
+# Step 6: Polished Gradio User Interface
+# Define a professional, clean theme for a science lab
+theme = gr.themes.Monochrome(
+    primary_hue="indigo",
+    secondary_hue="blue",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Space Mono"), "ui-sans-serif", "system-ui", "sans-serif"],
+).set(
+    body_background_fill="*neutral_50",
+    block_background_fill="*neutral_100",
+    block_border_width="1px",
+    block_shadow="*shadow_md",
+    button_shadow="*shadow_md",
+    button_primary_background_fill="*primary_500",
+    button_primary_background_fill_hover="*primary_600",
+    button_primary_text_color="white",
+)
+
+# List of example questions for users to click on
+example_questions = [
+    "What is the main focus of the Halassa Lab's research, in simple terms?",
+    "How does the thalamus affect our attention? Can you use an analogy?",
+    "What have you learned about sleep and how it affects the brain?",
+    "Explain the role of the thalamic reticular nucleus like I'm a high school student.",
+]
+
+with gr.Blocks(theme=theme, title="Halassa Lab AI Explainer") as demo:
+    with gr.Column():
+        # --- Header Section ---
+        gr.Image("https://d39w22sdwnt1s2.cloudfront.net/wp-content/uploads/2023/10/Halassa-M-2023-Option-2-1200x675.jpg", show_label=False, show_download_button=False, container=False)
+        gr.Markdown("# Welcome to the Halassa Lab AI Explainer")
+        gr.Markdown("I'm an AI assistant trained on the Halassa Lab's research. My job is to explain our work on how the brain processes information, pays attention, and sleeps. Ask me a question in plain English, or try one of the examples below!")
+
+        # --- Chatbot Interface ---
+        chatbot_ui = gr.Chatbot(label="Conversation", height=500, layout="panel", bubble_full_width=False)
+        message_box = gr.Textbox(label="Ask your question here...", lines=3, placeholder="e.g., How does the brain filter out distractions?")
+        
+        with gr.Row():
+            submit_button = gr.Button("Submit", variant="primary", scale=2)
+            clear_button = gr.ClearButton(components=[chatbot_ui, message_box], value="Clear Conversation", scale=1)
+
+        # --- Example Questions Section ---
+        gr.Examples(
+            examples=example_questions,
+            inputs=message_box,
+            label="Click an example to get started:"
+        )
+
+    # --- Logic to make the chatbot respond ---
+    def respond(message, history):
+        response_text = chat_with_rag(message, history, vector_store_data)
+        history.append((message, response_text))
+        return "", history
+
+    submit_button.click(respond, inputs=[message_box, chatbot_ui], outputs=[message_box, chatbot_ui])
+    message_box.submit(respond, inputs=[message_box, chatbot_ui], outputs=[message_box, chatbot_ui])
+
 
 # Step 7: Launch the app
-# On Hugging Face, demo.launch() is all you need.
-# It will run the app and make it accessible.
 if vector_store_data:
+    demo.launch()
+else:
+    # Display a clear error message in the UI if data loading failed
+    with gr.Blocks() as demo:
+        gr.Markdown("## ⚠️ Application Error")
+        gr.Markdown("Could not load the necessary data and models. Please check the Hugging Face Space logs for details. This may be due to missing data files or an incorrect API key.")
     demo.launch()