Upload app.py, version from Jiaqi

app.py

@@ -1,16 +1,11 @@
 import streamlit as st
-from streamlit_tags import st_tags, st_tags_sidebar
-from keytotext import pipeline
+
 from PIL import Image
 
 import json
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
-import gzip
-import os
-import torch
 import pickle
-import random
-import numpy as np
+import pandas as pd
 
 ############
 ## Main page
@@ -40,7 +35,7 @@ option1 = st.sidebar.selectbox(
      ('multi-qa-MiniLM-L6-cos-v1','null','null'))
 
 option2 = st.sidebar.selectbox(
-     'Which corss-encoder model would you like to be selected?',
+     'Which cross-encoder model would you like to be selected?',
      ('cross-encoder/ms-marco-MiniLM-L-6-v2','null','null'))
 
 st.sidebar.success("Load Successfully!")
@@ -49,22 +44,28 @@ st.sidebar.success("Load Successfully!")
 #    print("Warning: No GPU found. Please add GPU to your notebook")
 
 #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
-bi_encoder = SentenceTransformer(option1,device='cpu')
+@st.cache_resource
+def load_encoders(sentence_enc, cross_enc):
+    return SentenceTransformer(sentence_enc,device='cpu'), CrossEncoder(cross_enc,device='cpu')
+bi_encoder, cross_encoder = load_encoders(option1,option2)
 bi_encoder.max_seq_length = 256    #Truncate long passages to 256 tokens
 top_k = 32                          #Number of passages we want to retrieve with the bi-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(option2, device='cpu')
-
 passages = []
 
 # load pre-train embeedings files
+@st.cache_resource
+def load_pickle(path):
+    with open(path, "rb") as fIn:
+        cache_data = pickle.load(fIn)
+        passages = cache_data['sentences']
+        corpus_embeddings = cache_data['embeddings']
+    print("Load pre-computed embeddings from disc")
+    return passages,corpus_embeddings
+
 embedding_cache_path = 'etsy-embeddings-cpu.pkl'
-print("Load pre-computed embeddings from disc")
-with open(embedding_cache_path, "rb") as fIn:
-  cache_data = pickle.load(fIn)
-  passages = cache_data['sentences']
-  corpus_embeddings = cache_data['embeddings']
+passages,corpus_embeddings = load_pickle(embedding_cache_path)
+
 
 from rank_bm25 import BM25Okapi
 from sklearn.feature_extraction import _stop_words
@@ -75,15 +76,24 @@ import re
 
 import yake
 
-language = "en"
-max_ngram_size = 3
-deduplication_threshold = 0.9
-deduplication_algo = 'seqm'
-windowSize = 3
-numOfKeywords = 3
-
-custom_kw_extractor = yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold, dedupFunc=deduplication_algo, windowsSize=windowSize, top=numOfKeywords, features=None)
-
+@st.cache_resource
+def load_model():
+    language = "en"
+    max_ngram_size = 3
+    deduplication_threshold = 0.9
+    deduplication_algo = 'seqm'
+    windowSize = 3
+    numOfKeywords = 3
+    return yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold, dedupFunc=deduplication_algo, windowsSize=windowSize, top=numOfKeywords, features=None)
+custom_kw_extractor = load_model()
+# load query GMS information
+@st.cache_resource
+def load_json(path):
+    with open(path, 'r') as file:
+        query_gms_dict = json.load(file)
+    return query_gms_dict
+
+query_gms_dict = load_json('query_gms.json')
 # We lower case our text and remove stop-words from indexing
 def bm25_tokenizer(text):
     tokenized_doc = []
@@ -94,10 +104,14 @@ def bm25_tokenizer(text):
             tokenized_doc.append(token)
     return tokenized_doc
 
-tokenized_corpus = []
-for passage in tqdm(passages):
-    tokenized_corpus.append(bm25_tokenizer(passage))
+@st.cache_resource
+def get_tokenized_corpus(passages,_tokenizer):
+    tokenized_corpus = []
+    for passage in passages:
+        tokenized_corpus.append(_tokenizer(passage))
+    return tokenized_corpus
 
+tokenized_corpus = get_tokenized_corpus(passages,bm25_tokenizer)
 bm25 = BM25Okapi(tokenized_corpus)
 
 def word_len(s):
