Spaces:
Runtime error
Runtime error
File size: 3,792 Bytes
39caa01 c7496eb 39caa01 1d1d1d9 39caa01 1d1d1d9 39caa01 c1d264e 39caa01 927223b c1d264e 927223b f6c2115 927223b 39caa01 c1d264e 39caa01 1d1d1d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import streamlit as st
from streamlit_tags import st_tags, st_tags_sidebar
from keytotext import pipeline
from PIL import Image
import json
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import gzip
import os
import torch
import pickle
############
## Main page
############
st.write("# Code for Query Expansion")
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("# Parameter Selection")
maxtags_sidebar = st.sidebar.slider('Number of query allowed?', 1, 10, 1, key='ehikwegrjifbwreuk')
user_query = st_tags(
label='# Enter Query:',
text='Press enter to add more',
value=['Mother'],
suggestions=['five', 'six', 'seven', 'eight', 'nine', 'three', 'eleven', 'ten', 'four'],
maxtags=maxtags_sidebar,
key="aljnf")
# 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)
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)
# load pre-train embeedings files
embedding_cache_path = 'etsy-embeddings-cpu.pkl'
print("Load pre-computed embeddings from disc")
with open(embedding_cache_path, "rb") as fIn:
cache_data = pickle.load(fIn)
#corpus_sentences = cache_data['sentences']
corpus_embeddings = cache_data['embeddings']
# This function will search all wikipedia articles for passages that
# answer the query
def search(query):
print("Input question:", query)
##### 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()
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
print("\n-------------------------\n")
print("Top-10 Bi-Encoder Retrieval hits")
hits = sorted(hits, key=lambda x: x['score'], reverse=True)
for hit in hits[0:10]:
print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
# Output of top-10 hits from re-ranker
print("\n-------------------------\n")
print("Top-10 Cross-Encoder Re-ranker hits")
hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
for hit in hits[0:10]:
print("\t{:.3f}\t{}".format(hit['cross-score'], passages[hit['corpus_id']].replace("\n", " ")))
st.write("## Results:")
if st.button('Generate Sentence'):
out = search(query = user_query)
st.success(out) |