Update app.py

app.py

@@ -1,6 +1,7 @@
-import requests
+from flask import Flask, request, jsonify
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
-import os, re
+import os
+import re
 import torch
 from rank_bm25 import BM25Okapi
 from sklearn.feature_extraction import _stop_words
@@ -15,359 +16,41 @@ from PyPDF2 import PdfFileReader
 import validators
 import nltk
 import warnings
-import streamlit as st
-from PIL import Image
 
+app = Flask(__name__)
 
 nltk.download('punkt')
 
-from nltk import sent_tokenize
-
-warnings.filterwarnings("ignore")
+# ... [rest of your imports and functions]
 
 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,
-)
+# ... [rest of your functions]
 
-upload_doc = st.file_uploader("Upload a .txt, .pdf, .docx file",key="upload")
+@app.route('/semantic-search', methods=['POST'])
+def semantic_search():
+    try:
+        # Get data from the request
+        data = request.json
 
-search_query = st.text_input("Please Enter your search query here",value="What are the expectations for inflation for Australia?",key="text_input")
+        # Extract necessary data from the request
+        url_text = data.get('url_text', '')
+        upload_doc = data.get('upload_doc', '')
+        search_query = data.get('search_query', '')
 
-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)
+        # Perform semantic search
+        result = perform_semantic_search(url_text, upload_doc, search_query)
 
-col1, col2 = st.columns(2)
+        return jsonify(result)
 
-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')
+    except Exception as e:
+        return jsonify({'error': str(e)})
 
-if search:
-    if bi_encoder_type:
+def perform_semantic_search(url_text, upload_doc, search_query):
+    # Your existing logic for semantic search
+    # ...
 
-        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)
+    return {'result': 'Your result here'}
 
-st.markdown("""
-            """)
-                        
-st.markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=nickmuchi-semantic-search-with-retrieve-and-rerank)")
+if __name__ == '__main__':
+    app.run(debug=True)