File size: 17,706 Bytes
924fb11
e8ebf39
5b25803
e8ebf39
 
 
 
5b25803
844c34d
99f35d8
452072e
 
 
e8ebf39
9997b7b
844c34d
e8ebf39
0b28b48
 
5b25803
e8ebf39
 
88c1cba
e8ebf39
452ec4c
 
e8ebf39
88c1cba
 
 
e8ebf39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88c1cba
 
 
0f074cc
 
ae04b9d
99f35d8
a2bcc71
99f35d8
25f0178
 
 
9c4d6ae
9997b7b
9c4d6ae
 
924fb11
9c4d6ae
 
 
 
 
 
 
a70fbd3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f8386bf
9c4d6ae
e8ebf39
 
 
0f074cc
048eb6f
 
99f35d8
 
 
 
e8ebf39
9c4d6ae
88c1cba
9997b7b
ad304e2
844c34d
a2bcc71
9997b7b
 
 
 
 
 
a2bcc71
9997b7b
 
 
 
 
 
844c34d
 
b2f5314
844c34d
b2f5314
a2bcc71
63bcda0
 
 
7cdc620
41803fb
a2bcc71
41803fb
 
88c1cba
137e5e2
844c34d
ad304e2
e8ebf39
ae04b9d
5b25803
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ae04b9d
 
 
5b25803
 
ae04b9d
5b25803
e8ebf39
ae04b9d
e8ebf39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5b25803
e8ebf39
 
 
0f074cc
452ec4c
e8ebf39
88c1cba
7d19a12
 
3865d62
 
 
 
 
 
 
c5c0bc4
7d19a12
 
 
 
88c1cba
9c4d6ae
b8dee5c
9c4d6ae
63bcda0
0f074cc
 
 
9997b7b
0f074cc
 
88c1cba
844c34d
452ec4c
9c4d6ae
 
 
9997b7b
 
9c4d6ae
88c1cba
9997b7b
0f074cc
 
9997b7b
0f074cc
 
 
844c34d
9c4d6ae
 
 
 
9997b7b
 
 
88c1cba
 
e8ebf39
99f35d8
a70fbd3
d67901d
 
a2bcc71
 
99f35d8
924fb11
e8ebf39
924fb11
 
 
e8ebf39
924fb11
 
9c4d6ae
aeef8cf
a2bcc71
 
924fb11
 
 
e8ebf39
 
 
 
 
 
 
 
88c1cba
 
 
 
 
 
6915a03
88c1cba
 
 
 
b8dee5c
88c1cba
b8dee5c
88c1cba
 
 
 
e8ebf39
844c34d
7cdc620
9c4d6ae
e8ebf39
f36264a
 
ae04b9d
e8ebf39
 
844c34d
182ca2f
e8ebf39
844c34d
182ca2f
e8ebf39
924fb11
 
 
 
4037844
924fb11
 
e8ebf39
452ec4c
 
 
924fb11
25f0178
 
 
 
 
 
924fb11
f8386bf
d520110
a2bcc71
 
25f0178
 
e8ebf39
 
 
25f0178
 
 
 
924fb11
a70fbd3
 
 
 
 
 
 
 
 
f8386bf
 
 
 
 
 
 
 
 
 
924fb11
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88c1cba
924fb11
 
 
a2bcc71
924fb11
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ad304e2
 
 
 
 
99f35d8
924fb11
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
import base64
import os
import re
from hashlib import blake2b
from tempfile import NamedTemporaryFile

import dotenv
from grobid_quantities.quantities import QuantitiesAPI
from langchain.llms.huggingface_hub import HuggingFaceHub
from langchain.memory import ConversationBufferWindowMemory

dotenv.load_dotenv(override=True)

import streamlit as st
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings

from document_qa.document_qa_engine import DocumentQAEngine
from document_qa.grobid_processors import GrobidAggregationProcessor, decorate_text_with_annotations
from grobid_client_generic import GrobidClientGeneric

if 'rqa' not in st.session_state:
    st.session_state['rqa'] = {}

if 'model' not in st.session_state:
    st.session_state['model'] = None

if 'api_keys' not in st.session_state:
    st.session_state['api_keys'] = {}

if 'doc_id' not in st.session_state:
    st.session_state['doc_id'] = None

if 'loaded_embeddings' not in st.session_state:
    st.session_state['loaded_embeddings'] = None

