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 import numpy as np import pandas as pd from newspaper import Article import base64 import docx2txt from io import StringIO from PyPDF2 import PdfFileReader import validators import nltk import warnings import streamlit as st from PIL import Image nltk.download('punkt') from nltk import sent_tokenize warnings.filterwarnings("ignore") auth_token = os.environ.get("auth_token") 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("cp1252")) # 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 elif ( file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" ): file_text = docx2txt.process(file) return file_text, None def preprocess_plain_text(text,window_size=3): 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])) st.write(f"Sentences: {sum([len(p) for p in paragraphs])}") st.write(f"Passages: {len(passages)}") return 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,use_auth_token=auth_token) #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...'): if bi_enc == 'intfloat/e5-base-v2': corpus_embeddings = bi_encoder.encode(['passage: ' + sentence for sentence in passages], convert_to_tensor=True) else: corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True) st.success(f"Embeddings computed. Shape: {corpus_embeddings.shape}") return bi_encoder, corpus_embeddings @st.experimental_singleton(suppress_st_warning=True) def cross_encode(): global 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.experimental_memo(suppress_st_warning=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.experimental_singleton(suppress_st_warning=True) def bm25_api(passages): tokenized_corpus = [] for passage in 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",'intfloat/e5-base-v2',"neeva/query2query"] def display_df_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']) df['Score'] = round(df['Score'],2) return df #Streamlit App st.title("Semantic Search with Retrieve & Rerank 📝") """ [![](https://img.shields.io/twitter/follow/nickmuchi?label=@nickmuchi&style=social)](https://twitter.com/nickmuchi) """ 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, bi_encoder_type, 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}") ##### 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.subheader(f"Top-{top_k} lexical search (BM25) hits") bm25_df = display_df_as_table(bm25_hits,top_k) st.write(bm25_df.to_html(index=False), unsafe_allow_html=True) if bi_encoder_type == 'intfloat/e5-base-v2': query = 'query: ' + query ##### 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.cpu() hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k,score_function=util.dot_score) 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.subheader(f"Top-{top_k} Bi-Encoder Retrieval hits") hits = sorted(hits, key=lambda x: x['score'], reverse=True) cross_df = display_df_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') st.write(rerank_df.to_html(index=False), unsafe_allow_html=True) 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. - 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.""") st.markdown("""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), [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2) 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)""") st.markdown( """ Code and App Inspiration Source: [Sentence Transformers](https://www.sbert.net/examples/applications/retrieve_rerank/README.html)""") st.markdown( """ 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):""") st.markdown( """ - 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(Image.open('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. - Select the number of top hits to be generated. - Paste the URL with your corpus or upload your preferred document in txt, pdf or Word format. - Semantic Search away!! """ ) st.markdown("---") 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/2023/2023-05-02.html",key='text_url',on_change=clear_search_text) st.markdown( "

OR

", 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") 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: text, pdf_title = extract_text_from_file(upload_doc) passages = preprocess_plain_text(text,window_size=window_size) col1, col2 = st.columns(2) 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') 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) with st.spinner( text="Embedding completed, searching for relevant text for given query and hits..."): search_func(search_query,bi_encoder_type,top_k) st.markdown(""" """) st.markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=nickmuchi-semantic-search-with-retrieve-and-rerank)")