import streamlit as st import os from haystack.utils import fetch_archive_from_http, clean_wiki_text, convert_files_to_docs 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 from PIL import Image os.environ['TOKENIZERS_PARALLELISM'] = "false" # Haystack Components @st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None}, allow_output_mutation=True) def start_haystack(): document_store = InMemoryDocumentStore() load_and_write_data(document_store) retriever = TfidfRetriever(document_store=document_store) reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2-distilled", use_gpu=True) pipeline = ExtractiveQAPipeline(reader, retriever) return pipeline def load_and_write_data(document_store): doc_dir = './amazon_help_docs' docs = convert_files_to_docs(dir_path=doc_dir, clean_func=clean_wiki_text, split_paragraphs=True) document_store.write_documents(docs) pipeline = start_haystack() def set_state_if_absent(key, value): if key not in st.session_state: st.session_state[key] = value set_state_if_absent("question", "What is amazon music?") set_state_if_absent("results", None) def reset_results(*args): st.session_state.results = None # Streamlit App image = Image.open('got-haystack.png') st.image(image) st.markdown(""" 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 documents about Game of Thrones 👑 Go ahead and ask questions about the marvellous kingdom! """, unsafe_allow_html=True) question = st.text_input("", value=st.session_state.question, max_chars=100, on_change=reset_results) def ask_question(question): prediction = pipeline.run(query=question, params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}}) results = [] for answer in prediction["answers"]: answer = answer.to_dict() if answer["answer"]: results.append( { "context": "..." + answer["context"] + "...", "answer": answer["answer"], "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 if question: with st.spinner("👑    Performing semantic search on royal scripts..."): 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="ANSWER", background="#964448", color='#ffffff')) + 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!" )