if 'hash' not in st.session_state:
    st.session_state['hash'] = None

if 'git_rev' not in st.session_state:
    st.session_state['git_rev'] = "unknown"
    if os.path.exists("revision.txt"):
        with open("revision.txt", 'r') as fr:
            from_file = fr.read()
            st.session_state['git_rev'] = from_file if len(from_file) > 0 else "unknown"

if "messages" not in st.session_state:
    st.session_state.messages = []

if 'ner_processing' not in st.session_state:
    st.session_state['ner_processing'] = False

if 'uploaded' not in st.session_state:
    st.session_state['uploaded'] = False

if 'memory' not in st.session_state:
    st.session_state['memory'] = None

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",
        'About': "Upload a scientific article in PDF, ask questions, get insights."
    }
)

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
    st.session_state['doc_id'] = None
    st.session_state['uploaded'] = True
    if st.session_state['memory']:
        st.session_state['memory'].clear()


def clear_memory():
    st.session_state['memory'].clear()


# @st.cache_resource
def init_qa(model, api_key=None):
    ## For debug add: callbacks=[PromptLayerCallbackHandler(pl_tags=["langchain", "chatgpt", "document-qa"])])
    if model == 'chatgpt-3.5-turbo':
        st.session_state['memory'] = ConversationBufferWindowMemory(k=4)
        if api_key:
            chat = ChatOpenAI(model_name="gpt-3.5-turbo",
                              temperature=0,
                              openai_api_key=api_key,
                              frequency_penalty=0.1)
            embeddings = OpenAIEmbeddings(openai_api_key=api_key)

        else:
            chat = ChatOpenAI(model_name="gpt-3.5-turbo",
                              temperature=0,
                              frequency_penalty=0.1)
            embeddings = OpenAIEmbeddings()

    elif model == 'mistral-7b-instruct-v0.1':
        chat = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.1",
                              model_kwargs={"temperature": 0.01, "max_length": 4096, "max_new_tokens": 2048})
        embeddings = HuggingFaceEmbeddings(
            model_name="all-MiniLM-L6-v2")
        st.session_state['memory'] = ConversationBufferWindowMemory(k=4)

    elif model == 'zephyr-7b-beta':
        chat = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta",
                              model_kwargs={"temperature": 0.01, "max_length": 4096, "max_new_tokens": 2048})
        embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        st.session_state['memory'] = None
    else:
        st.error("The model was not loaded properly. Try reloading. ")
        st.stop()
        return

    return DocumentQAEngine(chat, embeddings, grobid_url=os.environ['GROBID_URL'], memory=st.session_state['memory'])


@st.cache_resource
def init_ner():
    quantities_client = QuantitiesAPI(os.environ['GROBID_QUANTITIES_URL'], check_server=True)

    materials_client = GrobidClientGeneric(ping=True)
    config_materials = {
        'grobid': {
            "server": os.environ['GROBID_MATERIALS_URL'],
            'sleep_time': 5,
            'timeout': 60,
            'url_mapping': {
                'processText_disable_linking': "/service/process/text?disableLinking=True",
                # 'processText_disable_linking': "/service/process/text"
            }
        }
    }

    materials_client.set_config(config_materials)

    gqa = GrobidAggregationProcessor(None,
                                     grobid_quantities_client=quantities_client,
                                     grobid_superconductors_client=materials_client
                                     )
    return gqa


gqa = init_ner()


def get_file_hash(fname):
    hash_md5 = blake2b()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def play_old_messages():
    if st.session_state['messages']:
        for message in st.session_state['messages']:
            if message['role'] == 'user':
                with st.chat_message("user"):
                    st.markdown(message['content'])
            elif message['role'] == 'assistant':
                with st.chat_message("assistant"):
                    if mode == "LLM":
                        st.markdown(message['content'], unsafe_allow_html=True)
                    else:
                        st.write(message['content'])


# is_api_key_provided = st.session_state['api_key']

