Fixed the pinecone retrieval issue

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+.streamlit/

app.py

@@ -6,12 +6,13 @@ import sys
 import uuid
 from json import JSONDecodeError
 from pathlib import Path
+from typing import List, Optional
 
 import pandas as pd
 import pinecone
 import streamlit as st
 from annotated_text import annotation
-from haystack import Document
+from haystack import BaseComponent, Document
 from haystack.document_stores import PineconeDocumentStore
 from haystack.nodes import (
     DocxToTextConverter,
@@ -26,22 +27,48 @@ from haystack.pipelines import ExtractiveQAPipeline, Pipeline
 from markdown import markdown
 from sentence_transformers import SentenceTransformer
 
+
+class PineconeSearch(BaseComponent):
+    outgoing_edges = 1
+
+    def run(self, query: str, top_k: Optional[int]):
+        # process the inputs
+        vector_embedding = emb_model.encode(query).tolist()
+        response = index.query([vector_embedding], top_k=top_k, include_metadata=True)
+        docs = [
+            Document(
+                content=d["metadata"]["text"],
+                meta={
+                    "title": d["metadata"]["filename"],
+                    "context": d["metadata"]["text"],
+                    "_split_id": d["metadata"]["_split_id"],
+                },
+            )
+            for d in response["matches"]
+        ]
+        output = {"documents": docs, "query": query}
+        return output, "output_1"
+
+    def run_batch(self, queries: List[str], top_k: Optional[int]):
+
+        return {}, "output_1"
+
+
 # connect to pinecone environment
-pinecone.init(
-    api_key=st.secrets["pinecone_apikey"],
-    environment="us-west1-gcp"
-)
+pinecone.init(api_key=st.secrets["pinecone_apikey"], environment="us-west1-gcp")
 index_name = "qa-demo-fast-384"
 # retriever_model = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
 retriever_model = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"
-embedding_dim=384
+emb_model = SentenceTransformer(retriever_model)
+
+embedding_dim = 384
 preprocessor = PreProcessor(
     clean_empty_lines=True,
     clean_whitespace=True,
     clean_header_footer=False,
     split_by="word",
     split_length=100,
-    split_respect_sentence_boundary=True
+    split_respect_sentence_boundary=True,
 )
 file_type_classifier = FileTypeClassifier()
 text_converter = TextConverter()
@@ -53,65 +80,50 @@ if index_name not in pinecone.list_indexes():
     # delete the current index and create the new index if it does not exist
     for delete_index in pinecone.list_indexes():
         pinecone.delete_index(delete_index)
-    pinecone.create_index(
-        index_name,
-        dimension=embedding_dim,
-        metric="cosine"
-    )
+    pinecone.create_index(index_name, dimension=embedding_dim, metric="cosine")
 
 # connect to abstractive-question-answering index we created
 index = pinecone.Index(index_name)
 
-FILE_UPLOAD_PATH= "./data/uploads/"
+FILE_UPLOAD_PATH = "./data/uploads/"
 os.makedirs(FILE_UPLOAD_PATH, exist_ok=True)
-# @st.cache
+
+
 def create_doc_store():
     document_store = PineconeDocumentStore(
-        api_key= st.secrets["pinecone_apikey"],
+        api_key=st.secrets["pinecone_apikey"],
         index=index_name,
         similarity="cosine",
-        embedding_dim=embedding_dim
+        embedding_dim=embedding_dim,
     )
     return document_store
 
-# @st.cache
-# def create_pipe(document_store):
-    # retriever = EmbeddingRetriever(
-    # document_store=document_store,
-    # embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1",
-    # model_format="sentence_transformers",
-    # )
-    # reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=False)
-    # pipe = ExtractiveQAPipeline(reader, retriever)
-    # return pipe
 
 def query(pipe, question, top_k_reader, top_k_retriever):
     res = pipe.run(
-        query=question, params={"Retriever": {"top_k": top_k_retriever}, "Reader": {"top_k": top_k_reader}}
+        query=question,
+        params={"Retriever": {"top_k": top_k_retriever}, "Reader": {"top_k": top_k_reader}},
     )
-    answer_df = []
-    # for r in res['answers']:
-    #     ans_dict = res['answers'][0].meta
-    #     ans_dict["answer"] = r.context
-    #     answer_df.append(ans_dict)
-    # result = pd.DataFrame(answer_df)
-    # result.columns = ["Source","Title","Year","Link","Answer"]
-    # result[["Answer","Link","Source","Title","Year"]] 
     return res
 
+
 document_store = create_doc_store()
 # pipe = create_pipe(document_store)
 
 retriever = EmbeddingRetriever(
-document_store=document_store,
-embedding_model=retriever_model,
-model_format="sentence_transformers",
+    document_store=document_store,
+    embedding_model=retriever_model,
+    model_format="sentence_transformers",
 )
 # load the retriever model from huggingface model hub
 sentence_encoder = SentenceTransformer(retriever_model)
 
 reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=False)
