Spaces:

nickmuchi
/

semantic-search-with-retrieve-and-rerank

Running

App Files Files Community

semantic-search-with-retrieve-and-rerank / app.py

nickmuchi

Update app.py

87e6a8a 8 months ago

raw history blame contribute delete

No virus

14.9 kB

	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)

	elif bi_enc == 'BAAI/bge-base-en-v1.5':

	instruction = "Represent this sentence for searching relevant passages: "

	corpus_embeddings = bi_encoder.encode([instruction + sentence for sentence in passages], normalize_embeddings=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 = ["BAAI/bge-base-en-v1.5","multi-qa-mpnet-base-dot-v1","all-mpnet-base-v2","multi-qa-MiniLM-L6-cos-v1",'intfloat/e5-base-v2']

	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 - [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5), [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(
	"<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")

	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)")