Update app.py

app.py

@@ -1,14 +1,14 @@
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
 import re
-import newspaper
+import pandas as pd
+from newspaper import Article
 import docx2txt
 from io import StringIO
 from PyPDF2 import PdfFileReader
 import validators
-import streamlit as st
 import nltk
-import pandas as pd
-import requests
+import warnings
+import streamlit as st
 
 nltk.download('punkt')
 
@@ -17,40 +17,62 @@ from nltk import sent_tokenize
 warnings.filterwarnings("ignore")
 
 
-def extarct_test_from_url(url: str):
-    article = Article(url,language-"en")
+def extract_text_from_url(url: str):
+    '''Extract text from url'''
+
+    article = Article(url)
     article.download()
-    article.parse
+    article.parse()
 
-    # receiving text
+    # get text
     text = article.text
+
+    # get article title
     title = article.title
+
     return title, text
 
 
 def extract_text_from_file(file):
+    '''Extract text from uploaded file'''
+
+    # read text file
     if file.type == "text/plain":
         # To convert to a string based IO:
         stringio = StringIO(file.getvalue().decode("utf-8"))
+
+        # To read file as string:
         file_text = stringio.read()
+
         return file_text, None
+
+    # read pdf file
     elif file.type == "application/pdf":
         pdfReader = PdfFileReader(file)
         count = pdfReader.numPages
         all_text = ""
         pdf_title = pdfReader.getDocumentInfo().title
+
         for i in range(count):
+
             try:
                 page = pdfReader.getPage(i)
                 all_text += page.extractText()
+
             except:
                 continue
+
         file_text = all_text
+
         return file_text, pdf_title
-        # read docx file
+
+    # read docx file
     elif (
-        file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"):
+            file.type
+            == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    ):
         file_text = docx2txt.process(file)
+
         return file_text, None
 
 
@@ -60,13 +82,18 @@ def preprocess_plain_text(text, window_size=3):
     text = re.sub(r"@\S+", " ", text)  # mentions
     text = re.sub(r"#\S+", " ", text)  # hastags
     text = re.sub(r"\s{2,}", " ", text)  # over spaces
-    text = re.sub("[^.,!?%$A-Za-z0-9]+", " ", text)  # special characters except .,!?
-    # removing spaces
+    # text = re.sub("[^.,!?%$A-Za-z0-9]+", " ", text)  # special characters except .,!?
+
+    # break into lines and remove leading and trailing space on each
     lines = [line.strip() for line in text.splitlines()]
-    # break multi-headlines into a line each
+
+    # #break multi-headlines into a line each
     chunks = [phrase.strip() for line in lines for phrase in line.split("  ")]
-    # drop blank lines
+
+    # # drop blank lines
     text = '\n'.join(chunk for chunk in chunks if chunk)
+
+    ## We split this article into paragraphs and then every paragraph into sentences
     paragraphs = []
     for paragraph in text.replace('\n', ' ').split("\n\n"):
         if len(paragraph.strip()) > 0:
@@ -79,13 +106,21 @@ def preprocess_plain_text(text, window_size=3):
             end_idx = min(start_idx + window_size, len(paragraph))
             passages.append(" ".join(paragraph[start_idx:end_idx]))
 
+    st.write(f"Sentences: {sum([len(p) for p in paragraphs])}")
+    st.write(f"Passages: {len(passages)}")
+
     return passages
 
 
-def biencode(bi_enc, passages):
+@st.experimental_memo(suppress_st_warning=True)
+def bi_encode(bi_enc, passages):
     global bi_encoder
+    # We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
     bi_encoder = SentenceTransformer(bi_enc)
 
+    # quantize the model
+    # bi_encoder = quantize_dynamic(model, {Linear, Embedding})
+
     # Compute the embeddings using the multi-process pool
     with st.spinner('Encoding passages into a vector space...'):
         corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True, show_progress_bar=True)
@@ -95,13 +130,17 @@ def biencode(bi_enc, passages):
     return bi_encoder, corpus_embeddings
 
 
+@st.experimental_singleton(suppress_st_warning=True)
 def cross_encode():
     global cross_encoder
-    # cross-encoder to re-rank the results list to improve the quality
+    # We use a cross-encoder, to re-rank the results list to improve the quality
     cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
     return cross_encoder
 
 
+bi_enc_options = ["multi-qa-mpnet-base-dot-v1", "all-mpnet-base-v2", "multi-qa-MiniLM-L6-cos-v1"]
+
+
 def display_as_table(model, top_k, score='score'):
     # Display the df with text and scores as a table
     df = pd.DataFrame([(hit[score], passages[hit['corpus_id']]) for hit in model[0:top_k]], columns=['Score', 'Text'])
@@ -110,77 +149,117 @@ def display_as_table(model, top_k, score='score'):
     return df
 
 
+# Streamlit App
+
+st.title("Semantic Search with Retrieve & Rerank 📝")
+
+window_size = st.sidebar.slider("Paragraph Window Size", min_value=1, max_value=10, value=3, key=
+'slider')
+
+bi_encoder_type = st.sidebar.selectbox("Bi-Encoder", options=bi_enc_options, key='sbox')
+
+top_k = st.sidebar.slider("Number of Top Hits Generated", min_value=1, max_value=5, value=2)
+
+
+# This function will search all wikipedia articles for passages that
+# answer the query
 def search_func(query, top_k=top_k):
     global bi_encoder, cross_encoder