-pipe = ExtractiveQAPipeline(reader, retriever)
+# pipe = ExtractiveQAPipeline(reader, retriever)
+# Custom built extractive QA pipeline
+pipe = Pipeline()
+pipe.add_node(component=PineconeSearch(), name="Retriever", inputs=["Query"])
+pipe.add_node(component=reader, name="Reader", inputs=["Retriever"])
 
 
 indexing_pipeline_with_classification = Pipeline()
@@ -133,20 +145,29 @@ indexing_pipeline_with_classification.add_node(
     inputs=["TextConverter", "PdfConverter", "DocxConverter"],
 )
 
+
 def set_state_if_absent(key, value):
     if key not in st.session_state:
         st.session_state[key] = value
 
+
 # Adjust to a question that you would like users to see in the search bar when they load the UI:
-DEFAULT_QUESTION_AT_STARTUP = os.getenv("DEFAULT_QUESTION_AT_STARTUP", "My blog post discusses remote work. Give me statistics.")
-DEFAULT_ANSWER_AT_STARTUP = os.getenv("DEFAULT_ANSWER_AT_STARTUP", "7% more remote workers have been at their current organization for 5 years or fewer")
+DEFAULT_QUESTION_AT_STARTUP = os.getenv(
+    "DEFAULT_QUESTION_AT_STARTUP", "My blog post discusses remote work. Give me statistics."
+)
+DEFAULT_ANSWER_AT_STARTUP = os.getenv(
+    "DEFAULT_ANSWER_AT_STARTUP",
+    "7% more remote workers have been at their current organization for 5 years or fewer",
+)
 
 # Sliders
 DEFAULT_DOCS_FROM_RETRIEVER = int(os.getenv("DEFAULT_DOCS_FROM_RETRIEVER", "3"))
 DEFAULT_NUMBER_OF_ANSWERS = int(os.getenv("DEFAULT_NUMBER_OF_ANSWERS", "3"))
 
 
-st.set_page_config(page_title="Haystack Demo", page_icon="https://haystack.deepset.ai/img/HaystackIcon.png")
+st.set_page_config(
+    page_title="Haystack Demo", page_icon="https://haystack.deepset.ai/img/HaystackIcon.png"
+)
 
 # Persistent state
 set_state_if_absent("question", DEFAULT_QUESTION_AT_STARTUP)
@@ -160,6 +181,7 @@ def reset_results(*args):
     st.session_state.results = None
     st.session_state.raw_json = None
 
