HarryLee/QueryExpansionForEtsy
0
1import streamlit as st2 3from PIL import Image4 5import json6from sentence_transformers import SentenceTransformer, CrossEncoder, util7import pickle8import pandas as pd9 10############11## Main page12############13 14st.write("# Demonstration for User Query Expansion（QE）")15 16st.markdown("***Idea is to build a model which will take query as inputs and generate expansion information as outputs.***")17image = Image.open('top.png')18st.image(image)19 20st.sidebar.write("# Top-N Selection")21maxtags_sidebar = st.sidebar.slider('Number of query allowed?', 1, 20, 1, key='ehikwegrjifbwreuk')22#user_query = st_tags(23# label='# Enter Query:',24# text='Press enter to add more',25# value=['Mother'],26# suggestions=['gift', 'nike', 'wool'],27# maxtags=maxtags_sidebar,28# key="aljnf")29 30user_query = st.text_input("Enter a query for the generated text: e.g., gift, home decoration ...")31 32# Add selectbox in streamlit33option1 = st.sidebar.selectbox(34 'Which transformers model would you like to be selected?',35 ('multi-qa-MiniLM-L6-cos-v1','null','null'))36 37option2 = st.sidebar.selectbox(38 'Which cross-encoder model would you like to be selected?',39 ('cross-encoder/ms-marco-MiniLM-L-6-v2','null','null'))40 41st.sidebar.success("Load Successfully!")42 43#if not torch.cuda.is_available():44# print("Warning: No GPU found. Please add GPU to your notebook")45 46#We use the Bi-Encoder to encode all passages, so that we can use it with sematic search47@st.cache_resource48def load_encoders(sentence_enc, cross_enc):49 return SentenceTransformer(sentence_enc,device='cpu'), CrossEncoder(cross_enc,device='cpu')50bi_encoder, cross_encoder = load_encoders(option1,option2)51bi_encoder.max_seq_length = 256 #Truncate long passages to 256 tokens52top_k = 32 #Number of passages we want to retrieve with the bi-encoder53 54passages = []55 56# load pre-train embeedings files57@st.cache_resource58def load_pickle(path):59 with open(path, "rb") as fIn:60 cache_data = pickle.load(fIn)61 passages = cache_data['sentences']62 corpus_embeddings = cache_data['embeddings']63 print("Load pre-computed embeddings from disc")64 return passages,corpus_embeddings65 66embedding_cache_path = 'etsy-embeddings-cpu.pkl'67passages,corpus_embeddings = load_pickle(embedding_cache_path)68 69 70from rank_bm25 import BM25Okapi71from sklearn.feature_extraction import _stop_words72import string73from tqdm.autonotebook import tqdm74import numpy as np75import re76 77import yake78 79@st.cache_resource80def load_model():81 language = "en"82 max_ngram_size = 383 deduplication_threshold = 0.984 deduplication_algo = 'seqm'85 windowSize = 386 numOfKeywords = 387 return yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold, dedupFunc=deduplication_algo, windowsSize=windowSize, top=numOfKeywords, features=None)88custom_kw_extractor = load_model()89# load query GMS information90@st.cache_resource91def load_json(path):92 with open(path, 'r') as file:93 query_gms_dict = json.load(file)94 return query_gms_dict95 96query_gms_dict = load_json('query_gms_mock_2M.json')97# We lower case our text and remove stop-words from indexing98def bm25_tokenizer(text):99 tokenized_doc = []100 for token in text.lower().split():101 token = token.strip(string.punctuation)102 103 if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:104 tokenized_doc.append(token)105 return tokenized_doc106 107@st.cache_resource108def get_tokenized_corpus(passages,_tokenizer):109 tokenized_corpus = []110 for passage in passages:111 tokenized_corpus.append(_tokenizer(passage))112 return tokenized_corpus113 114tokenized_corpus = get_tokenized_corpus(passages,bm25_tokenizer)115bm25 = BM25Okapi(tokenized_corpus)116 117def word_len(s):118 return len([i for i in s.split(' ') if i])119 120 121# This function will search all wikipedia articles for passages that122# answer the query123DEFAULT_SCORE = -100.0124def clean_string(input_string):125 string_sub1 = re.sub("([^\u0030-\u0039\u0041-\u007a])", ' ', input_string)126 string_sub2 = re.sub("\x20\x20", "\n", string_sub1)127 string_strip = string_sub2.strip().lower()128 output_string = []129 if len(string_strip) > 20:130 keywords = custom_kw_extractor.extract_keywords(string_strip)131 for tokens in keywords:132 string_clean = tokens[0]133 if word_len(string_clean) > 1:134 output_string.append(string_clean)135 else:136 output_string.append(string_strip)137 return output_string138 139# def add_gms_score_for_candidates(candidates, query_gms_dict):140# for query_candidate in candidates:141# value = candidates[query_candidate]142# value['gms'] = query_gms_dict.get(query_candidate, 0)143# candidates[query_candidate] = value144# return candidates145 146def generate_query_expansion_candidates(query):147 print("Input query:", query)148 expanded_query_set = {}149 150 ##### BM25 search (lexical search) #####151 bm25_scores = bm25.get_scores(bm25_tokenizer(query))152 # finds the indices of the top n scores153 top_n_indices = np.argpartition(bm25_scores, -5)[-5:]154 bm25_hits = [{'corpus_id': idx, 'bm25_score': bm25_scores[idx]} for idx in top_n_indices]155 # bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)156 157 158 ##### Sematic Search #####159 # Encode the query using the bi-encoder and find potentially relevant passages160 query_embedding = bi_encoder.encode(query, convert_to_tensor=True)161 # query_embedding = query_embedding.cuda()162 # Get the hits for the first query163 encoder_hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)[0]164 165 # For all retrieved passages, add the cross_encoder scores166 cross_inp = [[query, passages[hit['corpus_id']]] for hit in encoder_hits]167 cross_scores = cross_encoder.predict(cross_inp)168 for idx in range(len(cross_scores)):169 encoder_hits[idx]['cross_score'] = cross_scores[idx]170 171 candidates = {}172 for hit in bm25_hits:173 corpus_id = hit['corpus_id']174 if corpus_id not in candidates:175 candidates[corpus_id] = {'bm25_score': hit['bm25_score'], 'bi_score': DEFAULT_SCORE, 'cross_score': DEFAULT_SCORE}176 for hit in encoder_hits:177 corpus_id = hit['corpus_id']178 if corpus_id not in candidates:179 candidates[corpus_id] = {'bm25_score': DEFAULT_SCORE, 'bi_score': hit['score'], 'cross_score': hit['cross_score']}180 else:181 bm25_score = candidates[corpus_id]['bm25_score']182 candidates[corpus_id].update({'bm25_score': bm25_score, 'bi_score': hit['score'], 'cross_score': hit['cross_score']})183 184 final_candidates = {}185 for key, value in candidates.items():186 input_string = passages[key].replace("\n", "")187 string_set = set(clean_string(input_string))188 for item in string_set:189 final_candidates[item.replace("\n", " ")] = value190 # remove the query itself from candidates191 if query in final_candidates:192 del final_candidates[query]193 # print(final_candidates)194 # add gms column195 df = pd.DataFrame(final_candidates).T196 df['gms'] = [query_gms_dict.get(i,0) for i in df.index]197 # Total Results198 199 return df.to_dict('index')200 201def re_rank_candidates(query, candidates, method):202 if method == 'bm25':203 # Filter and sort by bm25_score204 filtered_sorted_result = sorted(205 [(k, v) for k, v in candidates.items() if v['bm25_score'] > DEFAULT_SCORE],206 key=lambda x: x[1]['bm25_score'],207 reverse=True208 )209 elif method == 'bi_encoder':210 # Filter and sort by bi_score211 filtered_sorted_result = sorted(212 [(k, v) for k, v in candidates.items() if v['bi_score'] > DEFAULT_SCORE],213 key=lambda x: x[1]['bi_score'],214 reverse=True215 )216 elif method == 'cross_encoder':217 # Filter and sort by cross_score218 filtered_sorted_result = sorted(219 [(k, v) for k, v in candidates.items() if v['cross_score'] > DEFAULT_SCORE],220 key=lambda x: x[1]['cross_score'],221 reverse=True222 )223 elif method == 'gms':224 filtered_sorted_by_encoder = sorted(225 [(k, v) for k, v in candidates.items() if (v['cross_score'] > DEFAULT_SCORE) & (v['bi_score'] > DEFAULT_SCORE)],226 key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],227 reverse=True228 )229 # first sort by cross_score + bi_score230 filtered_sorted_result = sorted(filtered_sorted_by_encoder, key=lambda x: x[1]['gms'], reverse=True231 )232 else:233 # use default method cross_score + bi_score234 # Filter and sort by cross_score + bi_score235 filtered_sorted_result = sorted(236 [(k, v) for k, v in candidates.items() if (v['cross_score'] > DEFAULT_SCORE) & (v['bi_score'] > DEFAULT_SCORE)],237 key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],238 reverse=True239 )240 data_dicts = [{'query': item[0], **item[1]} for item in filtered_sorted_result]241 # Convert the list of dictionaries into a DataFrame242 df = pd.DataFrame(data_dicts)243 return df244 245 246# st.write("## Raw Candidates:")247if st.button('Generated Expansion'):248 col1, col2 = st.columns(2)249 candidates = generate_query_expansion_candidates(query = user_query)250 251 with col1:252 st.subheader('Original Ranking')253 ranking_cross = re_rank_candidates(user_query, candidates, method='cross_encoder')254 ranking_cross.index = ranking_cross.index+1255 st.table(ranking_cross['query'][:maxtags_sidebar])256 257 with col2:258 st.subheader('GMS-sorted Ranking')259 ranking_gms = re_rank_candidates(user_query, candidates, method='gms')260 ranking_gms.index = ranking_gms.index + 1261 st.table(ranking_gms[['query', 'gms']][:maxtags_sidebar])262 263 ## convert into dataframe264 # data_dicts = [{'query': key, **values} for key, values in candidates.items()]265 # df = pd.DataFrame(data_dicts)266 # st.write(list(candidates.keys())[0:maxtags_sidebar])267 # st.write(df)268 # st.dataframe(df)269 # st.success(raw_candidates)270 271#if st.button('Rerank By GMS'):272 #candidates = generate_query_expansion_candidates(query = user_query)273 #df = re_rank_candidates(user_query, candidates, method='gms')274 #st.dataframe(df[['query', 'gms']][:maxtags_sidebar])