File size: 14,366 Bytes
0dd62db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
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(
    "<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("""
            """)