refactor

app.py

@@ -1,8 +1,5 @@
-# modified version of https://github.com/hwchase17/langchain-streamlit-template/blob/master/main.py
-
 import os
 import streamlit as st
-# from streamlit_chat import message
 
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.vectorstores.faiss import FAISS
@@ -12,6 +9,17 @@ from langchain import OpenAI
 from langchain import PromptTemplate
 
 
+BOOK_NAME = "1984"
+AUTHOR_NAME = "George Orwell"
+
+st.set_page_config(page_title="Talk2Book: 1984", page_icon="📖")
+st.title(f"Talk2Book: {BOOK_NAME}")
+st.markdown(f"#### Have a conversaion with {BOOK_NAME} by {AUTHOR_NAME} 🙊")
+
+
+
+
+##### functionss ####
 @st.experimental_singleton(show_spinner=False)
 def load_vectorstore():
     # download from hugging face
@@ -43,12 +51,9 @@ def load_vectorstore():
 
     return docsearch
 
-@st.experimental_singleton(show_spinner=False)
-def load_chain():
-
-    BOOK_NAME = "1984"
-    AUTHOR_NAME = "George Orwell"
 
+@st.experimental_memo(show_spinner=False):
+def load_prompt(book_name, author_name):
     prompt_template = f"""You're an AI version of {AUTHOR_NAME}'s book '{BOOK_NAME}' and are supposed to answer quesions people have for the book. Thanks to advancements in AI people can now talk directly to books.
     People have a lot of questions after reading {BOOK_NAME}, you are here to answer them as you think the author {AUTHOR_NAME} would, using context from the book.
     Where appropriate, briefly elaborate on your answer.
@@ -65,16 +70,22 @@ def load_chain():
         template=prompt_template, input_variables=["context", "question"]
     )
 
+    return PROMPT
+
+    
+@st.experimental_singleton(show_spinner=False)
+def load_chain():
     llm = OpenAI(temperature=0.2)
 
     chain = VectorDBQA.from_chain_type(
-        chain_type_kwargs = {"prompt": PROMPT},
+        chain_type_kwargs = {"prompt": load_prompt(book_name=BOOK_NAME, author_name=AUTHOR_NAME)},
         llm=llm,
         chain_type="stuff", 
         vectorstore=load_vectorstore(),
         k=8,
         return_source_documents=True,
         )
+    
     return chain
 
 
@@ -83,9 +94,7 @@ def get_answer(question):
     result = chain({"query": question})
 
     answer = result["result"]
-
-    # format sources
-
+    
     # pages
     unique_sources = set()
     for item in result['source_documents']:
@@ -95,7 +104,8 @@ def get_answer(question):
     for item in unique_sources:
         unique_pages += str(item) + ", "
 
-    pages = unique_pages # will look like 1, 2, 3,
+    # will look like 1, 2, 3,
+    pages = unique_pages[:-1] # removes the last comma
 
     # source text
     full_source = ""
@@ -110,11 +120,9 @@ def get_answer(question):
     return answer, pages, extract
 
 
-# From here down is all the StreamLit UI.
-st.set_page_config(page_title="Talk2Book: 1984", page_icon="📖")
-st.title("Talk2Book: 1984")
-st.markdown("#### Have a conversaion with 1984 by George Orwell 🙊")
+    
 
+##### sidebar ####
 with st.sidebar:
     api_key = st.text_input(label = "Paste your OpenAI API key here to get started", 
                             type = "password",
@@ -126,19 +134,11 @@ with st.sidebar:
 
     st.info("Based on [Talk2Book](https://github.com/batmanscode/Talk2Book)")
 
-# streamlit-chat not working
-# i get this error: https://discuss.streamlit.io/t/your-app-is-having-trouble-loading-the-xxx-component/25046
-# if "generated" not in st.session_state:
-#     st.session_state["generated"] = []
-
-# if "past" not in st.session_state:
-#     st.session_state["past"] = []
 
-def get_text():
-    user_input = st.text_input("Your question", "Who are you?", key="input")
-    return user_input
+    
 
-user_input = get_text()
+##### main ####
+user_input = st.text_input("Your question", "Who are you?", key="input")
 
 col1, col2 = st.columns([10, 1])
 
@@ -152,30 +152,18 @@ ask = col2.button("Ask")
 if ask:
 
     if api_key is "":
-        # output = "Whoops looks like you forgot your API key buddy"
-        st.write("**1984:** Whoops looks like you forgot your API key buddy")
+        st.write(f"**{BOOK_NAME}:** Whoops looks like you forgot your API key buddy")
         st.stop()
     else:
         with st.spinner("Um... excuse me but... this can take about a minute for your first question because some stuff have to be downloaded 🥺👉🏻👈🏻"):
             try:
                 answer, pages, extract = get_answer(question=user_input)
             except:
-                # output = "What's going on? That's not the right API key"
-                st.write("**1984:** What\'s going on? That's not the right API key")
+                st.write(f"**{BOOK_NAME}:** What\'s going on? That's not the right API key")
                 st.stop()
 
-    st.write(f"**1984:** {answer}")
+    st.write(f"**{BOOK_NAME}:** {answer}")
 
     # sources
     with st.expander(label = f"From pages: {pages}", expanded = False):
-        st.markdown(extract)
-
-# streamlit-chat
-#     st.session_state.past.append(user_input)
-#     st.session_state.generated.append(output)
-
-# if st.session_state["generated"]:
-
-#     for i in range(len(st.session_state["generated"]) - 1, -1, -1):
-#         message(st.session_state["generated"][i], key=str(i))
-#         message(st.session_state["past"][i], is_user=True, key=str(i) + "_user")
+        st.markdown(extract)