files updated

basecamp-dpr-contriever-embeddings.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:413837017c7b17e8e44556d9ab0cc9d42c9b24d3d28b29a39f3e7e143bd9f482
+size 856086

demo_dpr.py

@@ -7,7 +7,7 @@ from sentence_transformers.cross_encoder import CrossEncoder
 from st_aggrid import GridOptionsBuilder, AgGrid
 import pickle
 import torch
-from transformers import DPRQuestionEncoderTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel
 from pathlib import Path
 import base64
 import io
@@ -20,13 +20,11 @@ DATAFRAME_FILE_BSBS = 'basecamp.csv'
 selectbox_selections = {
     'Retrieve - Rerank (with fine-tuned cross-encoder)': 1,
     'Dense Passage Retrieval':2,
-    'Retrieve - Reranking with DPR':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',
-    'Retrieve - Reranking with DPR': 'Retrieve-rerank-DPR.png',
     'Retrieve - Rerank': 'retrieve-rerank.png'
 }
 
@@ -63,7 +61,7 @@ class CPU_Unpickler(pickle.Unpickler):
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def load_paragraphs(path):
     with open(path, "rb") as fIn:
-        cache_data = CPU_Unpickler(fIn).load()
+        cache_data = pickle.load(fIn)
         corpus_sentences = cache_data['contexes']
         corpus_embeddings = cache_data['embeddings']
 
@@ -84,45 +82,25 @@ def dot_product(question_output, context_output):
     result = torch.dot(mat1, mat2)
     return result
 
-def retrieve_rerank_DPR(question):
-    hits = retrieve(question)
-    return rerank_with_DPR(hits, question)
-
-def DPR_reranking(question, selected_contexes, selected_embeddings):
-    scores = []
-    tokenized_question = question_tokenizer(question, padding=True, truncation=True, return_tensors="pt",
-                                            add_special_tokens=True)
-    question_output = dpr_trained.model.question_model(**tokenized_question)
-    question_output = question_output['pooler_output']
-    for context_embedding in selected_embeddings:
-        score = dot_product(question_output, context_embedding)
-        scores.append(score.detach().cpu())
-
-    scores_index = sorted(range(len(scores)), key=lambda x: scores[x], reverse=True)
-    contexes_list = []
-    scores_final = []
-    for i, idx in enumerate(scores_index[:5]):
-        scores_final.append(scores[idx])
-        contexes_list.append(selected_contexes[idx])
-    return scores_final, contexes_list
-
 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 retrieve_rerank_DPR(question)
     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",
-                                            add_special_tokens=True)
+    tokenized_question = question_tokenizer(question, padding=True, truncation=True, return_tensors="pt")
     question_embedding = dpr_trained.model.question_model(**tokenized_question)
-    question_embedding = question_embedding['pooler_output']
+    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)
@@ -145,35 +123,13 @@ def retrieve(question):
     hits = hits[0]
     return hits
 
-def retrieve_with_dpr_embeddings(question):
-    # Semantic Search (Retrieve)
-    question_tokens = question_tokenizer(question, padding=True, truncation=True, return_tensors="pt",
-                                            add_special_tokens=True)
-
-    question_embedding = dpr_trained.model.question_model(**question_tokens)['pooler_output']
-    question_embedding = torch.squeeze(question_embedding, dim=0)
-    corpus_embeddings = torch.stack(dpr_context_embeddings)
-    corpus_embeddings = torch.squeeze(corpus_embeddings, dim=1)
-    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=100, score_function=util.dot_score)
-    if len(hits) == 0:
-        return []
-    hits = hits[0]
-    return hits, question_embedding
-
-def rerank_with_DPR(hits, question_embedding):
-    # Rerank - score all retrieved passages with cross-encoder
-    selected_contexes = [dpr_contexes[hit['corpus_id']] for hit in hits]
-    selected_embeddings = [dpr_context_embeddings[hit['corpus_id']] for hit in hits]
-    top_5_scores, top_5_contexes = DPR_reranking(question_embedding, selected_contexes, selected_embeddings)
-    return top_5_contexes, top_5_scores
-
 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][1]
+        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)
@@ -263,19 +219,22 @@ def qa_main_widgetsv2():
 
 @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 = DPRQuestionEncoderTokenizer.from_pretrained('facebook/dpr-question_encoder-single-nq-base')
+    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-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)
+

st-context-embeddings.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bd65fe793062375df1efd50218e9a7c35253fe06a24e5527de7855671a4f958c
-size 468299
+oid sha256:79e231244e12074d5e22f46cf3da70f4f1dd43cc6e82f36959d2c6817f2e2bf2
+size 441107