nickmuchi commited on
Commit
1281bfd
1 Parent(s): 786d711

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +303 -0
app.py ADDED
@@ -0,0 +1,303 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
7
+ import string
8
+ from tqdm.autonotebook import tqdm
9
+ import numpy as np
10
+ from bs4 import BeautifulSoup
11
+ from nltk import sent_tokenize
12
+ import time
13
+ from newspaper import Article
14
+ import base64
15
+ import docx2txt
16
+ from io import StringIO
17
+ from PyPDF2 import PdfFileReader
18
+ import validators
19
+
20
+ nltk.download('punkt')
21
+
22
+ from nltk import sent_tokenize
23
+
24
+ warnings.filterwarnings("ignore")
25
+
26
+ def extract_text_from_url(url: str):
27
+
28
+ '''Extract text from url'''
29
+
30
+ article = Article(url)
31
+ article.download()
32
+ article.parse()
33
+
34
+ # get text
35
+ text = article.text
36
+
37
+ # get article title
38
+ title = article.title
39
+
40
+ return title, text
41
+
42
+ def extract_text_from_file(file):
43
+
44
+ '''Extract text from uploaded file'''
45
+
46
+ # read text file
47
+ if file.type == "text/plain":
48
+ # To convert to a string based IO:
49
+ stringio = StringIO(file.getvalue().decode("utf-8"))
50
+
51
+ # To read file as string:
52
+ file_text = stringio.read()
53
+
54
+ # read pdf file
55
+ elif file.type == "application/pdf":
56
+ pdfReader = PdfFileReader(file)
57
+ count = pdfReader.numPages
58
+ all_text = ""
59
+
60
+ for i in range(count):
61
+ page = pdfReader.getPage(i)
62
+ all_text += page.extractText()
63
+ file_text = all_text
64
+
65
+ # read docx file
66
+ elif (
67
+ file.type
68
+ == "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
69
+ ):
70
+ file_text = docx2txt.process(file)
71
+
72
+ return file_text
73
+
74
+ def preprocess_plain_text(text,window_size=window_size):
75
+
76
+ text = text.encode("ascii", "ignore").decode() # unicode
77
+ text = re.sub(r"https*\S+", " ", text) # url
78
+ text = re.sub(r"@\S+", " ", text) # mentions
79
+ text = re.sub(r"#\S+", " ", text) # hastags
80
+ #text = re.sub(r"\s{2,}", " ", text) # over spaces
81
+ text = re.sub("[^.,!?%$A-Za-z0-9]+", " ", text) # special characters except .,!?
82
+
83
+ #break into lines and remove leading and trailing space on each
84
+ lines = [line.strip() for line in text.splitlines()]
85
+
86
+ # #break multi-headlines into a line each
87
+ chunks = [phrase.strip() for line in lines for phrase in line.split(" ")]
88
+
89
+ # # drop blank lines
90
+ text = '\n'.join(chunk for chunk in chunks if chunk)
91
+
92
+ ## We split this article into paragraphs and then every paragraph into sentences
93
+ paragraphs = []
94
+ for paragraph in text.replace('\n',' ').split("\n\n"):
95
+ if len(paragraph.strip()) > 0:
96
+ paragraphs.append(sent_tokenize(paragraph.strip()))
97
+
98
+ #We combine up to 3 sentences into a passage. You can choose smaller or larger values for window_size
99
+ #Smaller value: Context from other sentences might get lost
100
+ #Lager values: More context from the paragraph remains, but results are longer
101
+ window_size = window_size
102
+ passages = []
103
+ for paragraph in paragraphs:
104
+ for start_idx in range(0, len(paragraph), window_size):
105
+ end_idx = min(start_idx+window_size, len(paragraph))
106
+ passages.append(" ".join(paragraph[start_idx:end_idx]))
107
+
108
+ print("Paragraphs: ", len(paragraphs))
109
+ print("Sentences: ", sum([len(p) for p in paragraphs]))
110
+ print("Passages: ", len(passages))
111
+
112
+ return passages
113
+
114
+ @st.cache(allow_output_mutation=True)
115
+ def bi_encoder(bi_enc,passages):
116
+
117
+ #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
118
+ bi_encoder = SentenceTransformer(bi_enc)
119
+ #Start the multi-process pool on all available CUDA devices
120
+ pool = bi_encoder.start_multi_process_pool()
121
+
122
+ #Compute the embeddings using the multi-process pool
123
+ print('encoding passages into a vector space...')
124
+ corpus_embeddings = bi_encoder.encode_multi_process(passages, pool)
125
+ print("Embeddings computed. Shape:", corpus_embeddings.shape)
126
+
127
+ #Optional: Stop the proccesses in the pool
128
+ bi_encoder.stop_multi_process_pool(pool)
129
+
130
+ return corpus_embeddings
131
+
132
+ @st.cache(allow_output_mutation=True)
133
+ def cross_encoder():
134
+
135
+ #The bi-encoder will retrieve 100 documents. We use a cross-encoder, to re-rank the results list to improve the quality
136
+ cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
137
+ return cross_encoder
138
+
139
+ @st.cache(allow_output_mutation=True)
140
+ def bm25_tokenizer(text):
141
+
142
+ # We also compare the results to lexical search (keyword search). Here, we use
143
+ # the BM25 algorithm which is implemented in the rank_bm25 package.
144
+ # We lower case our text and remove stop-words from indexing
145
+ tokenized_doc = []
146
+ for token in text.lower().split():
147
+ token = token.strip(string.punctuation)
148
+
149
+ if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:
150
+ tokenized_doc.append(token)
151
+ return tokenized_doc
152
+
153
+ @st.cache(allow_output_mutation=True)
154
+ def bm25_api(passages):
155
+
156
+ tokenized_corpus = []
157
+ print('implementing BM25 algo for lexical search..')
158
+ for passage in tqdm(passages):
159
+ tokenized_corpus.append(bm25_tokenizer(passage))
160
+
161
+ bm25 = BM25Okapi(tokenized_corpus)
162
+
163
+ return bm25
164
+
165
+ bi_enc_options = ["multi-qa-mpnet-base-dot-v1","all-mpnet-base-v2","multi-qa-MiniLM-L6-cos-v1"]
166
+
167
+ # This function will search all wikipedia articles for passages that
168
+ # answer the query
169
+ def search(query, top_k=2):
170
+ st.write(f"Search Query: {query}")
171
+
172
+ st.write("Document Header: ")
173
+
174
+ ##### BM25 search (lexical search) #####
175
+ bm25_scores = bm25.get_scores(bm25_tokenizer(query))
176
+ top_n = np.argpartition(bm25_scores, -5)[-5:]
177
+ bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
178
+ bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
179
+
180
+ st.write(f"Top-{top_k} lexical search (BM25) hits")
181
+ for hit in bm25_hits[0:top_k]:
182
+ st.write("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
183
+
184
+ ##### Sematic Search #####
185
+ # Encode the query using the bi-encoder and find potentially relevant passages
186
+ question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
187
+ question_embedding = question_embedding.gpu()
188
+ hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
189
+ hits = hits[0] # Get the hits for the first query
190
+
191
+ ##### Re-Ranking #####
192
+ # Now, score all retrieved passages with the cross_encoder
193
+ cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
194
+ cross_scores = cross_encoder.predict(cross_inp)
195
+
196
+ # Sort results by the cross-encoder scores
197
+ for idx in range(len(cross_scores)):
198
+ hits[idx]['cross-score'] = cross_scores[idx]
199
+
200
+ # Output of top-3 hits from bi-encoder
201
+ st.markdown("\n-------------------------\n")
202
+ st.write(f"Top-{top_k} Bi-Encoder Retrieval hits")
203
+ hits = sorted(hits, key=lambda x: x['score'], reverse=True)
204
+ for hit in hits[0:top_k]:
205
+ st.write("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
206
+
207
+ # Output of top-3 hits from re-ranker
208
+ st.markdown("\n-------------------------\n")
209
+ st.write(f"Top-{top_k} Cross-Encoder Re-ranker hits")
210
+ hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
211
+ for hit in hits[0:top_k]:
212
+ st.write("\t{:.3f}\t{}".format(hit['cross-score'], passages[hit['corpus_id']].replace("\n", " ")))
213
+
214
+ #Streamlit App
215
+
216
+ st.title("Semantic Search with Retrieve & Rerank 📝")
217
+
218
+ window_size = st.sidebar.slider("Paragraph Window Size",min_value=1,max_value=10,value=3)
219
+
220
+ bi_encoder_type = st.sidebar.selectbox(
221
+ "Bi-Encoder", options=bi_enc_options
222
+ )
223
+
224
+ top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5,value=2)
225
+
226
+ st.markdown(
227
+ """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.
228
+ 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.
229
+ 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.
230
+
231
+ There models available to choose from:""")
232
+
233
+ st.markdown(
234
+ """Model Source:
235
+ 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)
236
+ Cross-Encoder - [cross-encoder/ms-marco-MiniLM-L-12-v2](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-12-v2)
237
+
238
+ Code and App Inspiration Source:
239
+ [Sentence Transformers](https://www.sbert.net/examples/applications/retrieve_rerank/README.html)
240
+
241
+ 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):
242
+
243
+ 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.
244
+
245
+ 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. """
246
+ )
247
+
248
+ st.image('encoder.png', caption='Retrieval and Re-Rank')
249
+
250
+ st.markdown("""
251
+ In order to use the app:
252
+ - Select the preferred Sentence Transformer model (Bi-Encoder).
253
+ - 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.
254
+ - Paste the URL with your corpus or upload your preferred document in txt, pdf or Word format
255
+ - Semantic Search away!! """
256
+ )
257
+
258
+ st.markdown("---")
259
+
260
+ url_text = st.text_input("Please Enter a url here")
261
+
262
+ st.markdown(
263
+ "<h3 style='text-align: center; color: red;'>OR</h3>",
264
+ unsafe_allow_html=True,
265
+ )
266
+
267
+ st.markdown(
268
+ "<h3 style='text-align: center; color: red;'>OR</h3>",
269
+ unsafe_allow_html=True,
270
+ )
271
+
272
+ upload_doc = st.file_uploader(
273
+ "Upload a .txt, .pdf, .docx file"
274
+ )
275
+
276
+ search_query = st.text_input("Please Enter your search query here")
277
+
278
+ if validators.url(url_text):
279
+ #if input is URL
280
+ title, text = extract_text_from_url(url_text)
281
+ passages = preprocess_plain_text(text,window_size=window_size)
282
+
283
+ elif upload_doc:
284
+
285
+ passages = preprocess_plain_text(extract_text_from_file(upload_doc),window_size=window_size)
286
+
287
+ search = st.button("Search")
288
+
289
+ if search:
290
+ if bi_encoder_type:
291
+
292
+ with st.spinner(
293
+ 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..."
294
+ ):
295
+ corpus_embeddings = bi_encoder(bi_encoder_type,passages)
296
+ cross_encoder = cross_encoder()
297
+ bm25 = bm25_api(passages)
298
+
299
+ with st.spinner(
300
+ text="Embedding completed, searching for relevant text for given query and hits..."):
301
+
302
+ search(search_query,top_k)
303
+