viboognesh commited on
Commit
a55873f
·
verified ·
1 Parent(s): dc5cd0f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -23
app.py CHANGED
@@ -1,5 +1,6 @@
1
  from io import BytesIO
2
  import streamlit as st
 
3
  import requests
4
  import os
5
  from PyPDF2 import PdfReader
@@ -68,6 +69,14 @@ def get_conversation_chain(vectorstore):
68
  )
69
  return conversation_chain
70
 
 
 
 
 
 
 
 
 
71
 
72
  def handle_userinput(user_question):
73
  response = st.session_state.conversation({'question': user_question})
@@ -85,8 +94,6 @@ def main():
85
  st.session_state.conversation = None
86
  if "chat_history" not in st.session_state:
87
  st.session_state.chat_history = None
88
- if "vectorstore_saved" not in st.session_state:
89
- st.session_state.vectorstore_saved = False
90
 
91
  if st.session_state.conversation is None:
92
  if os.path.isdir("./chroma_db"):
@@ -94,30 +101,15 @@ def main():
94
  with st.spinner("Loading vector store..."):
95
  vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=OpenAIEmbeddings())
96
  st.session_state.conversation = get_conversation_chain(vectorstore)
97
- st.session_state.vectorstore_saved = True
98
-
99
- if st.session_state.conversation is None:
100
- pdf_doc = getpdfdoc()
101
-
102
- # get pdf text
103
- raw_text = extract_text_from_pdf(pdf_doc)
104
-
105
- # get the text chunks
106
- text_chunks = get_text_chunks(raw_text)
107
-
108
- # create vector store
109
- vectorstore = get_vectorstore(text_chunks)
110
-
111
- # create conversation chain
112
- st.session_state.conversation = get_conversation_chain(vectorstore)
113
 
114
  if st.session_state.conversation is not None:
 
115
  st.header("Ask questions from 48 Laws of Power:books:")
116
- if st.session_state.vectorstore_saved:
117
- st.info("Chroma saved")
118
- else:
119
- st.info("Chroma failed to save")
120
- user_question = st.chat_input("Ask a question related 48 Laws of Power:")
121
  if user_question:
122
  handle_userinput(user_question)
123
 
 
1
  from io import BytesIO
2
  import streamlit as st
3
+ import shutil
4
  import requests
5
  import os
6
  from PyPDF2 import PdfReader
 
69
  )
70
  return conversation_chain
71
 
72
+ def retrain_model():
73
+ st.session_state.conversation = None
74
+ st.session_state.chat_history = None
75
+ pdf_doc = getpdfdoc() # get pdf
76
+ raw_text = extract_text_from_pdf(pdf_doc) # get pdf text
77
+ text_chunks = get_text_chunks(raw_text) # get the text chunks
78
+ vectorstore = get_vectorstore(text_chunks) # create vector store
79
+ st.session_state.conversation = get_conversation_chain(vectorstore) # create conversation chain
80
 
81
  def handle_userinput(user_question):
82
  response = st.session_state.conversation({'question': user_question})
 
94
  st.session_state.conversation = None
95
  if "chat_history" not in st.session_state:
96
  st.session_state.chat_history = None
 
 
97
 
98
  if st.session_state.conversation is None:
99
  if os.path.isdir("./chroma_db"):
 
101
  with st.spinner("Loading vector store..."):
102
  vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=OpenAIEmbeddings())
103
  st.session_state.conversation = get_conversation_chain(vectorstore)
104
+ else:
105
+ retrain_model()
106
+ else:
107
+ retrain_model()
 
 
 
 
 
 
 
 
 
 
 
 
108
 
109
  if st.session_state.conversation is not None:
110
+ st.sidebar.button("Retrain model", on_click=retrain_model)
111
  st.header("Ask questions from 48 Laws of Power:books:")
112
+ user_question = st.chat_input("Ask a question about your documents:")
 
 
 
 
113
  if user_question:
114
  handle_userinput(user_question)
115