working overhaul

app.py

@@ -1,21 +1,99 @@
 import streamlit as st
-import time
-from pipeline import pipeline
-
-st.title('Welcome to SRM RP explorer!')
-query = st.text_input('Enter a query to get started:', 'Applications of AI and deep learning')
-noof_results = st.slider('Number of results required', 1, 3, 1)
-
-start = time.time()
-pred = pipeline.run(query=query, params={"Retriever": {"top_k": 5}, "Reader": {"top_k": noof_results}})
-st.write('Time taken: %s s' % round(time.time()-start, 2))
-
-for i in pred['answers']:
-
-    ans, name, link, context, score = i.answer, i.meta['name'],i.meta['link'], i.context, i.score
-    st.write(f'Score: {score}')
-    st.write(f'Answer: {ans}')
-    st.write(f'Title: {name}')
-    # st.write(f'Abstract: {context}')
-    st.write(f'Link: {link}')
-    st.write(f'-------------')
+from markdown import markdown
+from annotated_text import annotation
+import logging
+
+from haystack.document_stores import InMemoryDocumentStore
+from haystack.nodes import TfidfRetriever
+from haystack.pipelines import ExtractiveQAPipeline
+from haystack.nodes import FARMReader
+import joblib
+
+@st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None},allow_output_mutation=True)
+def create_pipeline():
+    docs = joblib.load('docs.joblib')
+
+    document_store = InMemoryDocumentStore()
+    document_store.write_documents(docs)
+    
+    retriever = TfidfRetriever(document_store)
+    reader = FARMReader(model_name_or_path="ixa-ehu/SciBERT-SQuAD-QuAC")
+    
+    pipeline = ExtractiveQAPipeline(reader, retriever)
+    
+    return pipeline
+    
+pipeline = create_pipeline()
+
+def set_state_if_absent(key, value):
+    if key not in st.session_state:
+        st.session_state[key] = value
+
+set_state_if_absent("question", 'Applications of AI and deep learning')
+set_state_if_absent("results", None)
+
+def reset_results(*args):
+    st.session_state.results = None
+
+st.markdown('''#Welcome to **SRM RP explorer**!
+This QA demo uses a [Haystack Extractive QA Pipeline](https://haystack.deepset.ai/components/ready-made-pipelines#extractiveqapipeline) with 
+an [InMemoryDocumentStore](https://haystack.deepset.ai/components/document-store) which contains abstracts of 17k+ research papers associated with SRM university''')
+
+query = st.text_input('Enter a query to get started:', value=st.session_state.question, max_chars=100, on_change=reset_results)
+
+def ask_question(query):
+    start = time.time()
+    prediction = pipeline.run(query=query, params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}})
+    st.write('Time taken: %s s' % round(time.time()-start, 2))
+    
+    results = []
+    for answer in prediction["answers"]:
+        answer = answer.to_dict()
+        if answer["answer"]:
+            results.append(
+                {
+                    "title":answer["meta"]["name"],
+                    "link":answer["meta"]["link"],
+                    "context": "..." + answer["context"] + "...",
+                    "answer": answer["answer"],
+                    "score": round(answer["score"] * 100, 2),
+                    "offset_start_in_doc": answer["offsets_in_document"][0]["start"],
+                }
+            )
+        else:
+            results.append(
+                {
+                    "title":None,
+                    "link":None,
+                    "context": None,
+                    "answer": None,
+                    "relevance": round(answer["score"] * 100, 2),
+                }
+            )
+    return results
+
+if query:
+    with st.spinner("👑    Performing semantic search on abstracts..."):
+        try:
+            msg = 'Asked ' + question
+            logging.info(msg)
+            st.session_state.results = ask_question(question)    
+        except Exception as e:
+            logging.exception(e)
+
+if st.session_state.results:
+    st.write('## Top Results')
+    for count, result in enumerate(st.session_state.results):
+        if result["answer"]:
+            answer, context = result["answer"], result["context"]
+            start_idx = context.find(answer)
+            end_idx = start_idx + len(answer)
+            st.write(
+                markdown(context[:start_idx] + str(annotation(body=answer, label="RELEVANT", background="#964448", color='#ffffff')) + context[end_idx:]),
+                unsafe_allow_html=True,
+            )
+            st.markdown(f"**Title:** [{result['name']}]({result['link']})\n**Relevance:** {result['relevance']}")
+        else:
+            st.info(
+                "🤔    Haystack is unsure whether any of the documents contain an answer to your question. Try to reformulate it!"
+            )