Spaces:
Runtime error
Runtime error
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_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!" | |
) | |