Update app.py

app.py

@@ -1,39 +1,96 @@
 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 ExtractiveQAPipeline
-from haystack.nodes import FARMReader, TfidfRetriever
+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")
 
-os.environ['TOKENIZERS_PARALLELISM'] ="false"
 
-#Haystack Components
-@st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None},allow_output_mutation=True)
+# Haystack Components
+# @st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None}, allow_output_mutation=True)
+@st.cache_data
 def start_haystack():
-    document_store = InMemoryDocumentStore()
+    # 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)
-    reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2-distilled", use_gpu=True)
-    pipeline = ExtractiveQAPipeline(reader, retriever)
-    return pipeline
+    # 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)
+    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)
 
@@ -41,31 +98,38 @@ set_state_if_absent("results", None)
 def reset_results(*args):
     st.session_state.results = None
 
-#Streamlit App
+
+# 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 👑
+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)
+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}, "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": "..." + answer["context"] + "...",
+                    "context": "..." + str(answer["context"]) + "...",
                     "answer": answer["answer"],
-                    "relevance": round(answer["score"] * 100, 2),
-                    "offset_start_in_doc": answer["offsets_in_document"][0]["start"],
+                    # "relevance": round(answer["score"] * 100, 2),
+                    # "offset_start_in_doc": answer["offsets_in_document"][0]["start"],
                 }
             )
         else:
@@ -73,21 +137,20 @@ def ask_question(question):
                 {
                     "context": None,
                     "answer": None,
-                    "relevance": round(answer["score"] * 100, 2),
+                    # "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)    
+            st.session_state.results = ask_question(question)
         except Exception as e:
             logging.exception(e)
-    
-
 
 if st.session_state.results:
     st.write('## Top Results')
@@ -97,11 +160,14 @@ if st.session_state.results:
             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:]),
+                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']}")
+            # 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!"
+                "🤔    Haystack is unsure whether any of the documents contain an "
+                "answer to your question. Try to reformulate it!"
             )