update UI

app.py

@@ -1,6 +1,7 @@
 import os
 
 os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
+os.environ["WANDB_API_KEY"] = os.getenv("WANDB_API_KEY")
 
 import streamlit as st
 import weave
@@ -14,31 +15,25 @@ st.set_page_config(
     menu_items=None,
 )
 
-wandb_api_key = st.sidebar.text_input("WANDB_API_KEY", type="password")
-if len(wandb_api_key) >= 10:
-    os.environ["WANDB_API_KEY"] = wandb_api_key
-else:
-    st.stop()
-
 WANDB_PROJECT = "paper_reader"
 
 weave.init(f"{WANDB_PROJECT}")
 
 st.title("Chat with the Llama 3 paper 💬🦙")
 
+with st.spinner('Loading the RAG pipeline...'):
+    @st.cache_resource(show_spinner=False)
+    def load_rag_pipeline():
+        rag_pipeline = SimpleRAGPipeline()
+        rag_pipeline.build_query_engine()
 
-@st.cache_resource(show_spinner=False)
-def load_rag_pipeline():
-    rag_pipeline = SimpleRAGPipeline()
-    rag_pipeline.build_query_engine()
-
-    return rag_pipeline
+        return rag_pipeline
 
 
-if "rag_pipeline" not in st.session_state.keys():
-    st.session_state.rag_pipeline = load_rag_pipeline()
+    if "rag_pipeline" not in st.session_state.keys():
+        st.session_state.rag_pipeline = load_rag_pipeline()
 
-rag_pipeline = st.session_state["rag_pipeline"]
+    rag_pipeline = st.session_state["rag_pipeline"]
 
 
 def generate_response(query):