HarryLee/QueryExpansionForEtsy
0
1import streamlit as st2from streamlit_tags import st_tags, st_tags_sidebar3from keytotext import pipeline4from PIL import Image5 6import json7from sentence_transformers import SentenceTransformer, CrossEncoder, util8import gzip9import os10import torch11import pickle12import random13import numpy as np14import pandas as pd15 16############17## Main page18############19 20st.write("# Demonstration for Etsy Query Expansion（Etsy-QE）")21 22st.markdown("***Idea is to build a model which will take query as inputs and generate expansion information as outputs.***")23image = Image.open('etsy-shop-LLC.png')24st.image(image)25 26st.sidebar.write("# Top-N Selection")27maxtags_sidebar = st.sidebar.slider('Number of query allowed?', 1, 20, 1, key='ehikwegrjifbwreuk')28#user_query = st_tags(29# label='# Enter Query:',30# text='Press enter to add more',31# value=['Mother'],32# suggestions=['gift', 'nike', 'wool'],33# maxtags=maxtags_sidebar,34# key="aljnf")35 36user_query = st.text_input("Enter a query for the generated text: e.g., gift, home decoration ...")37 38# Add selectbox in streamlit39option1 = st.sidebar.selectbox(40 'Which transformers model would you like to be selected?',41 ('multi-qa-MiniLM-L6-cos-v1','null','null'))42 43option2 = st.sidebar.selectbox(44 'Which corss-encoder model would you like to be selected?',45 ('cross-encoder/ms-marco-MiniLM-L-6-v2','null','null'))46 47st.sidebar.success("Load Successfully!")48 49#if not torch.cuda.is_available():50# print("Warning: No GPU found. Please add GPU to your notebook")51 52#We use the Bi-Encoder to encode all passages, so that we can use it with sematic search53bi_encoder = SentenceTransformer(option1,device='cpu')54bi_encoder.max_seq_length = 256 #Truncate long passages to 256 tokens55top_k = 32 #Number of passages we want to retrieve with the bi-encoder56 57#The bi-encoder will retrieve 100 documents. We use a cross-encoder, to re-rank the results list to improve the quality58cross_encoder = CrossEncoder(option2, device='cpu')59 60passages = []61 62# load pre-train embeedings files63embedding_cache_path = 'etsy-embeddings-cpu.pkl'64print("Load pre-computed embeddings from disc")65with open(embedding_cache_path, "rb") as fIn:66 cache_data = pickle.load(fIn)67 passages = cache_data['sentences']68 corpus_embeddings = cache_data['embeddings']69 70from rank_bm25 import BM25Okapi71from sklearn.feature_extraction import _stop_words72import string73from tqdm.autonotebook import tqdm74import numpy as np75import re76 77import yake78 79language = "en"80max_ngram_size = 381deduplication_threshold = 0.982deduplication_algo = 'seqm'83windowSize = 384numOfKeywords = 385 86custom_kw_extractor = yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold, dedupFunc=deduplication_algo, windowsSize=windowSize, top=numOfKeywords, features=None)87# load query GMS information88with open('query_gms.json', 'r') as file:89 query_gms_dict = json.load(file)90 91# We lower case our text and remove stop-words from indexing92def bm25_tokenizer(text):93 tokenized_doc = []94 for token in text.lower().split():95 token = token.strip(string.punctuation)96 97 if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:98 tokenized_doc.append(token)99 return tokenized_doc100 101tokenized_corpus = []102for passage in tqdm(passages):103 tokenized_corpus.append(bm25_tokenizer(passage))104 105bm25 = BM25Okapi(tokenized_corpus)106 107def word_len(s):108 return len([i for i in s.split(' ') if i])109 110 111# This function will search all wikipedia articles for passages that112# answer the query113DEFAULT_SCORE = -100.0114def clean_string(input_string):115 string_sub1 = re.sub("([^\u0030-\u0039\u0041-\u007a])", ' ', input_string)116 string_sub2 = re.sub("\x20\x20", "\n", string_sub1)117 string_strip = string_sub2.strip().lower()118 output_string = []119 if len(string_strip) > 20:120 keywords = custom_kw_extractor.extract_keywords(string_strip)121 for tokens in keywords:122 string_clean = tokens[0]123 if word_len(string_clean) > 1:124 output_string.append(string_clean)125 else:126 output_string.append(string_strip)127 return output_string128 129def add_gms_score_for_candidates(candidates, query_gms_dict):130 for query_candidate in candidates:131 value = candidates[query_candidate]132 value['gms'] = query_gms_dict.get(query_candidate, 0)133 candidates[query_candidate] = value134 return candidates135 136def generate_query_expansion_candidates(query):137 print("Input query:", query)138 expanded_query_set = {}139 140 ##### BM25 search (lexical search) #####141 bm25_scores = bm25.get_scores(bm25_tokenizer(query))142 # finds the indices of the top n scores143 top_n_indices = np.argpartition(bm25_scores, -5)[-5:]144 bm25_hits = [{'corpus_id': idx, 'bm25_score': bm25_scores[idx]} for idx in top_n_indices]145 # bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)146 147 148 ##### Sematic Search #####149 # Encode the query using the bi-encoder and find potentially relevant passages150 query_embedding = bi_encoder.encode(query, convert_to_tensor=True)151 # query_embedding = query_embedding.cuda()152 # Get the hits for the first query153 encoder_hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)[0]154 155 # For all retrieved passages, add the cross_encoder scores156 cross_inp = [[query, passages[hit['corpus_id']]] for hit in encoder_hits]157 cross_scores = cross_encoder.predict(cross_inp)158 for idx in range(len(cross_scores)):159 encoder_hits[idx]['cross_score'] = cross_scores[idx]160 161 candidates = {}162 for hit in bm25_hits:163 corpus_id = hit['corpus_id']164 if corpus_id not in candidates:165 candidates[corpus_id] = {'bm25_score': hit['bm25_score'], 'bi_score': DEFAULT_SCORE, 'cross_score': DEFAULT_SCORE}166 for hit in encoder_hits:167 corpus_id = hit['corpus_id']168 if corpus_id not in candidates:169 candidates[corpus_id] = {'bm25_score': DEFAULT_SCORE, 'bi_score': hit['score'], 'cross_score': hit['cross_score']}170 else:171 bm25_score = candidates[corpus_id]['bm25_score']172 candidates[corpus_id].update({'bm25_score': bm25_score, 'bi_score': hit['score'], 'cross_score': hit['cross_score']})173 174 final_candidates = {}175 for key, value in candidates.items():176 input_string = passages[key].replace("\n", "")177 string_set = set(clean_string(input_string))178 for item in string_set:179 final_candidates[item] = value180 # remove the query itself from candidates181 if query in final_candidates: 182 del final_candidates[query]183 184 # add gms column185 for query_candidate in final_candidates:186 value = final_candidates[query_candidate]187 value['gms'] = query_gms_dict.get(query_candidate, 0)188 final_candidates[query_candidate] = value189 # Total Results190 st.write("E-Commerce Query Expansion Candidates: \n")191 return final_candidates192 193def re_rank_candidates(query, candidates, method):194 if method == 'bm25':195 # Filter and sort by bm25_score196 filtered_sorted_result = sorted(197 [(k, v) for k, v in candidates.items() if v['bm25_score'] > DEFAULT_SCORE],198 key=lambda x: x[1]['bm25_score'],199 reverse=True200 )201 elif method == 'bi_encoder':202 # Filter and sort by bi_score203 filtered_sorted_result = sorted(204 [(k, v) for k, v in candidates.items() if v['bi_score'] > DEFAULT_SCORE],205 key=lambda x: x[1]['bi_score'],206 reverse=True207 )208 elif method == 'cross_encoder':209 # Filter and sort by cross_score210 filtered_sorted_result = sorted(211 [(k, v) for k, v in candidates.items() if v['cross_score'] > DEFAULT_SCORE],212 key=lambda x: x[1]['cross_score'],213 reverse=True214 )215 elif method == 'gms':216 filtered_sorted_by_encoder = sorted(217 [(k, v) for k, v in candidates.items() if (v['cross_score'] > DEFAULT_SCORE) & (v['bi_score'] > DEFAULT_SCORE)],218 key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],219 reverse=True220 )221 # first sort by cross_score + bi_score222 filtered_sorted_result = sorted(filtered_sorted_by_encoder, key=lambda x: x[1]['gms'], reverse=True223 )224 else:225 # use default method cross_score + bi_score226 # Filter and sort by cross_score + bi_score227 filtered_sorted_result = sorted(228 [(k, v) for k, v in candidates.items() if (v['cross_score'] > DEFAULT_SCORE) & (v['bi_score'] > DEFAULT_SCORE)],229 key=lambda x: x[1]['cross_score'] + x[1]['bi_score'],230 reverse=True231 )232 data_dicts = [{'query': item[0], **item[1]} for item in filtered_sorted_result]233 # Convert the list of dictionaries into a DataFrame234 df = pd.DataFrame(data_dicts)235 return df236 237 238# st.write("## Raw Candidates:")239if st.button('Generated Expansion'): 240 candidates = generate_query_expansion_candidates(query = user_query)241 df = re_rank_candidates(user_query, candidates, method='cross_encoder')242 result = list(df['query'][:maxtags_sidebar])243 st.write(result)244 ## convert into dataframe245 # data_dicts = [{'query': key, **values} for key, values in candidates.items()]246 # df = pd.DataFrame(data_dicts)247 # st.write(list(candidates.keys())[0:maxtags_sidebar])248 # st.write(df)249 # st.dataframe(df)250 # st.success(raw_candidates)251 252if st.button('Rerank By GMS'):253 candidates = generate_query_expansion_candidates(query = user_query)254 df = re_rank_candidates(user_query, candidates, method='gms')255 st.dataframe(df[['query', 'gms']][:maxtags_sidebar])