test-Haystack / app.py
eruizgar91's picture
Update app.py
9975717
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!"
)