ABCASDFG98765432 commited on
Commit
05278ce
1 Parent(s): 4814168

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +343 -26
app.py CHANGED
@@ -1,7 +1,6 @@
1
- from flask import Flask, request, jsonify
2
  from sentence_transformers import SentenceTransformer, CrossEncoder, util
3
- import os
4
- import re
5
  import torch
6
  from rank_bm25 import BM25Okapi
7
  from sklearn.feature_extraction import _stop_words
@@ -16,41 +15,359 @@ from PyPDF2 import PdfFileReader
16
  import validators
17
  import nltk
18
  import warnings
 
 
19
 
20
- app = Flask(__name__)
21
 
22
  nltk.download('punkt')
23
 
24
- # ... [rest of your imports and functions]
 
 
25
 
26
  auth_token = os.environ.get("auth_token")
27
 
28
- # ... [rest of your functions]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
30
- @app.route('/semantic-search', methods=['POST'])
31
- def semantic_search():
32
- try:
33
- # Get data from the request
34
- data = request.json
35
 
36
- # Extract necessary data from the request
37
- url_text = data.get('url_text', '')
38
- upload_doc = data.get('upload_doc', '')
39
- search_query = data.get('search_query', '')
40
 
41
- # Perform semantic search
42
- result = perform_semantic_search(url_text, upload_doc, search_query)
 
 
 
 
 
 
 
43
 
44
- return jsonify(result)
45
 
46
- except Exception as e:
47
- return jsonify({'error': str(e)})
 
 
 
48
 
49
- def perform_semantic_search(url_text, upload_doc, search_query):
50
- # Your existing logic for semantic search
51
- # ...
52
 
53
- return {'result': 'Your result here'}
 
 
 
 
 
 
 
 
 
 
54
 
55
- if __name__ == '__main__':
56
- app.run(debug=True)
 
 
 
1
+ import requests
2
  from sentence_transformers import SentenceTransformer, CrossEncoder, util
3
+ import os, re
 
4
  import torch
5
  from rank_bm25 import BM25Okapi
6
  from sklearn.feature_extraction import _stop_words
 
15
  import validators
16
  import nltk
17
  import warnings
18
+ import streamlit as st
19
+ from PIL import Image
20
 
 
21
 
22
  nltk.download('punkt')
23
 
24
+ from nltk import sent_tokenize
25
+
26
+ warnings.filterwarnings("ignore")
27
 
28
  auth_token = os.environ.get("auth_token")
29
 
