import copy import streamlit as st import pandas as pd from sentence_transformers import SentenceTransformer, util from sentence_transformers.cross_encoder import CrossEncoder from st_aggrid import GridOptionsBuilder, AgGrid import pickle import torch from transformers import AutoTokenizer, AutoModel from pathlib import Path import base64 import io st.set_page_config(layout="wide") DATAFRAME_FILE_ORIGINAL = 'policyQA_original.csv' DATAFRAME_FILE_BSBS = 'basecamp.csv' selectbox_selections = { 'Retrieve - Rerank (with fine-tuned cross-encoder)': 1, 'Dense Passage Retrieval':2, 'Retrieve - Rerank':4 } imagebox_selections = { 'Retrieve - Rerank (with fine-tuned cross-encoder)': 'Retrieve-rerank-trained-cross-encoder.png', 'Dense Passage Retrieval': 'DPR_pipeline.png', 'Retrieve - Rerank': 'retrieve-rerank.png' } def retrieve_rerank(question): # Semantic Search (Retrieve) question_embedding = bi_encoder.encode(question, convert_to_tensor=True) hits = util.semantic_search(question_embedding, context_embeddings, top_k=100) if len(hits) == 0: return [] hits = hits[0] # Rerank - score all retrieved passages with cross-encoder cross_inp = [[question, contexes[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-5 hits from re-ranker hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) top_5_contexes = [] top_5_scores = [] for hit in hits[0:20]: top_5_contexes.append(contexes[hit['corpus_id']]) top_5_scores.append(hit['cross-score']) return top_5_contexes, top_5_scores class CPU_Unpickler(pickle.Unpickler): def find_class(self, module, name): if module == 'torch.storage' and name == '_load_from_bytes': return lambda b: torch.load(io.BytesIO(b), map_location='cpu') else: return super().find_class(module, name) @st.cache(show_spinner=False, allow_output_mutation=True) def load_paragraphs(path): with open(path, "rb") as fIn: cache_data = pickle.load(fIn) corpus_sentences = cache_data['contexes'] corpus_embeddings = cache_data['embeddings'] return corpus_embeddings, corpus_sentences @st.cache(show_spinner=False) def load_dataframes(): # data_original = pd.read_csv(DATAFRAME_FILE_ORIGINAL, index_col=0, sep='|') data_bsbs = pd.read_csv(DATAFRAME_FILE_BSBS, index_col=0, sep='|') data_bsbs.drop('context_id', axis=1, inplace=True) # data_original = data_original.sample(frac=1).reset_index(drop=True) data_bsbs = data_bsbs.sample(frac=1).reset_index(drop=True) return data_bsbs def dot_product(question_output, context_output): mat1 = torch.squeeze(question_output, 0) mat2 = torch.squeeze(context_output, 0) result = torch.dot(mat1, mat2) return result def search_pipeline(question, search_method): if search_method == 1: #Retrieve - rerank with fine-tuned cross encoder return retrieve_rerank_with_trained_cross_encoder(question) if search_method == 2: return custom_dpr_pipeline(question) # DPR only if search_method == 4: return retrieve_rerank(question) def mean_pooling(token_embeddings, mask): token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.) sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None] return sentence_embeddings def custom_dpr_pipeline(question): #paragraphs tokenized_question = question_tokenizer(question, padding=True, truncation=True, return_tensors="pt") question_embedding = dpr_trained.model.question_model(**tokenized_question) question_embedding = mean_pooling(question_embedding[0], tokenized_question['attention_mask']) # question_embedding = question_embedding['pooler_output'] results_list = [] for i,context_embedding in enumerate(dpr_context_embeddings): score = dot_product(question_embedding, context_embedding) results_list.append(score.detach().cpu()) hits = sorted(range(len(results_list)), key=lambda i: results_list[i], reverse=True) top_5_contexes = [] top_5_scores = [] for j in hits[0:5]: top_5_contexes.append(dpr_contexes[j]) top_5_scores.append(results_list[j]) return top_5_contexes, top_5_scores def retrieve(question): # Semantic Search (Retrieve) question_embedding = bi_encoder.encode(question, convert_to_tensor=True) hits = util.semantic_search(question_embedding, context_embeddings, top_k=100) if len(hits) == 0: return [] hits = hits[0] return hits def retrieve_rerank_with_trained_cross_encoder(question): hits = retrieve(question) cross_inp = [(question, contexes[hit['corpus_id']]) for hit in hits] cross_scores = trained_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-5 hits from re-ranker hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) top_5_contexes = [] top_5_scores = [] for hit in hits[0:5]: top_5_contexes.append(contexes[hit['corpus_id']]) top_5_scores.append(hit['cross-score']) return top_5_contexes, top_5_scores def interactive_table(dataframe): gb = GridOptionsBuilder.from_dataframe(dataframe) gb.configure_pagination(paginationAutoPageSize=True) gb.configure_side_bar() gb.configure_selection('single', rowMultiSelectWithClick=True, groupSelectsChildren="Group checkbox select children") # Enable multi-row selection gridOptions = gb.build() grid_response = AgGrid( dataframe, gridOptions=gridOptions, data_return_mode='AS_INPUT', update_mode='SELECTION_CHANGED', enable_enterprise_modules=False, fit_columns_on_grid_load=False, theme='streamlit', # Add theme color to the table height=350, width='100%', reload_data=False ) return grid_response def img_to_bytes(img_path): img_bytes = Path(img_path).read_bytes() encoded = base64.b64encode(img_bytes).decode() return encoded def qa_main_widgetsv2(): st.title("Question Answering Demo") st.markdown("""---""") option = st.selectbox("Select a search method:", list(selectbox_selections.keys())) header_html = "
".format( img_to_bytes(imagebox_selections[option]) ) st.markdown( header_html, unsafe_allow_html=True, ) st.markdown("""---""") col1, col3 = st.columns([1, 1]) with col1: form = st.form(key='first_form') question = form.text_area("What is your question?:", height=200) submit = form.form_submit_button('Submit') if "form_submit" not in st.session_state: st.session_state.form_submit = False if submit: st.session_state.form_submit = True if st.session_state.form_submit and question != '': with st.spinner(text='Related context search in progress..'): top_5_contexes, top_5_scores = search_pipeline(question.strip(), selectbox_selections[option]) if len(top_5_contexes) == 0: st.error("Related context not found!") st.session_state.form_submit = False else: for i, context in enumerate(top_5_contexes): st.markdown(f"## Related Context - {i + 1} (score: {top_5_scores[i]:.2f})") st.markdown(context) st.markdown("""---""") with col3: st.markdown("## Our Questions") grid_response = interactive_table(dataframe_bsbs) data2 = grid_response['selected_rows'] if "grid_click_2" not in st.session_state: st.session_state.grid_click_2 = False if len(data2) > 0: st.session_state.grid_click_2 = True if st.session_state.grid_click_2: selection = data2[0] # st.markdown("## Context & Answer:") st.markdown("### Context:") st.write(selection['context']) st.markdown("### Question:") st.write(selection['question']) st.markdown("### Answer:") st.write(selection['answer']) st.session_state.grid_click_2 = False @st.cache(show_spinner=False, allow_output_mutation = True) def load_models(dpr_model_path, auth_token, cross_encoder_model_path): dpr_trained = AutoModel.from_pretrained(dpr_model_path, use_auth_token=auth_token, trust_remote_code=True) bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') bi_encoder.max_seq_length = 500 trained_cross_encoder = CrossEncoder(cross_encoder_model_path) question_tokenizer = AutoTokenizer.from_pretrained('facebook/contriever-msmarco') return dpr_trained, bi_encoder, cross_encoder, trained_cross_encoder, question_tokenizer context_embeddings, contexes = load_paragraphs('st-context-embeddings.pkl') dpr_context_embeddings, dpr_contexes = load_paragraphs('basecamp-dpr-contriever-embeddings.pkl') dataframe_bsbs = load_dataframes() dpr_trained, bi_encoder, cross_encoder, trained_cross_encoder, question_tokenizer = copy.deepcopy(load_models(st.secrets["DPR_MODEL_PATH"], st.secrets["AUTH_TOKEN"], st.secrets["CROSS_ENCODER_MODEL_PATH"])) qa_main_widgetsv2() #if __name__ == '__main__': # top_5_contexes, top_5_scores = search_pipeline('What are the benefits of 37Signals Visa Card?', 1)