Fix loading of tokenizers in DPR (#2755)

deepset-ai · Jul 4, 2022 · dc48c44 · dc48c44
1 parent 2a8b129
commit dc48c44
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 15 deletions.
diff --git a/haystack/modeling/model/tokenization.py b/haystack/modeling/model/tokenization.py
@@ -308,14 +308,7 @@ def _infer_tokenizer_class_from_string(pretrained_model_name_or_path):
         elif "dpr-ctx_encoder" in pretrained_model_name_or_path.lower():
             tokenizer_class = "DPRContextEncoderTokenizer"
         else:
-            raise ValueError(
-                f"Could not infer tokenizer_class from model config or "
-                f"name '{pretrained_model_name_or_path}'. Set arg `tokenizer_class` "
-                f"in Tokenizer.load() to one of: AlbertTokenizer, XLMRobertaTokenizer, "
-                f"RobertaTokenizer, DistilBertTokenizer, BertTokenizer, XLNetTokenizer, "
-                f"CamembertTokenizer, ElectraTokenizer, DPRQuestionEncoderTokenizer,"
-                f"DPRContextEncoderTokenizer."
-            )
+            tokenizer_class = "AutoTokenizer"
 
         return tokenizer_class
 

diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py
@@ -152,7 +152,7 @@ def __init__(
             )
 
         self.infer_tokenizer_classes = infer_tokenizer_classes
-        tokenizers_default_classes = {"query": "AutoTokenizer", "passage": "AutoTokenizer"}
+        tokenizers_default_classes = {"query": "DPRQuestionEncoderTokenizer", "passage": "DPRContextEncoderTokenizer"}
         if self.infer_tokenizer_classes:
             tokenizers_default_classes["query"] = None  # type: ignore
             tokenizers_default_classes["passage"] = None  # type: ignore

diff --git a/test/nodes/test_retriever.py b/test/nodes/test_retriever.py
@@ -17,7 +17,7 @@
 from haystack.document_stores import MilvusDocumentStore
 from haystack.nodes.retriever.dense import DensePassageRetriever, EmbeddingRetriever, TableTextRetriever
 from haystack.nodes.retriever.sparse import BM25Retriever, FilterRetriever, TfidfRetriever
-from transformers import DPRContextEncoderTokenizerFast, DPRQuestionEncoderTokenizerFast, PreTrainedTokenizerFast
+from transformers import DPRContextEncoderTokenizerFast, DPRQuestionEncoderTokenizerFast
 
 from ..conftest import SAMPLES_PATH
 
@@ -277,8 +277,8 @@ def sum_params(model):
     assert loaded_retriever.processor.max_seq_len_query == 64
 
     # Tokenizer
-    assert isinstance(loaded_retriever.passage_tokenizer, PreTrainedTokenizerFast)
-    assert isinstance(loaded_retriever.query_tokenizer, PreTrainedTokenizerFast)
+    assert isinstance(loaded_retriever.passage_tokenizer, DPRContextEncoderTokenizerFast)
+    assert isinstance(loaded_retriever.query_tokenizer, DPRQuestionEncoderTokenizerFast)
     assert loaded_retriever.passage_tokenizer.do_lower_case == True
     assert loaded_retriever.query_tokenizer.do_lower_case == True
     assert loaded_retriever.passage_tokenizer.vocab_size == 30522
@@ -320,9 +320,9 @@ def sum_params(model):
     assert loaded_retriever.processor.max_seq_len_query == 64
 
     # Tokenizer
-    assert isinstance(loaded_retriever.passage_tokenizer, PreTrainedTokenizerFast)
-    assert isinstance(loaded_retriever.table_tokenizer, PreTrainedTokenizerFast)
-    assert isinstance(loaded_retriever.query_tokenizer, PreTrainedTokenizerFast)
+    assert isinstance(loaded_retriever.passage_tokenizer, DPRContextEncoderTokenizerFast)
+    assert isinstance(loaded_retriever.table_tokenizer, DPRContextEncoderTokenizerFast)
+    assert isinstance(loaded_retriever.query_tokenizer, DPRQuestionEncoderTokenizerFast)
     assert loaded_retriever.passage_tokenizer.do_lower_case == True
     assert loaded_retriever.table_tokenizer.do_lower_case == True
     assert loaded_retriever.query_tokenizer.do_lower_case == True