Spaces:
Runtime error
Runtime error
File size: 4,978 Bytes
39caa01 |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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
############
## 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'))
option2 = st.sidebar.selectbox(
'Which corss-encoder model would you like to be selected?',
('cross-encoder/ms-marco-MiniLM-L-6-v2'))
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)
# As dataset, we use Simple English Wikipedia. Compared to the full English wikipedia, it has only
# about 170k articles. We split these articles into paragraphs and encode them with the bi-encoder
etsy_filepath = '000000000001.json'
#if not os.path.exists(wikipedia_filepath):
# util.http_get('http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz', wikipedia_filepath)
passages = []
'''
with gzip.open(wikipedia_filepath, 'rt', encoding='utf8') as fIn:
for line in fIn:
data = json.loads(line.strip())
#Add all paragraphs
#passages.extend(data['paragraphs'])
#Only add the first paragraph
passages.append(data['paragraphs'][0])
'''
with open(etsy_filepath, 'r') as EtsyJson:
for line in EtsyJson:
data = json.loads(line.strip())
#passages.append(data['query'])
passages.append(data['title'])
print("Passages:", len(passages))
# We encode all passages into our vector space. This takes about 5 minutes (depends on your GPU speed)
corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True, show_progress_bar=True)
# This function will search all wikipedia articles for passages that
# answer the query
def search(query):
print("Input question:", query)
##### BM25 search (lexical search) #####
#bm25_scores = bm25.get_scores(bm25_tokenizer(query))
#top_n = np.argpartition(bm25_scores, -5)[-5:]
#bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
#bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
#print("Top-10 lexical search (BM25) hits")
#for hit in bm25_hits[0:10]:
# print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))
##### 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) |