Attempting sessions states and annotations

app.py

@@ -4,6 +4,9 @@ from haystack.schema import Answer
 from haystack.document_stores import InMemoryDocumentStore
 from haystack.pipelines import ExtractiveQAPipeline
 from haystack.nodes import FARMReader, TfidfRetriever
+import logging
+from markdown import markdown
+from annotated_text import annotation
 import validators
 import json
 
@@ -13,24 +16,82 @@ retriever = TfidfRetriever(document_store=document_store)
 reader = FARMReader(model_name_or_path="deepset/tinyroberta-squad2", use_gpu=True)
 pipeline = ExtractiveQAPipeline(reader, retriever)
 
+def set_state_if_absent(key, value):
+    if key not in st.session_state:
+        st.session_state[key] = value
+
+
+st.set_page_config(page_title="Game of Thrones QA with Haystack", page_icon="https://haystack.deepset.ai/img/HaystackIcon.png")
+
+set_state_if_absent("question", "Who is Arya's father")
+set_state_if_absent("results", None)
+
+
+def reset_results(*args):
+    st.session_state.results = None
+
 def load_and_write_data():
     doc_dir = './article_txt_got'
     docs = convert_files_to_docs(dir_path=doc_dir, clean_func=clean_wiki_text, split_paragraphs=True)
 
     document_store.write_documents(docs)
 
+
 #Streamlit App
 
 st.title('Game of Thrones QA with Haystack')
-question = st.text_input(label="Ask a Question about Game of Thromes", value="Who is Arya's father?")
+question = st.text_input("", value=st.session_state.question, max_chars=100, on_change=reset_results)
 
 load_and_write_data()
 
 def ask_question(question):
     prediction = pipeline.run(query=question, params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}})
-    st.write(prediction['answers'][0].to_dict())
-    st.write(prediction['answers'][1].to_dict())
-    st.write(prediction['answers'][2].to_dict())
+    results = []
+    for answer in prediction["answers"]:
+        if answer.get("answer", None):
+            results.append(
+                {
+                    "context": "..." + answer["context"] + "...",
+                    "answer": answer.get("answer", None),
+                    "relevance": round(answer["score"] * 100, 2),
+                    "offset_start_in_doc": answer["offsets_in_document"][0]["start"],
+                }
+            )
+        else:
+            results.append(
+                {
+                    "context": None,
+                    "answer": None,
+                    "relevance": round(answer["score"] * 100, 2),
+                }
+            )
+    return results
+    # st.write(prediction['answers'][0].to_dict())
+    # st.write(prediction['answers'][1].to_dict())
+    # st.write(prediction['answers'][2].to_dict())
+    
 
 if question:
-    ask_question(question)
+    try:
+        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(answer, "ANSWER", "#8ef")) + context[end_idx:]),
+                unsafe_allow_html=True,
+            )
+            st.markdown(f"**Relevance:** {result['relevance']}")
+        else:
+            st.info(
+                "🤔    Haystack is unsure whether any of the documents contain an answer to your question. Try to reformulate it!"
+            )

requirements.txt

@@ -1,2 +1,3 @@
 farm-haystack==1.4.0
-validators==0.18.2
+validators==0.18.2
+markdown