Merge pull request #16 from lfoppiano/pdf-render

.streamlit/config.toml

@@ -2,4 +2,7 @@
 level = "info"
 
 [browser]
-gatherUsageStats = false
+gatherUsageStats = true
+
+[ui]
+hideTopBar = true

document_qa/document_qa_engine.py

@@ -257,7 +257,7 @@ class DocumentQAEngine:
             self.embeddings_dict[hash] = Chroma.from_texts(texts,
                                                            embedding=self.embedding_function,
                                                            metadatas=metadata,
-                                                           collection_name=hash)
+                                                          collection_name=hash)
         else:
             # if 'documents' in self.embeddings_dict[hash].get() and len(self.embeddings_dict[hash].get()['documents']) == 0:
             #     self.embeddings_dict[hash].delete(ids=self.embeddings_dict[hash].get()['ids'])

streamlit_app.py

@@ -1,3 +1,4 @@
+import base64
 import os
 import re
 from hashlib import blake2b
@@ -55,10 +56,14 @@ if 'uploaded' not in st.session_state:
 if 'memory' not in st.session_state:
     st.session_state['memory'] = ConversationBufferWindowMemory(k=4)
 
+if 'binary' not in st.session_state:
+    st.session_state['binary'] = None
+
 st.set_page_config(
     page_title="Scientific Document Insights Q/A",
     page_icon="📝",
     initial_sidebar_state="expanded",
+    layout="wide",
     menu_items={
         'Get Help': 'https://github.com/lfoppiano/document-qa',
         'Report a bug': "https://github.com/lfoppiano/document-qa/issues",
@@ -66,6 +71,36 @@ st.set_page_config(
     }
 )
 
+css_modify_left_column = '''
+<style>                 
+    [data-testid="stHorizontalBlock"] > div:nth-child(1) {
+        overflow: hidden;
+        background-color: red;
+        height: 70vh;
+    }
+</style>
+'''
+css_modify_right_column = '''
+<style>                 
+    [data-testid="stHorizontalBlock"]> div:first-child {
+        background-color: red;
+        position: fixed
+        height: 70vh;
+    }
+</style>
+'''
+css_disable_scrolling_container = '''
+<style>
+    [data-testid="ScrollToBottomContainer"] {
+        overflow: hidden;
+    }
+</style>
+'''
+
+
+# st.markdown(css_lock_column_fixed, unsafe_allow_html=True)
+# st.markdown(css2, unsafe_allow_html=True)
+
 
 def new_file():
     st.session_state['loaded_embeddings'] = None
@@ -222,18 +257,21 @@ with st.sidebar:
         on_click=clear_memory,
         help="Clear the conversational memory. Currently implemented to retrain the 4 most recent messages.")
 
-st.title("📝 Scientific Document Insights Q/A")
-st.subheader("Upload a scientific article in PDF, ask questions, get insights.")
+left_column, right_column = st.columns([1, 1])
 
-st.markdown(
-    ":warning: Do not upload sensitive data. We **temporarily** store text from the uploaded PDF documents solely for the purpose of processing your request, and we **do not assume responsibility** for any subsequent use or handling of the data submitted to third parties LLMs.")
+with right_column:
+    st.title("📝 Scientific Document Insights Q/A")
+    st.subheader("Upload a scientific article in PDF, ask questions, get insights.")
 
