Spaces:
Runtime error
Runtime error
File size: 2,235 Bytes
ad256e1 0c39a08 ad256e1 c110f1d 9c417e1 c110f1d 9c417e1 45dd76d 0c39a08 9c417e1 0c39a08 3ca293a |
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 |
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import torch
import pickle
import pandas as pd
import gradio as gr
bi_encoder = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
corpus_embeddings=pd.read_pickle("corpus_embeddings_cpu.pkl")
corpus=pd.read_pickle("corpus.pkl")
def search(query,top_k=100):
print("Top 5 Answer by the NSE:")
print()
ans=[]
##### Sematic Search #####
# Encode the query using the bi-encoder and find potentially relevant passages
question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
hits = util.semantic_search(question_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, corpus[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]
hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
for idx, hit in enumerate(hits[0:5]):
ans.append(corpus[hit['corpus_id']])
return ans[0],ans[1],ans[2],ans[3],ans[4]
exp=["Who is steve jobs?","What is coldplay?","What is a turing test?","What is the most interesting thing about our universe?","What are the most beautiful places on earth?"]
desc="This is a semantic search engine powered by SentenceTransformers (Nils_Reimers) with a retrieval and reranking system on Wikipedia corous. This will return the top 5 results. So Quest on with Transformers."
inp=gr.Textbox(lines=1, placeholder=None,label="search you query here")
out1=gr.Textbox(type="text", label="Search result 1")
out2=gr.Textbox(type="text", label="Search result 2")
out3=gr.Textbox(type="text", label="Search result 3")
out4=gr.Textbox(type="text", label="Search result 4")
out5=gr.Textbox(type="text", label="Search result 5")
iface = gr.Interface(fn=search, inputs=inp, outputs=[out1,out2,out3,out4,out5],examples=exp,article=desc,title="Neural Search Engine")
iface.launch() |