demo_dpr update

demo_dpr.py

@@ -20,11 +20,13 @@ DATAFRAME_FILE_BSBS = 'basecamp.csv'
 selectbox_selections = {
     'Retrieve - Rerank (with fine-tuned cross-encoder)': 1,
     'Dense Passage Retrieval':2,
+   # 'Base Dense Passage Retrieval': 3,
     'Retrieve - Rerank':4
 }
 imagebox_selections = {
     'Retrieve - Rerank (with fine-tuned cross-encoder)': 'Retrieve-rerank-trained-cross-encoder.png',
     'Dense Passage Retrieval': 'DPR_pipeline.png',
+    'Base Dense Passage Retrieval': 'base-dpr.png',
     'Retrieve - Rerank': 'retrieve-rerank.png'
 }
 
@@ -71,7 +73,7 @@ def load_paragraphs(path):
 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_bsbs.drop(['context_id', 'answer', 'answer_start', 'answer_end'], 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
@@ -82,11 +84,31 @@ def dot_product(question_output, context_output):
     result = torch.dot(mat1, mat2)
     return result
 
+def base_dpr_pipeline(question):
+    tokenized_question = question_tokenizer(question, padding=True, truncation=True, return_tensors="pt")
+    question_embedding = base_dpr_context_encoder(**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(base_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(base_contexes[j])
+        top_5_scores.append(results_list[j])
+    return top_5_contexes, top_5_scores
+
 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 == 3:
+ #       return base_dpr_pipeline(question) # DPR only
     if search_method == 4:
         return retrieve_rerank(question)
 
@@ -213,8 +235,8 @@ def qa_main_widgetsv2():
             st.write(selection['context'])
             st.markdown("### Question:")
             st.write(selection['question'])
-            st.markdown("### Answer:")
-            st.write(selection['answer'])
+       #     st.markdown("### Answer:")
+       #     st.write(selection['answer'])
             st.session_state.grid_click_2 = False
 
 @st.cache(show_spinner=False, allow_output_mutation = True)
@@ -226,15 +248,17 @@ def load_models(dpr_model_path, auth_token, cross_encoder_model_path):
     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)
+    base_dpr_context_encoder = AutoModel.from_pretrained('facebook/contriever-msmarco')
     question_tokenizer = AutoTokenizer.from_pretrained('facebook/contriever-msmarco')
-    return dpr_trained, bi_encoder, cross_encoder, trained_cross_encoder, question_tokenizer
+    return dpr_trained, bi_encoder, cross_encoder, trained_cross_encoder, question_tokenizer, base_dpr_context_encoder
 
 context_embeddings, contexes = load_paragraphs('st-context-embeddings.pkl')
 dpr_context_embeddings, dpr_contexes = load_paragraphs('basecamp-dpr-contriever-embeddings.pkl')
+base_dpr_context_embeddings, base_contexes = load_paragraphs('basecamp-base-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"]))
+dpr_trained, bi_encoder, cross_encoder, trained_cross_encoder, question_tokenizer, base_dpr_context_encoder = 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)
+#    top_5_contexes, top_5_scores = search_pipeline('What contributions does 37Signals make to open-source projects?', 3)