Spaces:
Running
Running
import streamlit as st | |
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 validators | |
import json | |
#Haystack Components | |
document_store = InMemoryDocumentStore() | |
retriever = TfidfRetriever(document_store=document_store) | |
reader = FARMReader(model_name_or_path="deepset/tinyroberta-squad2", use_gpu=True) | |
pipeline = ExtractiveQAPipeline(reader, retriever) | |
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?") | |
load_and_write_data() | |
def ask_question(question): | |
prediction = pipeline.run(query=question, params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}}) | |
st.write(Answer(prediction['answers'][0]).to_dict()) | |
st.write(Answer(prediction['answers'][1]).to_dict()) | |
st.write(Answer(prediction['answers'][2]).to_dict()) | |
if question: | |
ask_question(question) |