feat(app): Run formatting, add caching and loading states

app.py

@@ -7,40 +7,111 @@ from haystack.document_stores.in_memory import InMemoryDocumentStore
 
 from utils import get_unique_docs
 
+
 # Load the dataset
-unique_docs = set()
-dataset = load_dataset("PedroCJardim/QASports", "basketball")
-docs_validation = get_unique_docs(dataset["validation"], unique_docs)
-docs_train = get_unique_docs(dataset["train"], unique_docs)
-docs_test = get_unique_docs(dataset["test"], unique_docs)
-docs_all = docs_validation + docs_train + docs_test
-
-# Create the Question Answering pipeline
-# Create in memory database
-document_store = InMemoryDocumentStore()
-document_store.write_documents(documents=docs_all)
-# Create the retriever and reader
-retriever = InMemoryBM25Retriever(document_store=document_store)
-reader = ExtractiveReader(model="laurafcamargos/distilbert-qasports-basket-small")
-reader.warm_up()
-# Create the pipeline
-pipe = Pipeline()
-pipe.add_component(instance=retriever, name="retriever")
-pipe.add_component(instance=reader, name="reader")
-pipe.connect("retriever.documents", "reader.documents")
+@st.cache_data(show_spinner=False)
+def load_documents():
+    """
+    Load the documents from the dataset considering only unique documents.
+
+    Returns:
+    - documents: list of dictionaries with the documents.
+    """
+    unique_docs = set()
+    dataset_name = "PedroCJardim/QASports"
+    dataset_split = "basketball"
+    st.caption(f'Fetching "{dataset_name}" dataset')
+    # build the dataset
+    dataset = load_dataset(dataset_name, dataset_split)
+    docs_validation = get_unique_docs(dataset["validation"], unique_docs)
+    docs_train = get_unique_docs(dataset["train"], unique_docs)
+    docs_test = get_unique_docs(dataset["test"], unique_docs)
+    documents = docs_validation + docs_train + docs_test
+    return documents
+
+
+@st.cache_data(show_spinner=False)
+def get_document_store(documents):
+    """
+    Index the files in the document store.
+
+    Args:
+    - files: list of dictionaries with the documents.
+    """
+    # Create in memory database
+    st.caption(f"Building the Document Store")
+    document_store = InMemoryDocumentStore()
+    document_store.write_documents(documents=documents)
+    return document_store
+
+
+@st.cache_data(show_spinner=False)
+def get_question_pipeline(_doc_store):
+    """
+    Create the pipeline with the retriever and reader components.
+
+    Args:
+    - doc_store: instance of the document store.
+
+    Returns:
+    - pipe: instance of the pipeline.
+    """
+    st.caption(f"Building the Question Answering pipeline")
+    # Create the retriever and reader
+    retriever = InMemoryBM25Retriever(document_store=_doc_store)
+    reader = ExtractiveReader(model="laurafcamargos/distilbert-qasports-basket-small")
+    reader.warm_up()
+    # Create the pipeline
+    pipe = Pipeline()
+    pipe.add_component(instance=retriever, name="retriever")
+    pipe.add_component(instance=reader, name="reader")
+    pipe.connect("retriever.documents", "reader.documents")
+    return pipe
+
+
+# # Create the retriever and reader
+# retriever = InMemoryBM25Retriever(document_store=document_store())
+# reader = ExtractiveReader(model="laurafcamargos/distilbert-qasports-basket-small")
+# reader.warm_up()
+# # Create the pipeline
+# pipe = Pipeline()
+# pipe.add_component(instance=retriever, name="retriever")
+# pipe.add_component(instance=reader, name="reader")
+# pipe.connect("retriever.documents", "reader.documents")
 
 # Streamlit interface
-st.markdown("""This website presents a collection of documents from the dataset named "QASports", the first large sports question answering dataset for open questions. QASports contains real data of players, teams and matches from the sports soccer, basketball and American football. It counts over 1.5 million questions and answers about 54k preprocessed, cleaned and organized documents from Wikipedia-like sources.""")
-st.subheader('QASports: Basketball', divider='rainbow')
-
-top_k = 3
-user_query = None
-user_query = st.text_input("Please, make a question about basketball:")
-
-if user_query:
-    answer = pipe.run(data={
-        "retriever": {"query": user_query, "top_k": 10},
-        "reader": {"query": user_query, "top_k": top_k},
-    })
-    # Display only the top k answers
-    st.json(answer["reader"]["answers"][0:top_k])
+with st.status(
+    "Downloading dataset...", expanded=st.session_state.get("expanded", True)
+) as status:
+    documents = load_documents()
+    status.update(label="Indexing documents...")
+    doc_store = get_document_store(documents)
+    status.update(label="Creating pipeline...")
+    pipe = get_question_pipeline(doc_store)
+    status.update(
+        label="Download and indexing complete!", state="complete", expanded=False
+    )
+    st.session_state["expanded"] = False
+
+st.subheader("🔎 QASports: Basketball", divider="rainbow")
+st.caption(
+    """This website presents a collection of documents from the dataset named "QASports", the first large sports question answering dataset for open questions. QASports contains real data of players, teams and matches from the sports soccer, basketball and American football. It counts over 1.5 million questions and answers about 54k preprocessed, cleaned and organized documents from Wikipedia-like sources."""
+)
+
+if user_query := st.text_input(
+    label="What do you want to know about Basketball?",
+    placeholder="How many field goals did Kobe Bryant score?",
+):
+    try:
+        top_k = 3
+        answer = pipe.run(
+            data={
+                "retriever": {"query": user_query, "top_k": 10},
+                "reader": {"query": user_query, "top_k": top_k},
+            }
+        )
+        # Display only the top k answers
+        max_k = len(answer["reader"]["answers"])
+        st.json(answer["reader"]["answers"][0 : min(top_k, max_k)])
+    except Exception as e:
+        st.error(f"Error: We do not have an answer for your question.")

utils.py

@@ -1,17 +1,18 @@
-'''This module contains utility functions for the project'''
+"""This module contains utility functions for the project"""
+
 import mmh3
 from haystack import Document
 
 
-def get_unique_docs(dataset, unique_docs:set):
-    '''Get unique documents from dataset
-    
+def get_unique_docs(dataset, unique_docs: set):
+    """Get unique documents from dataset
+
     Args:
     dataset: list of dictionaries
 
     Returns:
     docs: list of haystack.Document
-    '''
+    """
     docs = list()
     for doc in dataset:
         if doc["context"] is not None and doc["context_id"] not in unique_docs:
@@ -19,11 +20,12 @@ def get_unique_docs(dataset, unique_docs:set):
             document = Document(
                 content=doc["context"],
                 meta={
-                    'title': doc["context_title"],
-                    'context_id': doc["context_id"],
-                    'url': doc["url"],
-                    'source': 'QASports', 'category': 'basketball'
-                }
+                    "title": doc["context_title"],
+                    "context_id": doc["context_id"],
+                    "url": doc["url"],
+                    "source": "QASports",
+                    "category": "basketball",
+                },
             )
             docs.append(document)
     return docs