algomuffin/neural-search-engine

1from sentence_transformers import SentenceTransformer, CrossEncoder, util2import torch3import pickle4import pandas as pd5import gradio as gr6bi_encoder = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")7cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")8corpus_embeddings=pd.read_pickle("corpus_embeddings_cpu.pkl")9corpus=pd.read_pickle("corpus.pkl")10def search(query,top_k=100):11    print("Top 5 Answer by the NSE:")12    print()13    ans=[]14    ##### Sematic Search #####15    # Encode the query using the bi-encoder and find potentially relevant passages16    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)17    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)18    hits = hits[0]  # Get the hits for the first query19    ##### Re-Ranking #####20    # Now, score all retrieved passages with the cross_encoder21    cross_inp = [[query, corpus[hit['corpus_id']]] for hit in hits]22    cross_scores = cross_encoder.predict(cross_inp)23    # Sort results by the cross-encoder scores24    for idx in range(len(cross_scores)):25        hits[idx]['cross-score'] = cross_scores[idx]26    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)27    28    for idx, hit in enumerate(hits[0:5]):29        ans.append(corpus[hit['corpus_id']])30    return ans[0],ans[1],ans[2],ans[3],ans[4]31iface = gr.Interface(fn=search, inputs=["text"], outputs=["textbox","textbox","textbox","textbox","textbox"],examples=["How big is London?", "Where is Rome?","Who is steve jobs?","What is the most interesting thing about our universe?"],article="This is a semantic search engine powered by SentenceTransformers (Nils_Reimers) with a retrieval and reranking system on Wikipedia corpus. It will show the top 5 results",title="Neural Search Engine").launch()32