|
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=[] |
|
|
|
|
|
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] |
|
|
|
|
|
|
|
cross_inp = [[query, corpus[hit['corpus_id']]] for hit in hits] |
|
cross_scores = cross_encoder.predict(cross_inp) |
|
|
|
|
|
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] |
|
|
|
iface = gr.Interface(fn=search, inputs=["text"], outputs=["textbox","textbox","textbox","textbox","textbox"]).launch(share=True) |