+
 # Title
 st.write("# Haystack Search Demo")
 st.markdown(
@@ -187,12 +209,16 @@ for data_file in data_files:
             f.write(data_file.getbuffer())
         ALL_FILES.append(file_path)
         st.sidebar.write(str(data_file.name) + "    ✅ ")
-        META_DATA.append({"filename":data_file.name})
-        
+        META_DATA.append({"filename": data_file.name})
+
+data_files = []
+
 
 if len(ALL_FILES) > 0:
     # document_store.update_embeddings(retriever, update_existing_embeddings=False)
-    docs = indexing_pipeline_with_classification.run(file_paths=ALL_FILES, meta=META_DATA)["documents"]
+    docs = indexing_pipeline_with_classification.run(file_paths=ALL_FILES, meta=META_DATA)[
+        "documents"
+    ]
     index_name = "qa_demo"
     # we will use batches of 64
     batch_size = 128
@@ -204,7 +230,7 @@ if len(ALL_FILES) > 0:
     upload_count = 0
     for i in range(0, len(docs), batch_size):
         # find end of batch
-        i_end = min(i+batch_size, len(docs))
+        i_end = min(i + batch_size, len(docs))
         # extract batch
         batch = [doc.content for doc in docs[i:i_end]]
         # generate embeddings for batch
@@ -222,10 +248,10 @@ if len(ALL_FILES) > 0:
         to_upsert = list(zip(ids, emb, meta))
         # upsert/insert these records to pinecone
         _ = index.upsert(vectors=to_upsert)
-        upload_count+=batch_size
-        upload_percentage = min(int((upload_count/len(docs))*100), 100)
+        upload_count += batch_size
+        upload_percentage = min(int((upload_count / len(docs)) * 100), 100)
         my_bar.progress(upload_percentage)
-    
+
 top_k_reader = st.sidebar.slider(
     "Max. number of answers",
     min_value=1,
@@ -251,12 +277,12 @@ top_k_retriever = st.sidebar.slider(
 #         raw_json = upload_doc(data_file)
 
 question = st.text_input(
-        value=st.session_state.question,
-        max_chars=100,
-        on_change=reset_results,
-        label="question",
-        label_visibility="hidden",
-    )
+    value=st.session_state.question,
+    max_chars=100,
+    on_change=reset_results,
+    label="question",
+    label_visibility="hidden",
+)
 col1, col2 = st.columns(2)
 col1.markdown("<style>.stButton button {width:100%;}</style>", unsafe_allow_html=True)
 col2.markdown("<style>.stButton button {width:100%;}</style>", unsafe_allow_html=True)
@@ -265,23 +291,21 @@ col2.markdown("<style>.stButton button {width:100%;}</style>", unsafe_allow_html
 run_pressed = col1.button("Run")
 if run_pressed:
 
-    run_query = (
-        run_pressed or question != st.session_state.question
-    )
+    run_query = run_pressed or question != st.session_state.question
     # Get results for query
     if run_query and question:
         reset_results()
         st.session_state.question = question
 
-        with st.spinner(
-            "🧠 &nbsp;&nbsp; Performing neural search on documents... \n "
-        ):
+        with st.spinner("🧠 &nbsp;&nbsp; Performing neural search on documents... \n "):
             try:
-                st.session_state.results  = query(
+                st.session_state.results = query(
                     pipe, question, top_k_reader=top_k_reader, top_k_retriever=top_k_retriever
                 )
             except JSONDecodeError as je:
-                st.error("👓 &nbsp;&nbsp; An error occurred reading the results. Is the document store working?")
+                st.error(
+                    "👓 &nbsp;&nbsp; An error occurred reading the results. Is the document store working?"
+                )
             except Exception as e:
                 logging.exception(e)
                 if "The server is busy processing requests" in str(e) or "503" in str(e):
@@ -294,23 +318,24 @@ if st.session_state.results:
 
     st.write("## Results:")
 
-    for count, result in enumerate(st.session_state.results['answers']):
+    for count, result in enumerate(st.session_state.results["answers"]):
         answer, context = result.answer, result.context
         start_idx = context.find(answer)
         end_idx = start_idx + len(answer)
         # Hack due to this bug: https://github.com/streamlit/streamlit/issues/3190
         try:
-            source = f"[{result.meta['Title']}]({result.meta['link']})"
+            filename = result.meta["title"]
             st.write(
-            markdown(f'**Source:** {source} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '),
-            unsafe_allow_html=True,
-             )
+                markdown(
+                    f'From file:  {filename} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '
+                ),
+                unsafe_allow_html=True,
+            )
         except:
-            filename = result.meta.get('filename', "")
+            filename = result.meta.get("filename", "")
             st.write(
-            markdown(f'From file: {filename} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '),
-            unsafe_allow_html=True,
+                markdown(
+                    f'From file: {filename} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '
+                ),
+                unsafe_allow_html=True,
             )
-
-        
-