Spaces:
Sleeping
Sleeping
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 | |
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 | |
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 | |
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 | |
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) | |