with st.sidebar:
    st.session_state['model'] = model = st.selectbox(
        "Model:",
        options=[
            "chatgpt-3.5-turbo",
            "mistral-7b-instruct-v0.1",
            "zephyr-7b-beta",
            "gpt-4",
            "gpt-4-1106-preview"
        ],
        index=2,
        placeholder="Select model",
        help="Select the LLM model:",
        disabled=st.session_state['doc_id'] is not None or st.session_state['uploaded']
    )

    st.markdown(
        ":warning: [Usage disclaimer](https://github.com/lfoppiano/document-qa/tree/review-interface#disclaimer-on-data-security-and-privacy-%EF%B8%8F) :warning: ")

    if (model == 'mistral-7b-instruct-v0.1' or model == 'zephyr-7b-beta') and model not in st.session_state['api_keys']:
        if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
            api_key = st.text_input('Huggingface API Key', type="password")

            st.markdown("Get it [here](https://huggingface.co/docs/hub/security-tokens)")
        else:
            api_key = os.environ['HUGGINGFACEHUB_API_TOKEN']

        if api_key:
            # st.session_state['api_key'] = is_api_key_provided = True
            if model not in st.session_state['rqa'] or model not in st.session_state['api_keys']:
                with st.spinner("Preparing environment"):
                    st.session_state['api_keys'][model] = api_key
                    # if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
                    #     os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
                    st.session_state['rqa'][model] = init_qa(model)

    elif model == 'chatgpt-3.5-turbo' and model not in st.session_state['api_keys']:
        if 'OPENAI_API_KEY' not in os.environ:
            api_key = st.text_input('OpenAI API Key', type="password")
            st.markdown("Get it [here](https://platform.openai.com/account/api-keys)")
        else:
            api_key = os.environ['OPENAI_API_KEY']

        if api_key:
            if model not in st.session_state['rqa'] or model not in st.session_state['api_keys']:
                with st.spinner("Preparing environment"):
                    st.session_state['api_keys'][model] = api_key
                    if 'OPENAI_API_KEY' not in os.environ:
                        st.session_state['rqa'][model] = init_qa(model, api_key)
                    else:
                        st.session_state['rqa'][model] = init_qa(model)
    # else:
    #     is_api_key_provided = st.session_state['api_key']

    st.button(
        'Reset chat memory.',
        key="reset-memory-button",
        on_click=clear_memory,
        help="Clear the conversational memory. Currently implemented to retrain the 4 most recent messages.",
        disabled=model in st.session_state['rqa'] and st.session_state['rqa'][model].memory is None)

left_column, right_column = st.columns([1, 1])

with right_column:
    st.title("πŸ“ Scientific Document Insights Q/A")
    st.subheader("Upload a scientific article in PDF, ask questions, get insights.")

    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. ")

question = st.chat_input(
    "Ask something about the article",
    # placeholder="Can you give me a short summary?",
    disabled=not uploaded_file
)

with st.sidebar:
    st.header("Settings")
    mode = st.radio("Query mode", ("LLM", "Embeddings"), disabled=not uploaded_file, index=0, horizontal=True,
                    help="LLM will respond the question, Embedding will show the "
                         "paragraphs relevant to the question in the paper.")
    chunk_size = st.slider("Chunks size", 100, 2000, value=250,
                           help="Size of chunks in which the document is partitioned",
                           disabled=uploaded_file is not None)
    context_size = st.slider("Context size", 3, 10, value=4,
                             help="Number of chunks to consider when answering a question",
                             disabled=not uploaded_file)

    st.session_state['ner_processing'] = st.checkbox("Identify materials and properties.")
    st.markdown(
        'The LLM responses undergo post-processing to extract <span style="color:orange">physical quantities, measurements</span>, and <span style="color:green">materials</span> mentions.',
        unsafe_allow_html=True)

    st.divider()

    st.header("Documentation")
    st.markdown("https://github.com/lfoppiano/document-qa")
    st.markdown(
        """Upload a scientific article as PDF document. Once the spinner stops, you can proceed to ask your questions.""")

    if st.session_state['git_rev'] != "unknown":
        st.markdown("**Revision number**: [" + st.session_state[
            'git_rev'] + "](https://github.com/lfoppiano/document-qa/commit/" + st.session_state['git_rev'] + ")")

    st.header("Query mode (Advanced use)")
    st.markdown(
        """By default, the mode is set to LLM (Language Model) which enables question/answering. You can directly ask questions related to the document content, and the system will answer the question using content from the document.""")

    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 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 = []

    # timestamp = datetime.utcnow()

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 len(st.session_state.messages) > 1:
        #     last_answer = st.session_state.messages[len(st.session_state.messages)-1]
        #     if last_answer['role'] == "assistant":
        #         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()