Update app.py

app.py

@@ -73,9 +73,12 @@ def chunk_text(text: str, max_chunk_size: int = 1000) :
 
 # Streamlit UI
 st.title("CUDA Documentation QA System")
+
 # Initialize global variables
-vector_store = None
-documents_loaded = False
+if 'vector_store' not in st.session_state:
+    st.session_state.vector_store = None
+if 'documents_loaded' not in st.session_state:
+    st.session_state.documents_loaded = False
 
 # Crawling and processing the data
 if st.button('Crawl CUDA Documentation'):
@@ -94,12 +97,13 @@ if st.button('Crawl CUDA Documentation'):
                                            model_kwargs={'device': 'cpu'})
 
         # Store embeddings in FAISS
-        vector_store = FAISS.from_texts(texts, embeddings)
+        st.session_state.vector_store = FAISS.from_texts(texts, embeddings)
+        st.session_state.documents_loaded = True
         st.write("Embeddings stored in FAISS.")
 
 # Asking questions
 query = st.text_input("Enter your question about CUDA:")
-if query:
+if query and st.session_state.documents_loaded:
     with st.spinner('Searching for an answer...'):
         # Initialize Google Generative AI
         llm = GoogleGenerativeAI(model='gemini-1.0-pro', google_api_key="AIzaSyC1AvHnvobbycU8XSCXh-gRq3DUfG0EP98")
@@ -112,7 +116,7 @@ if query:
         # Create the retrieval QA chain
         qa_chain = RetrievalQA.from_chain_type(
             chain_type="map_rerank",
-            retriever=vector_store.as_retriever(),
+            retriever=st.session_state.vector_store.as_retriever(),
             combine_documents_chain=qa_prompt,
             llm=llm
         )
@@ -122,3 +126,6 @@ if query:
         st.write(response['answer'])
         st.write("**Source:**")
         st.write(response['source'])
+elif query:
+    st.warning("Please crawl the CUDA documentation first.")
+