30
+ def extract_text_from_url(url: str):
31
+
32
+ '''Extract text from url'''
33
+
34
+ article = Article(url)
35
+ article.download()
36
+ article.parse()
37
+
38
+ # get text
39
+ text = article.text
40
+
41
+ # get article title
42
+ title = article.title
43
+
44
+ return title, text
45
+
46
+ def extract_text_from_file(file):
47
+
48
+ '''Extract text from uploaded file'''
49
+
50
+ # read text file
51
+ if file.type == "text/plain":
52
+ # To convert to a string based IO:
53
+ stringio = StringIO(file.getvalue().decode("cp1252"))
54
+
55
+ # To read file as string:
56
+ file_text = stringio.read()
57
+
58
+ return file_text, None
59
+
60
+ # read pdf file
61
+ elif file.type == "application/pdf":
62
+ pdfReader = PdfFileReader(file)
63
+ count = pdfReader.numPages
64
+ all_text = ""
65
+ pdf_title = pdfReader.getDocumentInfo().title
66
+
67
+ for i in range(count):
68
+
69
+ try:
70
+ page = pdfReader.getPage(i)
71
+ all_text += page.extractText()
72
+
73
+ except:
74
+ continue
75
+
76
+ file_text = all_text
77
+
78
+ return file_text, pdf_title
79
+
80
+ # read docx file
81
+ elif (
82
+ file.type
83
+ == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
84
+ ):
85
+ file_text = docx2txt.process(file)
86
+
87
+ return file_text, None
88
+
89
+ def preprocess_plain_text(text,window_size=3):
90
+
91
+ text = text.encode("ascii", "ignore").decode() # unicode
92
+ text = re.sub(r"https*\S+", " ", text) # url
93
+ text = re.sub(r"@\S+", " ", text) # mentions
94
+ text = re.sub(r"#\S+", " ", text) # hastags
95
+ text = re.sub(r"\s{2,}", " ", text) # over spaces
96
+ #text = re.sub("[^.,!?%$A-Za-z0-9]+", " ", text) # special characters except .,!?
97
+
98
+ #break into lines and remove leading and trailing space on each
99
+ lines = [line.strip() for line in text.splitlines()]
100
+
101
+ # #break multi-headlines into a line each
102
+ chunks = [phrase.strip() for line in lines for phrase in line.split(" ")]
103
+
104
+ # # drop blank lines
105
+ text = '\n'.join(chunk for chunk in chunks if chunk)
106
+
107
+ ## We split this article into paragraphs and then every paragraph into sentences
108
+ paragraphs = []
109
+ for paragraph in text.replace('\n',' ').split("\n\n"):
110
+ if len(paragraph.strip()) > 0:
111
+ paragraphs.append(sent_tokenize(paragraph.strip()))
112
+
113
+ #We combine up to 3 sentences into a passage. You can choose smaller or larger values for window_size
114
+ #Smaller value: Context from other sentences might get lost
115
+ #Lager values: More context from the paragraph remains, but results are longer
116
+ window_size = window_size
117
+ passages = []
118
+ for paragraph in paragraphs:
119
+ for start_idx in range(0, len(paragraph), window_size):
120
+ end_idx = min(start_idx+window_size, len(paragraph))
121
+ passages.append(" ".join(paragraph[start_idx:end_idx]))
122
+
123
+ st.write(f"Sentences: {sum([len(p) for p in paragraphs])}")
124
+ st.write(f"Passages: {len(passages)}")
125
+
126
+ return passages
127
+
128
+ @st.experimental_memo(suppress_st_warning=True)
129
+ def bi_encode(bi_enc,passages):
130
+
131
+ global bi_encoder
132
+ #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
133
+ bi_encoder = SentenceTransformer(bi_enc,use_auth_token=auth_token)
134
+
135
+ #quantize the model
136
+ #bi_encoder = quantize_dynamic(model, {Linear, Embedding})
137
+
138
+ #Compute the embeddings using the multi-process pool
139
+ with st.spinner('Encoding passages into a vector space...'):
140
+
141
+ if bi_enc == 'intfloat/e5-base-v2':
142
+
143
+ corpus_embeddings = bi_encoder.encode(['passage: ' + sentence for sentence in passages], convert_to_tensor=True)
144
+
145
+ elif bi_enc == 'BAAI/bge-base-en-v1.5':
146
+
147
+ instruction = "Represent this sentence for searching relevant passages: "
148
+
149
+ corpus_embeddings = bi_encoder.encode([instruction + sentence for sentence in passages], normalize_embeddings=True)
150
+
151
+ else:
152
+
153
+ corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True)
154
+
155
+
156
+ st.success(f"Embeddings computed. Shape: {corpus_embeddings.shape}")
157
+
158
+ return bi_encoder, corpus_embeddings
159
+
160
+ @st.experimental_singleton(suppress_st_warning=True)
161
+ def cross_encode():
162
+
163
+ global cross_encoder
164
+ #The bi-encoder will retrieve 100 documents. We use a cross-encoder, to re-rank the results list to improve the quality
165
+ cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
166
+ return cross_encoder
167
+
168
+ @st.experimental_memo(suppress_st_warning=True)
169
+ def bm25_tokenizer(text):
170
+
171
+ # We also compare the results to lexical search (keyword search). Here, we use
172
+ # the BM25 algorithm which is implemented in the rank_bm25 package.
173
+ # We lower case our text and remove stop-words from indexing
174
+ tokenized_doc = []
175
+ for token in text.lower().split():
176
+ token = token.strip(string.punctuation)
177
+
178
+ if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:
179
+ tokenized_doc.append(token)
180
+ return tokenized_doc
181
+
182
+ @st.experimental_singleton(suppress_st_warning=True)
183
+ def bm25_api(passages):
184
+
185
+ tokenized_corpus = []
186
+
187
+ for passage in passages:
188
+ tokenized_corpus.append(bm25_tokenizer(passage))
189
+
190
+ bm25 = BM25Okapi(tokenized_corpus)
191
+
192
+ return bm25
193
+
194
+ 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']
195
+
196
+ def display_df_as_table(model,top_k,score='score'):
197
+ # Display the df with text and scores as a table
198
+ df = pd.DataFrame([(hit[score],passages[hit['corpus_id']]) for hit in model[0:top_k]],columns=['Score','Text'])
199
+ df['Score'] = round(df['Score'],2)
200
+
201
+ return df
202
+
203
+ #Streamlit App
204
+
205
+ st.title("Semantic Search with Retrieve & Rerank 📝")
206
+
207
+ """
208
+ [![](https://img.shields.io/twitter/follow/nickmuchi?label=@nickmuchi&style=social)](https://twitter.com/nickmuchi)
209
+ """
210
+
211
+ window_size = st.sidebar.slider("Paragraph Window Size",min_value=1,max_value=10,value=3,key=
212
+ 'slider')
213
+
214
+ bi_encoder_type = st.sidebar.selectbox("Bi-Encoder", options=bi_enc_options, key='sbox')
215
+
216
+ top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5,value=2)
217
+
218
+ # This function will search all wikipedia articles for passages that
219
+ # answer the query
220
+ def search_func(query, bi_encoder_type, top_k=top_k):
221
+
222
+ global bi_encoder, cross_encoder
223
+
224
+ st.subheader(f"Search Query: {query}")
225
+
226
+ if url_text:
227
+
228
+ st.write(f"Document Header: {title}")
229
+
230
+ elif pdf_title:
231
+
232
+ st.write(f"Document Header: {pdf_title}")
233
+
234
+ ##### BM25 search (lexical search) #####
235
+ bm25_scores = bm25.get_scores(bm25_tokenizer(query))
236
+ top_n = np.argpartition(bm25_scores, -5)[-5:]
237
+ bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
238
+ bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
239
+
240
+ st.subheader(f"Top-{top_k} lexical search (BM25) hits")
241
+
242
+ bm25_df = display_df_as_table(bm25_hits,top_k)
243
+ st.write(bm25_df.to_html(index=False), unsafe_allow_html=True)
244
+
245
+ if bi_encoder_type == 'intfloat/e5-base-v2':
246
+ query = 'query: ' + query
247
+ ##### Sematic Search #####
248
+ # Encode the query using the bi-encoder and find potentially relevant passages
249
+ question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
250
+ question_embedding = question_embedding.cpu()
251
+ hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k,score_function=util.dot_score)
252
+ hits = hits[0] # Get the hits for the first query
253
+
254
+ ##### Re-Ranking #####
255
+ # Now, score all retrieved passages with the cross_encoder
256
+ cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
257
+ cross_scores = cross_encoder.predict(cross_inp)
258
+
259
+ # Sort results by the cross-encoder scores
260
+ for idx in range(len(cross_scores)):
261
+ hits[idx]['cross-score'] = cross_scores[idx]
262
+
263
+ # Output of top-3 hits from bi-encoder
264
+ st.markdown("\n-------------------------\n")
265
+ st.subheader(f"Top-{top_k} Bi-Encoder Retrieval hits")
266
+ hits = sorted(hits, key=lambda x: x['score'], reverse=True)
267
+
268
+ cross_df = display_df_as_table(hits,top_k)
269
+ st.write(cross_df.to_html(index=False), unsafe_allow_html=True)
270
+
271
+ # Output of top-3 hits from re-ranker
272
+ st.markdown("\n-------------------------\n")
273
+ st.subheader(f"Top-{top_k} Cross-Encoder Re-ranker hits")
274
+ hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
275
+
276
+ rerank_df = display_df_as_table(hits,top_k,'cross-score')
277
+ st.write(rerank_df.to_html(index=False), unsafe_allow_html=True)
278
+
279
+ st.markdown(
280
+ """
281
+ - 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.
282
+ - 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.
283
+ - 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.""")
284
+
285
+ st.markdown("""There models available to choose from:""")
286
+
287
+ st.markdown(
288
+ """
289
+ Model Source:
290
+ - 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)
291
+ - Cross-Encoder - [cross-encoder/ms-marco-MiniLM-L-12-v2](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-12-v2)""")
292
+
293
+ st.markdown(
294
+ """
295
+ Code and App Inspiration Source: [Sentence Transformers](https://www.sbert.net/examples/applications/retrieve_rerank/README.html)""")
296
+
297
+ st.markdown(
298
+ """
299
+ 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):""")
300
+
301
+ st.markdown(
302
+ """
303
+ - 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.
304
+ - 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.""")
305
+
306
+ st.image(Image.open('encoder.png'), caption='Retrieval and Re-Rank')
307
+
308
+ st.markdown("""
309
+ In order to use the app:
310
+ - Select the preferred Sentence Transformer model (Bi-Encoder).
311
+ - 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.
312
+ - Select the number of top hits to be generated.
313
+ - Paste the URL with your corpus or upload your preferred document in txt, pdf or Word format.
314
+ - Semantic Search away!! """
315
+ )
316
+
317
+ st.markdown("---")
318
+
319
+ def clear_text():
320
+ st.session_state["text_url"] = ""
321
+ st.session_state["text_input"]= ""
322
+
323
+ def clear_search_text():
324
+ st.session_state["text_input"]= ""
325
+
326
+ 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)
327
+
328
+ st.markdown(
329
+ "<h3 style='text-align: center; color: red;'>OR</h3>",
330
+ unsafe_allow_html=True,
331
+ )
332
 
