File size: 6,219 Bytes
2b95436
 
b0df4b4
 
2b95436
 
 
b0df4b4
 
 
 
 
2b95436
 
 
 
b0df4b4
 
 
 
 
2b95436
 
b0df4b4
9975717
 
2b95436
b0df4b4
 
 
 
 
2b95436
b0df4b4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2b95436
 
 
b0df4b4
 
2b95436
 
 
b0df4b4
2b95436
 
b0df4b4
2b95436
 
 
 
b0df4b4
2b95436
 
 
 
 
 
 
b0df4b4
 
2b95436
 
 
 
b0df4b4
 
 
 
 
2b95436
 
 
b0df4b4
 
 
2b95436
 
b0df4b4
 
2b95436
 
 
 
b0df4b4
2b95436
 
b0df4b4
2b95436
b0df4b4
 
2b95436
 
 
 
 
 
 
b0df4b4
2b95436
 
 
 
b0df4b4
2b95436
 
 
 
 
b0df4b4
2b95436
 
 
 
 
 
 
 
 
 
 
b0df4b4
 
 
2b95436
 
b0df4b4
2b95436
 
b0df4b4
 
2b95436
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
import streamlit as st
import os

from haystack import Pipeline
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 DocumentSearchPipeline, ExtractiveQAPipeline, GenerativeQAPipeline
from haystack.nodes import (DensePassageRetriever, EmbeddingRetriever, FARMReader,
                            OpenAIAnswerGenerator, Seq2SeqGenerator,
                            TfidfRetriever)
from haystack.nodes import RAGenerator
import logging
from markdown import markdown
from annotated_text import annotation
from PIL import Image
logging.basicConfig(format="%(levelname)s - %(name)s -  %(message)s", level=logging.WARNING)
logging.getLogger("haystack").setLevel(logging.INFO)

os.environ['TOKENIZERS_PARALLELISM'] = "false"
MY_API_KEY = os.environ.get("MY_API_KEY")


# Haystack Components
@st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None}, allow_output_mutation=True)
# @st.cache_data
def start_haystack():
    # document_store = InMemoryDocumentStore()
    # For dense retriever
    document_store = InMemoryDocumentStore(embedding_dim=128)
    # For OPEN AI retriever
    # document_store = InMemoryDocumentStore(embedding_dim=1024)
    load_and_write_data(document_store)
    # retriever = TfidfRetriever(document_store=document_store)
    retriever = DensePassageRetriever(
        document_store=document_store,
        query_embedding_model="vblagoje/dpr-question_encoder-single-lfqa-wiki",
        passage_embedding_model="vblagoje/dpr-ctx_encoder-single-lfqa-wiki",
    )

    # retriever = EmbeddingRetriever(
    #     document_store=document_store,
    #     embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1",
    #     model_format="sentence_transformers",
    # )
    document_store.update_embeddings(retriever)

    # OPEN AI
    # retriever = EmbeddingRetriever(
    #     document_store=document_store,
    #     batch_size=8,
    #     embedding_model="ada",
    #     api_key=MY_API_KEY,
    #     max_seq_len=1024
    # )
    # document_store.update_embeddings(retriever)

    # reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=True)
    # pipeline = ExtractiveQAPipeline(reader, retriever)

    generator = Seq2SeqGenerator(model_name_or_path="vblagoje/bart_lfqa")
    # generator = OpenAIAnswerGenerator(
    #     api_key=MY_API_KEY,
    #     model="text-davinci-003",
    #     max_tokens=50,
    #     presence_penalty=0.1,
    #     frequency_penalty=0.1,
    #     top_k=3,
    #     temperature=0.9
    # )
    # pipe.add_node(component=retriever, name="Retriever", inputs=["Query"])
    # pipe.add_node(component=generator, name="prompt_node", inputs=["Query"])
    pipe = GenerativeQAPipeline(generator=generator, retriever=retriever)

    return pipe


def load_and_write_data(document_store):
    doc_dir = './dao_data'
    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 the goal of VitaDAO?")
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}})
    prediction = pipeline.run(query=question, params={"Retriever": {"top_k": 10}, "Generator": {"top_k": 1}})
    results = []
    for answer in prediction["answers"]:
        answer = answer.to_dict()
        if answer["answer"]:
            print(answer)
            results.append(
                {
                    "context": "..." + str(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!"
            )