import requests from sentence_transformers import SentenceTransformer, CrossEncoder, util import os, re import torch from rank_bm25 import BM25Okapi from sklearn.feature_extraction import _stop_words import string import numpy as np import pandas as pd import base64 from io import StringIO import validators import nltk import warnings import streamlit as st from PIL import Image from beir.datasets.data_loader_hf import HFDataLoader from beir.reranking.models.mono_t5 import MonoT5 warnings.filterwarnings("ignore") auth_token = os.environ.get("auth_token") @st.cache_data() def load_data(dataset_type): corpus, queries, qrels = HFDataLoader(hf_repo="clarin-knext/"+dataset_type, streaming=False, keep_in_memory=False).load(split="test") corpus = [ doc['text']for doc in corpus] queries = [ query['text']for query in queries] return queries, corpus @st.cache_data() def bi_encode(bi_encoder_name,passages, dataset_name='scifact-pl'): global bi_encoder #We use the Bi-Encoder to encode all passages, so that we can use it with sematic search bi_encoder = SentenceTransformer(bi_encoder_name,use_auth_token=auth_token) # Thos code would be used if we would embed the passages, but here to make it fast we will load already embedded tensors: # with st.spinner('Encoding passages into a vector space...'): # if bi_encoder_name == 'intfloat/multilingual-e5-base': # corpus_embeddings = bi_encoder.encode(['passage: ' + sentence for sentence in passages], convert_to_tensor=True) # else: # corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True) with st.spinner('Loading encoded passages...'): if bi_encoder_name == "sentence-transformers/distiluse-base-multilingual-cased-v1": name = 'distiluse-base-multilingual-cased-v1' elif bi_encoder_name == 'intfloat/multilingual-e5-base': name = 'multilingual-e5-base' elif bi_encoder_name == 'nthakur/mcontriever-base-msmarco': name = 'mcontriever' corpus_embeddings_name = "-".join([name, dataset_name, "corpus"]) corpus_embeddings = torch.load(corpus_embeddings_name, map_location=torch.device('cpu')) st.success(f"Embeddings computed. Shape: {corpus_embeddings.shape}") return bi_encoder, corpus_embeddings @st.cache_resource() def cross_encode(cross_encoder_name): global cross_encoder #The bi-encoder will retrieve 100 documents. We use a cross-encoder, to re-rank the results list to improve the quality if cross_encoder_name == "clarin-knext/plt5-base-msmarco": cross_encoder = MonoT5(cross_encoder_name, use_amp=False, token_true='▁prawda', token_false='▁fałsz') else: cross_encoder = CrossEncoder(cross_encoder_name)#('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1') return cross_encoder @st.cache_data() def bm25_tokenizer(text): # We also compare the results to lexical search (keyword search). Here, we use # the BM25 algorithm which is implemented in the rank_bm25 package. # We lower case our text and remove stop-words from indexing tokenized_doc = [] for token in text.lower().split(): token = token.strip(string.punctuation) if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS: tokenized_doc.append(token) return tokenized_doc @st.cache_resource() def bm25_api(passages): tokenized_corpus = [] for passage in passages: tokenized_corpus.append(bm25_tokenizer(passage)) bm25 = BM25Okapi(tokenized_corpus) return bm25 bi_enc_options = ["sentence-transformers/distiluse-base-multilingual-cased-v1", 'intfloat/multilingual-e5-base', 'nthakur/mcontriever-base-msmarco'] # "all-mpnet-base-v2","multi-qa-MiniLM-L6-cos-v1",'intfloat/e5-base-v2',"neeva/query2query" cross_enc_options = [ 'clarin-knext/plt5-base-msmarco', 'clarin-knext/herbert-base-reranker-msmarco', 'cross-encoder/mmarco-mMiniLMv2-L12-H384-v1'] datasets_options = ["nfcorpus-pl", "scifact-pl", "fiqa-pl"] def display_df_as_table(model,top_k,score='score'): # Display the df with text and scores as a table df = pd.DataFrame([(hit[score], passages[hit['corpus_id']]) for hit in model[0:top_k]],columns=['Score','Text']) df['Score'] = round(df['Score'],2) return df #Streamlit App st.title("Retrieval BEIR-PL Demo") """ Example of retrieval over BEIR-PL dataset. """ # window_size = st.sidebar.slider("Paragraph Window Size",min_value=1,max_value=10,value=3,key= # 'slider') st.sidebar.title("Menu") dataset_type = st.sidebar.selectbox("Dataset", options=datasets_options, key='dataset_select') bi_encoder_type = st.sidebar.selectbox("Bi-Encoder", options=bi_enc_options, key='bi_select') cross_encoder_type = st.sidebar.selectbox("Cross-Encoder", options=cross_enc_options, key='cross_select') top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5,value=2) hide_bm25 = st.sidebar.checkbox("Hide BM25 results?") hide_biencoder = st.sidebar.checkbox("Hide Bi-Encoder results?") hide_crossencoder = st.sidebar.checkbox("Hide Cross-Encoder results?") # This function will search all wikipedia articles for passages that # answer the query def search_func(query, bi_encoder_type, top_k=top_k): global bi_encoder, cross_encoder st.subheader(f"Search Query:\n_{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) if not hide_bm25: st.subheader(f"Top-{top_k} lexical search (BM25) hits") bm25_df = display_df_as_table(bm25_hits,top_k) st.write(bm25_df.to_html(index=False), unsafe_allow_html=True) ##### Sematic Search ##### # Encode the query using the bi-encoder and find potentially relevant passages if bi_encoder_type == 'intfloat/multilingual-e5-base': question_embedding = bi_encoder.encode("query: " + query, convert_to_tensor=True) else: question_embedding = bi_encoder.encode(query, convert_to_tensor=True) question_embedding = question_embedding.cpu() HITS_NUM=20 # Number of hits for reranker hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=HITS_NUM,score_function=util.dot_score) 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] if not hide_biencoder: # Output of top-k hits from bi-encoder st.markdown("\n-------------------------\n") st.subheader(f"Top-{top_k} Bi-Encoder Retrieval hits") hits = sorted(hits, key=lambda x: x['score'], reverse=True) biencoder_df = display_df_as_table(hits,top_k) st.write(biencoder_df.to_html(index=False), unsafe_allow_html=True) if not hide_crossencoder: # Output of top-3 hits from re-ranker st.markdown("\n-------------------------\n") st.subheader(f"Top-{top_k} Cross-Encoder Re-ranker hits") hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) rerank_df = display_df_as_table(hits,top_k,'cross-score') st.write(rerank_df.to_html(index=False), unsafe_allow_html=True) st.markdown("---") def clear_text(): st.session_state["text_input"]= "" question, passages = load_data(dataset_type) st.write(pd.DataFrame(question[:5], columns=["Example queries from dataset"]).to_html(index=False, justify='center'), unsafe_allow_html=True) search_query = st.text_input("Ask your question:", value=question[0], key="text_input") col1, col2 = st.columns(2) with col1: search = st.button("Search",key='search_but', help='Click to Search!') with col2: clear = st.button("Clear Text Input", on_click=clear_text,key='clear',help='Click to clear the search query') if search: if bi_encoder_type: with st.spinner( text=f"Loading {bi_encoder_type} bi-encoder and embedding document into vector space. This might take a few seconds depending on the length of your document..." ): bi_encoder, corpus_embeddings = bi_encode(bi_encoder_type,passages, dataset_name=dataset_type) cross_encoder = cross_encode(cross_encoder_type) bm25 = bm25_api(passages) with st.spinner( text="Embedding completed, searching for relevant text for given query and hits..."): search_func(search_query,bi_encoder_type,top_k) st.markdown(""" """)