app.py

import streamlit as st

from PIL import Image

import json
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import pickle
import pandas as pd

############
## Main page
############

st.write("# Demonstration for User Query Expansion（QE）")

st.markdown("***Idea is to build a model which will take query as inputs and generate expansion information as outputs.***")
image = Image.open('top.png')
st.image(image)

st.sidebar.write("# Top-N Selection")
maxtags_sidebar = st.sidebar.slider('Number of query allowed?', 1, 20, 1, key='ehikwegrjifbwreuk')
#user_query = st_tags(
#    label='# Enter Query:',
#    text='Press enter to add more',
#    value=['Mother'],
#    suggestions=['gift', 'nike', 'wool'],
#    maxtags=maxtags_sidebar,
#    key="aljnf")

user_query = st.text_input("Enter a query for the generated text: e.g., gift, home decoration ...")

# Add selectbox in streamlit
option1 = st.sidebar.selectbox(
     'Which transformers model would you like to be selected?',
     ('multi-qa-MiniLM-L6-cos-v1','null','null'))

option2 = st.sidebar.selectbox(
     '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!")

#if not torch.cuda.is_available():
#    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
@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

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'
passages,corpus_embeddings = load_pickle(embedding_cache_path)


from rank_bm25 import BM25Okapi
from sklearn.feature_extraction import _stop_words
import string
from tqdm.autonotebook import tqdm
import numpy as np
import re

import yake

@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_mock_2M.json')
# We lower case our text and remove stop-words from indexing
def bm25_tokenizer(text):
    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.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):
    return len([i for i in s.split(' ') if i])


# This function will search all wikipedia articles for passages that
# answer the 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)
    expanded_query_set = {}

    ##### BM25 search (lexical search) #####
    bm25_scores = bm25.get_scores(bm25_tokenizer(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)


    ##### Sematic Search #####
    # Encode the query using the bi-encoder and find potentially relevant passages
    query_embedding = bi_encoder.encode(query, convert_to_tensor=True)
    # 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]

    # 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)
    for idx in range(len(cross_scores)):
        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

    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])