@@ -106,205 +120,155 @@ def word_len(s):
 
 # This function will search all wikipedia articles for passages that
 # answer the query
-def search(query):
+DEFAULT_SCORE = -100.0
+def clean_string(input_string):
+    string_sub1 = re.sub("([^\u0030-\u0039\u0041-\u007a])", ' ', input_string)
+    string_sub2 = re.sub("\x20\x20", "\n", string_sub1)
+    string_strip = string_sub2.strip().lower()
+    output_string = []
+    if len(string_strip) > 20:
+        keywords = custom_kw_extractor.extract_keywords(string_strip)
+        for tokens in keywords:
+            string_clean = tokens[0]
+            if word_len(string_clean) > 1:
+                output_string.append(string_clean)
+    else:
+        output_string.append(string_strip)
+    return output_string
+
+# def add_gms_score_for_candidates(candidates, query_gms_dict):
+#     for query_candidate in candidates:
+#         value = candidates[query_candidate]
+#         value['gms'] = query_gms_dict.get(query_candidate, 0)
+#         candidates[query_candidate] = value
+#     return candidates
+
+def generate_query_expansion_candidates(query):
     print("Input query:", query)
-    total_qe = []
+    expanded_query_set = {}
 
     ##### 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)
-    
-    #print("Top-10 lexical search (BM25) hits")
-    qe_string = []
-    for hit in bm25_hits[0:1000]:
-      if passages[hit['corpus_id']].replace("\n", " ") not in qe_string:
-        qe_string.append(passages[hit['corpus_id']].replace("\n", ""))
-
-    sub_string = []
-    for item in qe_string:
-      for sub_item in item.split(","):
-        sub_string.append(sub_item)
-    #print(sub_string)
-    total_qe.append(sub_string)
-
-    ##### Sematic Search #####
-    # Encode the query using the bi-encoder and find potentially relevant passages
-    query_embedding = bi_encoder.encode(query, convert_to_tensor=True)
-    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)
-    hits = hits[0]  # Get the hits for the first query
+    # finds the indices of the top n scores
+    top_n_indices = np.argpartition(bm25_scores, -5)[-5:]
+    bm25_hits = [{'corpus_id': idx, 'bm25_score': bm25_scores[idx]} for idx in top_n_indices]
+    # bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
 
-    ##### 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-10 hits from bi-encoder
-    #print("\n-------------------------\n")
-    #print("Top-N Bi-Encoder Retrieval hits")
-    hits = sorted(hits, key=lambda x: x['score'], reverse=True)
-    qe_string = []
-    for hit in hits[0:1000]:
-      if passages[hit['corpus_id']].replace("\n", " ") not in qe_string:
-        qe_string.append(passages[hit['corpus_id']].replace("\n", ""))
-    #print(qe_string)
-    total_qe.append(qe_string)
-
-    # Output of top-10 hits from re-ranker
-    #print("\n-------------------------\n")
-    #print("Top-N Cross-Encoder Re-ranker hits")
-    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
-    qe_string = []
-    for hit in hits[0:1000]:
-      if passages[hit['corpus_id']].replace("\n", " ") not in qe_string:
-        qe_string.append(passages[hit['corpus_id']].replace("\n", ""))
-    #print(qe_string)
-    total_qe.append(qe_string)
-
-    # Total Results
-    total_qe.append(qe_string)
-    st.write("E-Commerce Query Expansion Results: \n")
-
-    res = []
-    for sub_list in total_qe:
-      for i in sub_list:
-        rs = re.sub("([^\u0030-\u0039\u0041-\u007a])", ' ', i)
-        rs_final = re.sub("\x20\x20", "\n", rs)
-        #st.write(rs_final.strip())
-        res.append(rs_final.strip())
-
-    res_clean = []
-    for out in res:
-      if len(out) > 20:
-        keywords = custom_kw_extractor.extract_keywords(out)
-        for key in keywords:
-          res_clean.append(key[0])
-      else:
-        res_clean.append(out)
-
-    show_out = []
-    for i in res_clean:
-      num = word_len(i)
-      if num > 1:
-        show_out.append(i)
-    unique_list = list(set(show_out))
-    new_unique_list = [item for item in unique_list if item != query]
-    Lowercasing_list = [item.lower() for item in new_unique_list]
-    st.write(Lowercasing_list[0:maxtags_sidebar])
-
-    return Lowercasing_list
-
-def search_nolog(query):
-    total_qe = []
-    ##### 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)
-    
-    qe_string = []
-    for hit in bm25_hits[0:1000]:
-      if passages[hit['corpus_id']].replace("\n", " ") not in qe_string:
-        qe_string.append(passages[hit['corpus_id']].replace("\n", ""))
-
-    sub_string = []
-    for item in qe_string:
-      for sub_item in item.split(","):
-        sub_string.append(sub_item)
-    total_qe.append(sub_string)
 
     ##### Sematic Search #####
     # Encode the query using the bi-encoder and find potentially relevant passages
     query_embedding = bi_encoder.encode(query, convert_to_tensor=True)
