app.py

import streamlit as st
from streamlit_tags import st_tags, st_tags_sidebar
from keytotext import pipeline
from PIL import Image
from tabulate import tabulate

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
############

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

st.markdown("***Idea is to build a model which will take query as inputs and generate expansion information as outputs.***")
image = Image.open('etsy-shop-LLC.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 corss-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
bi_encoder = SentenceTransformer(option1,device='cpu')
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
print("Load pre-computed embeddings from disc")
# embedding_cache_path = 'embeddings.pt'
# corpus_embeddings = torch.load(embedding_cache_path)
# with open('sentences.json', 'r') as file:
#     passages = json.load(file)

embedding_cache_path = 'etsy-embeddings-cpu.pkl'
# embedding_cache_path = 'etsy-embeddings-cpu-3parts-0530.pkl'
with open(embedding_cache_path, "rb") as fIn:
  cache_data = pickle.load(fIn)
  passages = cache_data['sentences']
  corpus_embeddings = cache_data['embeddings']

from sklearn.feature_extraction import _stop_words
import string
from tqdm.autonotebook import tqdm
import numpy as np
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)
# load query GMS information
with open('query_gms_mock_2M.json', 'r') as file:
    query_gms_dict = json.load(file)
    

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

# add gms column
def add_gms_score_for_candidates(candidates):
    candidates_final = {}
    for key, value in candidates.items():
        gms_value = query_gms_dict.get(key, 0)
        candidates_final[key] = {'gms': gms_value, 'bi_score': value['bi_score'], 'cross_score': value['cross_score']}
    return candidates_final
    
def generate_query_expansion_candidates(query):
    print("Input query:", query)
    expanded_query_set = {}

    ##### 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 encoder_hits:
        corpus_id = hit['corpus_id']
        candidates[corpus_id] = {'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]

    # add gms column
    for query_candidate in final_candidates:
        value = final_candidates[query_candidate]
        value['gms'] = query_gms_dict.get(query_candidate, 0)
        final_candidates[query_candidate] = value
    # Total Results
    # st.write("E-Commerce Query Expansion Candidates: \n")
    return final_candidates

def re_rank_candidates(query, candidates, method):
    if 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
        )
    return filtered_sorted_result


if st.button('Generated Expansion'): 
    st.write("E-Commerce Query Expansion Candidates: \n")
    col1, col2 = st.columns(2)
    candidates = generate_query_expansion_candidates(query = user_query)
    with col1:
        st.subheader('Raw Candidates:')
        candidates_rerank = re_rank_candidates(user_query, candidates, method='cross_encoder')[:maxtags_sidebar]
        result = [item[0] for item in candidates_rerank]
        st.write(result)
    with col2:
        st.subheader('Rerank By GMS:')
        candidates_gms = add_gms_score_for_candidates(candidates)
        candidates_rerank = re_rank_candidates(user_query, candidates_gms, method='gms')[:maxtags_sidebar]
        data_dicts = [{'query': item[0], 'GMS Value': item[1]['gms']} for item in candidates_rerank]
        df = pd.DataFrame.from_dict(data_dicts)
        st.write(df)