Upload demo_dpr.py

demo_dpr.py

@@ -10,8 +10,7 @@ import torch
 from transformers import DPRQuestionEncoderTokenizer, AutoModel
 from pathlib import Path
 import base64
-import regex
-import tokenizers
+import io
 
 st.set_page_config(layout="wide")
 
@@ -55,12 +54,16 @@ def retrieve_rerank(question):
         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)
+        cache_data = CPU_Unpickler(fIn).load()
         corpus_sentences = cache_data['contexes']
         corpus_embeddings = cache_data['embeddings']
 
@@ -276,6 +279,3 @@ dpr_trained, bi_encoder, cross_encoder, trained_cross_encoder, question_tokenize
 
 qa_main_widgetsv2()
 
-#if __name__ == '__main__':
-#    search_pipeline('Life insurance is paid by insurance companies that pay for what?', 1)
-