Create app.py

app.py

@@ -0,0 +1,303 @@
+import requests
+from sentence_transformers import SentenceTransformer, CrossEncoder, util
+import os, re
+import torch
+from rank_bm25 import BM25Okapi
+from sklearn.feature_extraction import _stop_words
+import string
+from tqdm.autonotebook import tqdm
+import numpy as np
+from bs4 import BeautifulSoup
+from nltk import sent_tokenize
+import time
+from newspaper import Article
+import base64 
+import docx2txt
+from io import StringIO
+from PyPDF2 import PdfFileReader
+import validators
+
+nltk.download('punkt')
+
+from nltk import sent_tokenize
+
+warnings.filterwarnings("ignore")
+
+def extract_text_from_url(url: str):
+    
+    '''Extract text from url'''
+    
+    article = Article(url)
+    article.download()
+    article.parse()
+    
+    # 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()
+
+    # read pdf file
+    elif file.type == "application/pdf":
+        pdfReader = PdfFileReader(file)
+        count = pdfReader.numPages
+        all_text = ""
+        
+        for i in range(count):
+            page = pdfReader.getPage(i)
+            all_text += page.extractText()
+            file_text = all_text
+
+    # read docx file
+    elif (
+        file.type
+        == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    ):
+        file_text = docx2txt.process(file)
+
+    return file_text
+    
+def preprocess_plain_text(text,window_size=window_size):
+
+    text = text.encode("ascii", "ignore").decode()  # unicode
+    text = re.sub(r"https*\S+", " ", text)  # url
+    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 .,!?
+    
+    #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
+    chunks = [phrase.strip() for line in lines for phrase in line.split("  ")]
+    
+    # # 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:
+            paragraphs.append(sent_tokenize(paragraph.strip()))
+
+    #We combine up to 3 sentences into a passage. You can choose smaller or larger values for window_size
+    #Smaller value: Context from other sentences might get lost
+    #Lager values: More context from the paragraph remains, but results are longer
+    window_size = window_size
+    passages = []
+    for paragraph in paragraphs:
+        for start_idx in range(0, len(paragraph), window_size):
+            end_idx = min(start_idx+window_size, len(paragraph))
+            passages.append(" ".join(paragraph[start_idx:end_idx]))
+            
+    print("Paragraphs: ", len(paragraphs))
+    print("Sentences: ", sum([len(p) for p in paragraphs]))
+    print("Passages: ", len(passages))
+
+    return passages
+
+@st.cache(allow_output_mutation=True)
+def bi_encoder(bi_enc,passages):
+    
+    #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
+    bi_encoder = SentenceTransformer(bi_enc)
+    #Start the multi-process pool on all available CUDA devices
+    pool = bi_encoder.start_multi_process_pool()
+
+    #Compute the embeddings using the multi-process pool
+    print('encoding passages into a vector space...')
+    corpus_embeddings = bi_encoder.encode_multi_process(passages, pool)
+    print("Embeddings computed. Shape:", corpus_embeddings.shape)
+
+    #Optional: Stop the proccesses in the pool
+    bi_encoder.stop_multi_process_pool(pool)
+    
+    return corpus_embeddings
+    
+@st.cache(allow_output_mutation=True)
+def cross_encoder():
+    
+    #The bi-encoder will retrieve 100 documents. 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
+    
+@st.cache(allow_output_mutation=True)
+def bm25_tokenizer(text):
+    
+# We also compare the results to lexical search (keyword search). Here, we use 
+# the BM25 algorithm which is implemented in the rank_bm25 package.
+# We lower case our text and remove stop-words from indexing
+    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
+
+@st.cache(allow_output_mutation=True)
+def bm25_api(passages):
+
+    tokenized_corpus = []
+    print('implementing BM25 algo for lexical search..')
+    for passage in tqdm(passages):
+        tokenized_corpus.append(bm25_tokenizer(passage))
+
+    bm25 = BM25Okapi(tokenized_corpus)
+    
+    return bm25
+
+bi_enc_options = ["multi-qa-mpnet-base-dot-v1","all-mpnet-base-v2","multi-qa-MiniLM-L6-cos-v1"]
+
+# This function will search all wikipedia articles for passages that
+# answer the query
+def search(query, top_k=2):
+    st.write(f"Search Query: {query}")
+    
+    st.write("Document Header: ")
+
+    ##### BM25 search (lexical search) #####
+    bm25_scores = bm25.get_scores(bm25_tokenizer(query))
+    top_n = np.argpartition(bm25_scores, -5)[-5:]
+    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)
+    
+    st.write(f"Top-{top_k} lexical search (BM25) hits")
+    for hit in bm25_hits[0:top_k]:
+        st.write("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
+
+    ##### Sematic Search #####
+    # 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.gpu()
+    hits = util.semantic_search(question_embedding, corpus_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, 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 from bi-encoder
+    st.markdown("\n-------------------------\n")
+    st.write(f"Top-{top_k} Bi-Encoder Retrieval hits")
+    hits = sorted(hits, key=lambda x: x['score'], reverse=True)
+    for hit in hits[0:top_k]:
+        st.write("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
+
+    # Output of top-3 hits from re-ranker
+    st.markdown("\n-------------------------\n")
+    st.write(f"Top-{top_k} Cross-Encoder Re-ranker hits")
+    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
+    for hit in hits[0:top_k]:
+        st.write("\t{:.3f}\t{}".format(hit['cross-score'], passages[hit['corpus_id']].replace("\n", " ")))
+        
+#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)
+
+bi_encoder_type = st.sidebar.selectbox(
+    "Bi-Encoder", options=bi_enc_options
+)
+
+top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5,value=2)
+
+st.markdown(
+    """The app supports asymmetric Semantic search which seeks to improve search accuracy of documents/URL by understanding the content of the search query in contrast to traditional search engines which only find documents based on lexical matches, semantic search can also find synonyms. 
+    The idea behind semantic search is to embed all entries in your corpus, whether they be sentences, paragraphs, or documents, into a vector space. At search time, the query is embedded into the same vector space and the closest embeddings from your corpus are found. These entries should have a high semantic overlap with the query.
+    The all-* models where trained on all available training data (more than 1 billion training pairs) and are designed as general purpose models. The all-mpnet-base-v2 model provides the best quality, while all-MiniLM-L6-v2 is 5 times faster and still offers good quality. The models used have been trained on broad datasets, however, if your document/corpus is specialised, such as for science or economics, the results returned might be unsatisfactory.
+     
+    There models available to choose from:""")
+    
+st.markdown(
+    """Model Source: 
+    Bi-Encoders - [multi-qa-mpnet-base-dot-v1](https://huggingface.co/sentence-transformers/multi-qa-mpnet-base-dot-v1), [all-mpnet-base-v2](https://huggingface.co/sentence-transformers/all-mpnet-base-v2), [multi-qa-MiniLM-L6-cos-v1](https://huggingface.co/sentence-transformers/multi-qa-MiniLM-L6-cos-v1) and [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
+    Cross-Encoder - [cross-encoder/ms-marco-MiniLM-L-12-v2](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-12-v2)
+    
+    Code and App Inspiration Source:
+    [Sentence Transformers](https://www.sbert.net/examples/applications/retrieve_rerank/README.html)
+    
+    Quick summary of the purposes of a Bi and Cross-encoder below, the image and info were adapted from [www.sbert.net](https://www.sbert.net/examples/applications/semantic-search/README.html):
+    
+    Bi-Encoder (Retrieval): The Bi-encoder is responsible for independently embedding the sentences and search queries into a vector space. The result is then passed to the cross-encoder for checking the relevance/similarity between the query and sentences.
+    
+    Cross-Encoder (Re-Ranker): A re-ranker based on a Cross-Encoder can substantially improve the final results for the user. The query and a possible document is passed simultaneously to transformer network, which then outputs a single score between 0 and 1 indicating how relevant the document is for the given query. The cross-encoder further boost the performance, especially when you search over a corpus for which the bi-encoder was not trained for.   """
+)
+
+st.image('encoder.png', caption='Retrieval and Re-Rank')
+
+st.markdown("""   
+    In order to use the app:
+    - Select the preferred Sentence Transformer model (Bi-Encoder).
+    - Select the number of sentences per paragraph to partition your corpus (Window-Size), if you choose a small value the context from the other sentences might get lost and for larger values the results might take longer to generate.
+    - Paste the URL with your corpus or upload your preferred document in txt, pdf or Word format
+    - Semantic Search away!! """
+)
+
+st.markdown("---")
+
+url_text = st.text_input("Please Enter a url here")
+
+st.markdown(
+    "<h3 style='text-align: center; color: red;'>OR</h3>",
+    unsafe_allow_html=True,
+)
+
+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"
+)
+
+search_query = st.text_input("Please Enter your search query here")
+
+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:
+    
+    passages = preprocess_plain_text(extract_text_from_file(upload_doc),window_size=window_size)
+
+search = st.button("Search")
+
+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..."
+        ):
+            corpus_embeddings = bi_encoder(bi_encoder_type,passages)
+            cross_encoder = cross_encoder()
+            bm25 = bm25_api(passages)
+            
+         with st.spinner(
+            text="Embedding completed, searching for relevant text for given query and hits..."):
+            
+            search(search_query,top_k)
+