-uploaded_file = st.file_uploader("Upload an article",
+    st.markdown(
+        ":warning: Do not upload sensitive data. We **temporarily** store text from the uploaded PDF documents solely for the purpose of processing your request, and we **do not assume responsibility** for any subsequent use or handling of the data submitted to third parties LLMs.")
+
+    uploaded_file = st.file_uploader("Upload an article",
                                  type=("pdf", "txt"),
                                  on_change=new_file,
-                                 disabled=st.session_state['model'] is not None and st.session_state['model'] not in
-                                          st.session_state['api_keys'],
-                                 help="The full-text is extracted using Grobid. ")
+                                     disabled=st.session_state['model'] is not None and st.session_state['model'] not in
+                                              st.session_state['api_keys'],
+                                     help="The full-text is extracted using Grobid. ")
 
 question = st.chat_input(
     "Ask something about the article",
@@ -276,65 +314,99 @@ with st.sidebar:
     st.markdown(
         """If you switch the mode to "Embedding," the system will return specific chunks from the document that are semantically related to your query. This mode helps to test why sometimes the answers are not satisfying or incomplete. """)
 
+
+@st.cache_resource
+def get_pdf_display(binary):
+    base64_pdf = base64.b64encode(binary).decode('utf-8')
+    return F'<embed src="data:application/pdf;base64,{base64_pdf}" width="100%" height="700" type="application/pdf"></embed>'
+
+
 if uploaded_file and not st.session_state.loaded_embeddings:
     if model not in st.session_state['api_keys']:
         st.error("Before uploading a document, you must enter the API key. ")
         st.stop()
-    with st.spinner('Reading file, calling Grobid, and creating memory embeddings...'):
-        binary = uploaded_file.getvalue()
-        tmp_file = NamedTemporaryFile()
-        tmp_file.write(bytearray(binary))
-        # hash = get_file_hash(tmp_file.name)[:10]
-        st.session_state['doc_id'] = hash = st.session_state['rqa'][model].create_memory_embeddings(tmp_file.name,
-                                                                                                    chunk_size=chunk_size,
+
+    with right_column:
+        with st.spinner('Reading file, calling Grobid, and creating memory embeddings...'):
+            binary = uploaded_file.getvalue()
+            tmp_file = NamedTemporaryFile()
+            tmp_file.write(bytearray(binary))
+            st.session_state['binary'] = binary
+
+            st.session_state['doc_id'] = hash = st.session_state['rqa'][model].create_memory_embeddings(tmp_file.name,
+                                                                                                        chunk_size=chunk_size,
                                                                                                     perc_overlap=0.1,
                                                                                                     include_biblio=True)
-        st.session_state['loaded_embeddings'] = True
-        st.session_state.messages = []
+            st.session_state['loaded_embeddings'] = True
+            st.session_state.messages = []
 
     # timestamp = datetime.utcnow()
 
-if st.session_state.loaded_embeddings and question and len(question) > 0 and st.session_state.doc_id:
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            if message['mode'] == "LLM":
-                st.markdown(message["content"], unsafe_allow_html=True)
-            elif message['mode'] == "Embeddings":
-                st.write(message["content"])
-    if model not in st.session_state['rqa']:
-        st.error("The API Key for the " + model + " is  missing. Please add it before sending any query. `")
-        st.stop()
-
-    with st.chat_message("user"):
-        st.markdown(question)
-        st.session_state.messages.append({"role": "user", "mode": mode, "content": question})
-
-    text_response = None
-    if mode == "Embeddings":
-        with st.spinner("Generating LLM response..."):
-            text_response = st.session_state['rqa'][model].query_storage(question, st.session_state.doc_id,
-                                                                         context_size=context_size)
-    elif mode == "LLM":
-        with st.spinner("Generating response..."):
-            _, text_response = st.session_state['rqa'][model].query_document(question, st.session_state.doc_id,
+with left_column:
+    if st.session_state['binary']:
+        left_column.markdown(get_pdf_display(st.session_state['binary']), unsafe_allow_html=True)
+
+with right_column:
+    # css = '''
+    #     <style>
+    #         [data-testid="column"] {
+    #             overflow: auto;
+    #             height: 70vh;
+    #         }
+    #     </style>
+    #     '''
+    # st.markdown(css, unsafe_allow_html=True)
+
+    # st.markdown(
+    #     """
+    #     <script>
+    #     document.querySelectorAll('[data-testid="column"]').scrollIntoView({behavior: "smooth"});
+    #     </script>
+    #     """,
+    #     unsafe_allow_html=True,
+    # )
+
+    if st.session_state.loaded_embeddings and question and len(question) > 0 and st.session_state.doc_id:
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                if message['mode'] == "LLM":
+                    st.markdown(message["content"], unsafe_allow_html=True)
+                elif message['mode'] == "Embeddings":
+                    st.write(message["content"])
+        if model not in st.session_state['rqa']:
+            st.error("The API Key for the " + model + " is  missing. Please add it before sending any query. `")
+            st.stop()
+
+        with st.chat_message("user"):
+            st.markdown(question)
+            st.session_state.messages.append({"role": "user", "mode": mode, "content": question})
+
+        text_response = None
+        if mode == "Embeddings":
+            with st.spinner("Generating LLM response..."):
+                text_response = st.session_state['rqa'][model].query_storage(question, st.session_state.doc_id,
+                                                                             context_size=context_size)
+        elif mode == "LLM":
+            with st.spinner("Generating response..."):
+                _, text_response = st.session_state['rqa'][model].query_document(question, st.session_state.doc_id,
                                                                              context_size=context_size)
 
-    if not text_response:
-        st.error("Something went wrong. Contact Luca Foppiano (Foppiano.Luca@nims.co.jp) to report the issue.")
-
-    with st.chat_message("assistant"):
-        if mode == "LLM":
-            if st.session_state['ner_processing']:
-                with st.spinner("Processing NER on LLM response..."):
-                    entities = gqa.process_single_text(text_response)
-                    decorated_text = decorate_text_with_annotations(text_response.strip(), entities)
-                    decorated_text = decorated_text.replace('class="label material"', 'style="color:green"')
-                    decorated_text = re.sub(r'class="label[^"]+"', 'style="color:orange"', decorated_text)
-                    text_response = decorated_text
-            st.markdown(text_response, unsafe_allow_html=True)
-        else:
-            st.write(text_response)
-        st.session_state.messages.append({"role": "assistant", "mode": mode, "content": text_response})
+        if not text_response:
+            st.error("Something went wrong. Contact Luca Foppiano (Foppiano.Luca@nims.co.jp) to report the issue.")
+
+        with st.chat_message("assistant"):
+            if mode == "LLM":
+                if st.session_state['ner_processing']:
+                    with st.spinner("Processing NER on LLM response..."):
+                        entities = gqa.process_single_text(text_response)
+                        decorated_text = decorate_text_with_annotations(text_response.strip(), entities)
+                        decorated_text = decorated_text.replace('class="label material"', 'style="color:green"')
+                        decorated_text = re.sub(r'class="label[^"]+"', 'style="color:orange"', decorated_text)
+                        text_response = decorated_text
+                st.markdown(text_response, unsafe_allow_html=True)
+            else:
+                st.write(text_response)
+            st.session_state.messages.append({"role": "assistant", "mode": mode, "content": text_response})
 
         # if len(st.session_state.messages) > 1:
         #     last_answer = st.session_state.messages[len(st.session_state.messages)-1]
@@ -342,5 +414,5 @@ if st.session_state.loaded_embeddings and question and len(question) > 0 and st.
         #         last_question = st.session_state.messages[len(st.session_state.messages)-2]
         #         st.session_state.memory.save_context({"input": last_question['content']}, {"output": last_answer['content']})
 
-elif st.session_state.loaded_embeddings and st.session_state.doc_id:
-    play_old_messages()
+    elif st.session_state.loaded_embeddings and st.session_state.doc_id:
+        play_old_messages()