Update app.py

app.py

@@ -1,5 +1,6 @@
 from io import BytesIO
 import streamlit as st
+import shutil
 import requests
 import os
 from PyPDF2 import PdfReader
@@ -68,6 +69,14 @@ def get_conversation_chain(vectorstore):
         )
         return conversation_chain
 
+def retrain_model():
+    st.session_state.conversation = None
+    st.session_state.chat_history = None
+    pdf_doc = getpdfdoc() # get pdf
+    raw_text = extract_text_from_pdf(pdf_doc) # get pdf text
+    text_chunks = get_text_chunks(raw_text) # get the text chunks    
+    vectorstore = get_vectorstore(text_chunks) # create vector store
+    st.session_state.conversation = get_conversation_chain(vectorstore) # create conversation chain
 
 def handle_userinput(user_question):
     response = st.session_state.conversation({'question': user_question})
@@ -85,8 +94,6 @@ def main():
         st.session_state.conversation = None
     if "chat_history" not in st.session_state:
         st.session_state.chat_history = None
-    if "vectorstore_saved" not in st.session_state:
-        st.session_state.vectorstore_saved = False
     
     if st.session_state.conversation is None:
         if os.path.isdir("./chroma_db"):
@@ -94,30 +101,15 @@ def main():
                 with st.spinner("Loading vector store..."):
                     vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=OpenAIEmbeddings())
                     st.session_state.conversation = get_conversation_chain(vectorstore)
-                    st.session_state.vectorstore_saved = True
-
-    if st.session_state.conversation is None:
-        pdf_doc = getpdfdoc()
-
-        # get pdf text
-        raw_text = extract_text_from_pdf(pdf_doc)
-
-        # get the text chunks
-        text_chunks = get_text_chunks(raw_text)
-
-        # create vector store
-        vectorstore = get_vectorstore(text_chunks)
-
-        # create conversation chain
-        st.session_state.conversation = get_conversation_chain(vectorstore)
+            else:
+                retrain_model()
+        else:
+            retrain_model()
 
     if st.session_state.conversation is not None:
+        st.sidebar.button("Retrain model", on_click=retrain_model)
         st.header("Ask questions from 48 Laws of Power:books:")
-        if st.session_state.vectorstore_saved:
-            st.info("Chroma saved")
-        else:
-            st.info("Chroma failed to save")
-        user_question = st.chat_input("Ask a question related 48 Laws of Power:")
+        user_question = st.chat_input("Ask a question about your documents:")
         if user_question:
             handle_userinput(user_question)