Update app.py

app.py

@@ -106,200 +106,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 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)
+    # 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)
+    
     
-    #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
+    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
-    #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)
-
+        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] = value
+    # remove the query itself from candidates
+    if query in final_candidates: 
+        del final_candidates[query]
+        
     # Total Results
-    total_qe.append(qe_string)
     st.write("E-Commerce Query Expansion Results: \n")
+    st.write(list(final_candidates.keys())[0:maxtags_sidebar])
+    return final_candidates
 
-    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)
+with open('query_gms.json', 'r') as file:
+    query_gms_dict = json.load(file)
     
-    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
-
-    ##### 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
-    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)
-
-    # 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 = []
-    rerank_list = search_nolog(query = user_query)
-    unique_list = list(set(rerank_list))
-    new_unique_list = [item for item in unique_list if item != user_query]
-    Lowercasing_list = [item.lower() for item in new_unique_list]
-
-    # st.write("E-Commerce Query Expansion Results: \n")
-    st.write(Lowercasing_list[0:maxtags_sidebar])
-
-    for i in Lowercasing_list[0:maxtags_sidebar]:
-        reres.append(i)
-    np.random.seed(7)
-    np.random.shuffle(reres)
-    test_res = {'front door': 0.5, 'family':0.3}
+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 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 == 'encoder':
+        # 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
+        )
+    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
+        )
     st.write("Reranking Results: \n")
-    st.write(test_res)
-
-st.write("## Results:")
+    st.write(filtered_sorted_result)
+
+# def reranking():
+#     rerank_list = []
+#     reres = []
+#     rerank_list = search_nolog(query = user_query)
+#     unique_list = list(set(rerank_list))
+#     new_unique_list = [item for item in unique_list if item != user_query]
+#     Lowercasing_list = [item.lower() for item in new_unique_list]
+
+#     # st.write("E-Commerce Query Expansion Results: \n")
+#     st.write(Lowercasing_list[0:maxtags_sidebar])
+
+#     for i in Lowercasing_list[0:maxtags_sidebar]:
+#         reres.append(i)
+#     np.random.seed(7)
+#     np.random.shuffle(reres)
+#     test_res = {'front door': 0.5, 'family':0.3}
+#     st.write("Reranking Results: \n")
+#     st.write(test_res)
+
+raw_candidates = generate_query_expansion_candidates(query = user_query)
+candidates = add_gms_score_for_candidates(raw_candidates, query_gms_dict)
+
+st.write("## Raw Candidates:")
 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)
+    out_res = raw_candidates
+    st.success(out_res)
+
+if st.button('Rerank By Encoder'):
+    out_res = re_rank_candidates(user_query, candidates, method='encoder')
+    st.write("Reranking By Encoder: \n")
+    st.write(out_res)
+    st.success(out_res)