+
     st.subheader(f"Search Query: {query}")
+
     if url_text:
+
         st.write(f"Document Header: {title}")
+
     elif pdf_title:
+
         st.write(f"Document Header: {pdf_title}")
+
+    # Encode the query using the bi-encoder and find potentially relevant passages
     question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
     question_embedding = question_embedding.cpu()
     hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k, score_function=util.dot_score)
-    hits = hits[0]
-    # score all retrieved passages with the cross_encoder
+    hits = hits[0]  # Get the hits for the first query
+
+    # Now, score all retrieved passages with the cross_encoder
     cross_inp = [[query, passages[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]
-    # Output of top-3 hits
+
+    # Output of top-3 hits from bi-encoder
+    st.markdown("\n-------------------------\n")
+    st.subheader(f"Top-{top_k} Bi-Encoder Retrieval hits")
+    hits = sorted(hits, key=lambda x: x['score'], reverse=True)
+
+    cross_df = display_as_table(hits, top_k)
+    st.write(cross_df.to_html(index=False), unsafe_allow_html=True)
+
+    # Output of top-3 hits from re-ranker
     st.markdown("\n-------------------------\n")
     st.subheader(f"Top-{top_k} Cross-Encoder Re-ranker hits")
     hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
 
-    rerank_df = display_df_as_table(hits, top_k, 'cross-score')
+    rerank_df = display_as_table(hits, top_k, 'cross-score')
     st.write(rerank_df.to_html(index=False), unsafe_allow_html=True)
 
-    def clear_text():
-        st.session_state["text_url"] = ""
-        st.session_state["text_input"] = ""
 
-    def clear_search_text():
-        st.session_state["text_input"] = ""
+def clear_text():
+    st.session_state["text_url"] = ""
+    st.session_state["text_input"] = ""
+
+
+def clear_search_text():
+    st.session_state["text_input"] = ""
+
+
+url_text = st.text_input("Please Enter a url here",
+                         value="https://www.rba.gov.au/monetary-policy/rba-board-minutes/2022/2022-05-03.html",
+                         key='text_url', on_change=clear_search_text)
+
+st.markdown(
+    "<h3 style='text-align: center; color: red;'>OR</h3>",
+    unsafe_allow_html=True,
+)
+
+upload_doc = st.file_uploader("Upload a .txt, .pdf, .docx file", key="upload")
 
-    url_text = st.text_input("Enter a url here", value="https://en.wikipedia.org/wiki/Virat_Kohli", key='text_url',
-                             on_change=clear_search_text)
-    st.markdown(
-        "<h3 style='text-align: center; color: red;'>OR</h3>",
-        unsafe_allow_html=True, )
-    upload_doc = st.file_uploader("Upload a .txt, .pdf, .docx file", key="upload")
+search_query = st.text_input("Please Enter your search query here",
+                             value="What are the expectations for inflation for Australia?", key="text_input")
 
-    search_query = st.text_input("Enter your search query here", value="How many Centuries Virat Kohli scored?",
-                                 key="text_input")
-    if validators.url(url_text):
-        # if input is URL
-        title, text = extract_text_from_url(url_text)
-        passages = preprocess_plain_text(text, window_size=window_size)
+if validators.url(url_text):
+    # if input is URL
+    title, text = extract_text_from_url(url_text)
+    passages = preprocess_plain_text(text, window_size=window_size)
 
-    elif upload_doc:
+elif upload_doc:
 
-        text, pdf_title = extract_text_from_file(upload_doc)
-        passages = preprocess_plain_text(text, window_size=window_size)
+    text, pdf_title = extract_text_from_file(upload_doc)
+    passages = preprocess_plain_text(text, window_size=window_size)
 
-    col1, col2 = st.columns(2)
+col1, col2 = st.columns(2)
 
-    with col1:
-        search = st.button("Search", key='search_but', help='Click to Search!!')
+with col1:
+    search = st.button("Search", key='search_but', help='Click to Search!!')
 
-    with col2:
-        clear = st.button("Clear Text Input", on_click=clear_text, key='clear',
-                          help='Click to clear the URL input and search query')
+with col2:
+    clear = st.button("Clear Text Input", on_click=clear_text, key='clear',
+                      help='Click to clear the URL input and search query')
 
-    if search:
-        if bi_encoder_type:
-            with st.spinner(
-                    text=f"Loading {bi_encoder_type} bi-encoder and embedding document into vector space. This might take a few seconds depending on the length of your document..."
-            ):
-                bi_encoder, corpus_embeddings = bi_encode(bi_encoder_type, passages)
-                cross_encoder = cross_encode()
-                bm25 = bm25_api(passages)
+if search:
+    if bi_encoder_type:
+        with st.spinner(
+                text=f"Loading {bi_encoder_type} bi-encoder and embedding document into vector space. This might take a few seconds depending on the length of your document..."
+        ):
+            bi_encoder, corpus_embeddings = bi_encode(bi_encoder_type, passages)
+            cross_encoder = cross_encode()
 
-            with st.spinner(
-                    text="Embedding completed, searching for relevant text for given query and hits..."):
-                search_func(search_query, top_k)
+        with st.spinner(
+                text="Embedding completed, searching for relevant text for given query and hits..."):
+            search_func(search_query, top_k)
 
-    st.markdown("""  """)
+st.markdown("""
+            """)