Spaces:
Running
Running
| import streamlit as st | |
| from sentence_transformers import SentenceTransformer, CrossEncoder, util | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| embedding_path = "abokbot/wikipedia-embedding" | |
| st.header("Wikipedia Search Engine app") | |
| st_model_load = st.text('Loading wikipedia embedding...') | |
| def load_embedding(): | |
| print("Loading embedding...") | |
| hf_hub_download(repo_id="abokbot/wikipedia-embedding", filename="simple_wikipedia_embedding.pt") | |
| wikipedia_embedding = torch.load("simple_wikipedia_embedding.pt") | |
| print("Embedding loaded!") | |
| return wikipedia_embedding | |
| wikipedia_embedding = load_embedding() | |
| st.success('Embedding loaded!') | |
| st_model_load.text("") | |
| def load_encoders(): | |
| print("Loading encoders...") | |
| bi_encoder = SentenceTransformer('msmarco-MiniLM-L-6-v3') | |
| bi_encoder.max_seq_length = 256 #Truncate long passages to 256 tokens | |
| top_k = 32 | |
| cross_encoder = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L-2-v2') | |
| return bi_encoder, cross_encoder | |
| bi_encoder, cross_encoder = load_encoders() | |
| st.success('Encoders loaded!') | |
| st_model_load.text("") | |
| if 'text' not in st.session_state: | |
| st.session_state.text = "" | |
| st_text_area = st.text_area( | |
| 'Enter query (e.g. What is the capital city of Kenya? or Number of deputees in French parliement)', | |
| value=st.session_state.text, | |
| height=100 | |
| ) | |
| """ | |
| #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search | |
| # cf https://www.sbert.net/docs/pretrained-models/msmarco-v3.html | |
| bi_encoder = SentenceTransformer('msmarco-MiniLM-L-6-v3') | |
| 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('cross-encoder/ms-marco-TinyBERT-L-2-v2') | |
| def search(query): | |
| print("Input question:", query) | |
| ##### Sematic Search ##### | |
| # Encode the query using the bi-encoder and find potentially relevant passages | |
| question_embedding = bi_encoder.encode(query, convert_to_tensor=True) | |
| question_embedding = question_embedding.cuda() | |
| 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, dataset["text"][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-3 hits from re-ranker | |
| print("\n-------------------------\n") | |
| print("Top-3 Cross-Encoder Re-ranker hits") | |
| hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) | |
| for hit in hits[0:3]: | |
| print("score: ", round(hit['cross-score'], 3),"\n", | |
| "title: ", dataset["title"][hit['corpus_id']], "\n", | |
| "substract: ", dataset["text"][hit['corpus_id']].replace("\n", " "), "\n", | |
| "link: ", dataset["url"][hit['corpus_id']],"\n") | |
| """ |