333
+ upload_doc = st.file_uploader("Upload a .txt, .pdf, .docx file",key="upload")
 
 
 
 
334
 
335
+ search_query = st.text_input("Please Enter your search query here",value="What are the expectations for inflation for Australia?",key="text_input")
 
 
 
336
 
337
+ if validators.url(url_text):
338
+ #if input is URL
339
+ title, text = extract_text_from_url(url_text)
340
+ passages = preprocess_plain_text(text,window_size=window_size)
341
+
342
+ elif upload_doc:
343
+
344
+ text, pdf_title = extract_text_from_file(upload_doc)
345
+ passages = preprocess_plain_text(text,window_size=window_size)
346
 
347
+ col1, col2 = st.columns(2)
348
 
349
+ with col1:
350
+ search = st.button("Search",key='search_but', help='Click to Search!!')
351
+
352
+ with col2:
353
+ clear = st.button("Clear Text Input", on_click=clear_text,key='clear',help='Click to clear the URL input and search query')
354
 
355
+ if search:
356
+ if bi_encoder_type:
 
357
 
358
+ with st.spinner(
359
+ 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..."
360
+ ):
361
+ bi_encoder, corpus_embeddings = bi_encode(bi_encoder_type,passages)
362
+ cross_encoder = cross_encode()
363
+ bm25 = bm25_api(passages)
364
+
365
+ with st.spinner(
366
+ text="Embedding completed, searching for relevant text for given query and hits..."):
367
+
368
+ search_func(search_query,bi_encoder_type,top_k)
369
 
370
+ st.markdown("""
371
+ """)
372
+
373
+ st.markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=nickmuchi-semantic-search-with-retrieve-and-rerank)")