-    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)
-    hits = hits[0]  # Get the hits for the first query
+    # query_embedding = query_embedding.cuda()
+    # Get the hits for the first query
+    encoder_hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)[0]
 
-    ##### Re-Ranking #####
-    # Now, score all retrieved passages with the cross_encoder
-    cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
+    # For all retrieved passages, add the cross_encoder scores
+    cross_inp = [[query, passages[hit['corpus_id']]] for hit in encoder_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-10 hits from bi-encoder
-    hits = sorted(hits, key=lambda x: x['score'], reverse=True)
-    qe_string = []
-    for hit in hits[0:1000]:
-      if passages[hit['corpus_id']].replace("\n", " ") not in qe_string:
-        qe_string.append(passages[hit['corpus_id']].replace("\n", ""))
-    total_qe.append(qe_string)
-
-    # Output of top-10 hits from re-ranker
-    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
-    qe_string = []
-    for hit in hits[0:1000]:
-      if passages[hit['corpus_id']].replace("\n", " ") not in qe_string:
-        qe_string.append(passages[hit['corpus_id']].replace("\n", ""))
-    total_qe.append(qe_string)
-
+        encoder_hits[idx]['cross_score'] = cross_scores[idx]
+
+    candidates = {}
+    for hit in bm25_hits:
+        corpus_id = hit['corpus_id']
+        if  corpus_id not in candidates:
+            candidates[corpus_id] = {'bm25_score': hit['bm25_score'], 'bi_score': DEFAULT_SCORE, 'cross_score': DEFAULT_SCORE}
+    for hit in encoder_hits:
+        corpus_id = hit['corpus_id']
+        if corpus_id not in candidates:
+            candidates[corpus_id] = {'bm25_score': DEFAULT_SCORE, 'bi_score': hit['score'], 'cross_score': hit['cross_score']}
+        else:
+            bm25_score = candidates[corpus_id]['bm25_score']
+            candidates[corpus_id].update({'bm25_score': bm25_score, 'bi_score': hit['score'], 'cross_score': hit['cross_score']})
+
+    final_candidates = {}
+    for key, value in candidates.items():
+        input_string = passages[key].replace("\n", "")
+        string_set = set(clean_string(input_string))
+        for item in string_set:
+            final_candidates[item.replace("\n", " ")] = value
+    # remove the query itself from candidates
+    if query in final_candidates:
+        del final_candidates[query]
+    # print(final_candidates)
+    # add gms column
+    df = pd.DataFrame(final_candidates).T
+    df['gms'] = [query_gms_dict.get(i,0) for i in df.index]
     # Total Results
-    total_qe.append(qe_string)
-
-    res = []
-    for sub_list in total_qe:
-      for i in sub_list:
-        rs = re.sub("([^\u0030-\u0039\u0041-\u007a])", ' ', i)
-        rs_final = re.sub("\x20\x20", "\n", rs)
-        res.append(rs_final.strip())
-
-    res_clean = []
-    for out in res:
-      if len(out) > 20:
-        keywords = custom_kw_extractor.extract_keywords(out)
-        for key in keywords:
-          res_clean.append(key[0])
-      else:
-        res_clean.append(out)
-
-    show_out = []
-    for i in res_clean:
-      num = word_len(i)
-      if num > 1:
-        show_out.append(i)
-
-    return show_out
-
-def reranking():
-    rerank_list = []
-    reres = []
-    remove_dup = []
-    rerank_list = search_nolog(query = user_query)
-    unique_list = list(set(rerank_list))
-    Lowercasing_list = [item.lower() for item in unique_list]
-    new_unique_list = [item for item in Lowercasing_list if item != user_query]
-
-    for i in new_unique_list:
-        clean_string = i.strip()
-        if clean_string not in remove_dup:
-            remove_dup.append(clean_string)
-
-    st.write("E-Commerce Query Expansion Results: \n")
-    st.write(remove_dup[0:maxtags_sidebar])
-
-    for i in remove_dup[0:maxtags_sidebar]:
-        reres.append(i)
-    np.random.seed(7)
-    np.random.shuffle(reres)
-    st.write("Reranking Results: \n")
-    st.write(reres)
-
-st.write("## Results:")
-if st.button('Generated Expansion'):
-    out_res = search(query = user_query)
-    #st.success(out_res)
 
-if st.button('Rerank'):
-    out_res = reranking()
-    #st.success(out_res)
+    return df.to_dict('index')
+
+def re_rank_candidates(query, candidates, method):
+    if method == 'bm25':
+        # Filter and sort by bm25_score
+        filtered_sorted_result = sorted(
+            [(k, v) for k, v in candidates.items() if v['bm25_score'] > DEFAULT_SCORE],
+            key=lambda x: x[1]['bm25_score'],
+            reverse=True
+        )
+    elif method == 'bi_encoder':
+        # Filter and sort by bi_score
+        filtered_sorted_result = sorted(
+            [(k, v) for k, v in candidates.items() if v['bi_score'] > DEFAULT_SCORE],
+            key=lambda x: x[1]['bi_score'],
+            reverse=True
+        )
+    elif method == 'cross_encoder':
+        # Filter and sort by cross_score
+        filtered_sorted_result = sorted(
+            [(k, v) for k, v in candidates.items() if v['cross_score'] > DEFAULT_SCORE],
+            key=lambda x: x[1]['cross_score'],
+            reverse=True
+        )
+    elif method == 'gms':
+        filtered_sorted_by_encoder = sorted(
+            [(k, v) for k, v in candidates.items() if (v['cross_score'] > DEFAULT_SCORE) & (v['bi_score'] > DEFAULT_SCORE)],
+            key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],
+            reverse=True
+        )
+        # first sort by cross_score + bi_score
+        filtered_sorted_result = sorted(filtered_sorted_by_encoder, key=lambda x: x[1]['gms'], reverse=True
+        )
+    else:
+        # use default method cross_score + bi_score
+        # Filter and sort by cross_score + bi_score
+        filtered_sorted_result = sorted(
+            [(k, v) for k, v in candidates.items() if (v['cross_score'] > DEFAULT_SCORE) & (v['bi_score'] > DEFAULT_SCORE)],
+            key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],
+            reverse=True
+        )
+    data_dicts = [{'query': item[0], **item[1]} for item in filtered_sorted_result]
+    # Convert the list of dictionaries into a DataFrame
+    df = pd.DataFrame(data_dicts)
+    return df
+
+
+# st.write("## Raw Candidates:")
+if st.button('Generated Expansion'):
+    col1, col2 = st.columns(2)
+    candidates = generate_query_expansion_candidates(query = user_query)
+
+    with col1:
+        st.subheader('Original Ranking')
+        ranking_cross = re_rank_candidates(user_query, candidates, method='cross_encoder')
+        ranking_cross.index = ranking_cross.index+1
+        st.table(ranking_cross['query'][:maxtags_sidebar])
+
+    with col2:
+        st.subheader('GMS-sorted Ranking')
+        ranking_gms = re_rank_candidates(user_query, candidates, method='gms')
+        ranking_gms.index = ranking_gms.index + 1
+        st.table(ranking_gms[['query', 'gms']][:maxtags_sidebar])
+
+    ## convert into dataframe
+    # data_dicts = [{'query': key, **values} for key, values in candidates.items()]
+    # df = pd.DataFrame(data_dicts)
+    # st.write(list(candidates.keys())[0:maxtags_sidebar])
+    # st.write(df)
+    # st.dataframe(df)
+    # st.success(raw_candidates)
+
+#if st.button('Rerank By GMS'):
+    #candidates = generate_query_expansion_candidates(query = user_query)
+    #df = re_rank_candidates(user_query, candidates, method='gms')
+    #st.dataframe(df[['query', 'gms']][:maxtags_sidebar])