Update appStore/keyword_search.py

appStore/keyword_search.py

@@ -24,6 +24,43 @@ import numpy as np
 import tempfile
 import sqlite3
 
+#Haystack Components
+@st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None},allow_output_mutation=True)
+def start_haystack(temp.name, file):
+    document_store = InMemoryDocumentStore()
+    documents = pre.load_document(temp.name, file)
+    documents_processed = pre.preprocessing(documents)
+    document_store.write_documents(documents_processed)
+    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
+
+
+def ask_question(question):
+    prediction = pipeline.run(query=question, params={"Retriever": {"top_k": 10}, "Reader": {"top_k": 5}})
+    results = []
+    for answer in prediction["answers"]:
+        answer = answer.to_dict()
+        if answer["answer"]:
+            results.append(
+                {
+                    "context": "..." + 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
+
 def app():
 
     with st.container():
@@ -58,151 +95,45 @@ def app():
                 st.write("Filename: ", file.name)
                 
                 # load document
-                docs = pre.load_document(temp.name, file)
+                pipeline = start_haystack(temp.name, file)
+                #docs = pre.load_document(temp.name, file)
 
                 # preprocess document
-                haystackDoc, dataframeDoc, textData, paraList = clean.preprocessing(docs)
+                #haystackDoc, dataframeDoc, textData, paraList = clean.preprocessing(docs)
                 
-                # testing
-                # st.write(len(all_text))
-                # for i in par_list:
-                #     st.write(i)
+
                 
-                keyword = st.text_input("Please enter here what you want to search, we will look for similar context in the document.",
+                question = st.text_input("Please enter your question here, we will look for the answer in the document.",
                                          value="floods",)
                                          
-                @st.cache(allow_output_mutation=True)
-                def load_sentenceTransformer(name):
-                    return SentenceTransformer(name)
-    
-                bi_encoder = load_sentenceTransformer('msmarco-distilbert-cos-v5') # multi-qa-MiniLM-L6-cos-v1
-                bi_encoder.max_seq_length = 64     #Truncate long passages to 256 tokens
-                top_k = 32
-                
-                #@st.cache(allow_output_mutation=True)
-                #def load_crossEncoder(name):
-                 #   return CrossEncoder(name)
                 
-                # cross_encoder = load_crossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
-                document_embeddings = bi_encoder.encode(paraList, convert_to_tensor=True, show_progress_bar=False)
-
-                def bm25_tokenizer(text):
-                    tokenized_doc = []
-                    for token in text.lower().split():
-                        token = token.strip(string.punctuation)
-
-                        if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:
-                            tokenized_doc.append(token)
-                    return tokenized_doc
+                if st.button("Find them."):
+                     with st.spinner("👑 Performing semantic search on"):#+file.name+"..."):
+                        try:
+                            msg = 'Asked ' + question
+                            logging.info(msg)
+                            st.session_state.results = ask_question(question)    
+                        except Exception as e:
+                            logging.exception(e)
                     
-                def bm25TokenizeDoc(paraList):
-                    tokenized_corpus = []
-                    for passage in tqdm(paraList):
-                        if len(passage.split()) >256:
-                            temp  = " ".join(passage.split()[:256])
-                            tokenized_corpus.append(bm25_tokenizer(temp))
-                            temp  = " ".join(passage.split()[256:])
-                            tokenized_corpus.append(bm25_tokenizer(temp))
-                        else:
-                            tokenized_corpus.append(bm25_tokenizer(passage))
-                           
-                    return tokenized_corpus
                 
-                tokenized_corpus = bm25TokenizeDoc(paraList)
                 
-
-                document_bm25 = BM25Okapi(tokenized_corpus)
-                
-                
-                def search(keyword):
-                    ##### BM25 search (lexical search) #####
-                    bm25_scores = document_bm25.get_scores(bm25_tokenizer(keyword))
-                    top_n = np.argpartition(bm25_scores, -10)[-10:]
-                    bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
-                    bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
-                    
-                    ##### Sematic Search #####
-                    # Encode the query using the bi-encoder and find potentially relevant passages
-                    #query = "Does document contain {} issues ?".format(keyword)
-                    question_embedding = bi_encoder.encode(keyword, convert_to_tensor=True)
-              
-                    hits = util.semantic_search(question_embedding, document_embeddings, top_k=top_k)
-                    hits = hits[0]  # Get the hits for the first query
-                    
-                    
-                    ##### Re-Ranking #####
-                    # Now, score all retrieved passages with the cross_encoder
-                    #cross_inp = [[query, paraList[hit['corpus_id']]] for hit in hits]
-                    #cross_scores = cross_encoder.predict(cross_inp)
-                    
-                    # Sort results by the cross-encoder scores
-                    #for idx in range(len(cross_scores)):
-                     #   hits[idx]['cross-score'] = cross_scores[idx]
-                     
-                    
-                    return bm25_hits, hits
-    
-    
-                if st.button("Find them."):
-                   bm25_hits, hits = search(keyword)     
-
-                   st.markdown("""
-                               We will provide with 2 kind of results. The 'lexical search' and the semantic search. 
-                               """)  
-                    # In the semantic search part we provide two kind of results one with only Retriever (Bi-Encoder) and other the ReRanker (Cross Encoder)           
-                   st.markdown("Top few lexical search (BM25) hits")
-                   for hit in bm25_hits[0:5]:
-                       if hit['score'] > 0.00:   
-                           st.write("\t Score: {:.3f}:  \t{}".format(hit['score'], paraList[hit['corpus_id']].replace("\n", " ")))
-                   
-                   
-                   
-                   
-                   
-                 #   st.table(bm25_hits[0:3])
-                   
-                   st.markdown("\n-------------------------\n")
-                   st.markdown("Top few Bi-Encoder Retrieval hits")
-                   
-                   hits = sorted(hits, key=lambda x: x['score'], reverse=True)
-                   for hit in hits[0:5]:
-                     #  if hit['score'] > 0.45:
-                       st.write("\t Score: {:.3f}:  \t{}".format(hit['score'], paraList[hit['corpus_id']].replace("\n", " ")))
-                   #st.table(hits[0:3]
-                   
-                   #st.markdown("-------------------------")
-                   
-                   #hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
-                   #st.markdown("Top few Cross-Encoder Re-ranker hits")
-                   #for hit in hits[0:3]:
-                      # st.write("\t Score: {:.3f}:  \t{}".format(hit['cross-score'], paraList[hit['corpus_id']].replace("\n", " ")))    
-                   #st.table(hits[0:3]
-                   
-                   
-                   
-                   
-    
-    #for hit in bm25_hits[0:3]:
-      #  print("\t{:.3f}\t{}".format(hit['score'], paraList[hit['corpus_id']].replace("\n", " ")))
-
-    
-
-
-
-
-
-    # Output of top-5 hits from bi-encoder
-    #print("\n-------------------------\n")
-    #print("Top-3 Bi-Encoder Retrieval hits")
-    #hits = sorted(hits, key=lambda x: x['score'], reverse=True)
-    #for hit in hits[0:3]:
-     #   print("\t{:.3f}\t{}".format(hit['score'], paraList[hit['corpus_id']].replace("\n", " ")))
-
-    # Output of top-5 hits from re-ranker
-   # print("\n-------------------------\n")
-    #print("Top-3 Cross-Encoder Re-ranker hits")
-   # hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
-   # for hit in hits[0:3]:
-     #   print("\t{:.3f}\t{}".format(hit['cross-score'], paraList[hit['corpus_id']].replace("\n", " ")))    
+                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!"
+                            )
+