DPR - enable loading of custom BERT models (#746)

* Add loading of other bert type models

* Fix loading and saving, add test

Author: Timoeller

farm/data_handler/processor.py

@@ -2868,6 +2868,9 @@ def file_to_dicts(self, file: str) -> [dict]:
             ...]}
         """
         dicts = read_dpr_json(file, max_samples=self.max_samples, num_hard_negatives=self.num_hard_negatives, num_positives=self.num_positives, shuffle_negatives=self.shuffle_negatives, shuffle_positives=self.shuffle_positives)
+
+        # shuffle dicts to make sure that similar positive passages do not end up in one batch
+        dicts = random.sample(dicts, len(dicts))
         return dicts
 
     def dataset_from_dicts(self, dicts, indices=None, return_baskets = False):

farm/modeling/language_model.py

@@ -1447,16 +1447,23 @@ def load(cls, pretrained_model_name_or_path, language=None, **kwargs):
             dpr_question_encoder.model = transformers.DPRQuestionEncoder.from_pretrained(farm_lm_model, config=dpr_config, **kwargs)
             dpr_question_encoder.language = dpr_question_encoder.model.config.language
         else:
-            model_type = AutoConfig.from_pretrained(pretrained_model_name_or_path).model_type
-            if model_type == "dpr":
+            original_model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+            if original_model_config.model_type == "dpr":
                 # "pretrained dpr model": load existing pretrained DPRQuestionEncoder model
                 dpr_question_encoder.model = transformers.DPRQuestionEncoder.from_pretrained(
                     str(pretrained_model_name_or_path), **kwargs)
             else:
                 # "from scratch": load weights from different architecture (e.g. bert) into DPRQuestionEncoder
-                dpr_question_encoder.model = transformers.DPRQuestionEncoder(config=transformers.DPRConfig(**kwargs))
+                # but keep config values from original architecture
+                # TODO test for architectures other than BERT, e.g. Electra
+                if original_model_config.model_type != "bert":
+                    logger.warning(f"Using a model of type '{original_model_config.model_type}' which might be incompatible with DPR encoders."
+                                   f"Bert based encoders are supported that need input_ids,token_type_ids,attention_mask as input tensors.")
+                original_config_dict = vars(original_model_config)
+                original_config_dict.update(kwargs)
+                dpr_question_encoder.model = transformers.DPRQuestionEncoder(config=transformers.DPRConfig(**original_config_dict))
                 dpr_question_encoder.model.base_model.bert_model = AutoModel.from_pretrained(
-                    str(pretrained_model_name_or_path), **kwargs)
+                    str(pretrained_model_name_or_path), **original_config_dict)
             dpr_question_encoder.language = cls._get_or_infer_language_from_name(language, pretrained_model_name_or_path)
 
         return dpr_question_encoder
@@ -1541,16 +1548,25 @@ def load(cls, pretrained_model_name_or_path, language=None, **kwargs):
             dpr_context_encoder.language = dpr_context_encoder.model.config.language
         else:
             # Pytorch-transformer Style
-            model_type = AutoConfig.from_pretrained(pretrained_model_name_or_path).model_type
-            if model_type == "dpr":
+            original_model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+            if original_model_config.model_type == "dpr":
                 # "pretrained dpr model": load existing pretrained DPRContextEncoder model
                 dpr_context_encoder.model = transformers.DPRContextEncoder.from_pretrained(
                     str(pretrained_model_name_or_path), **kwargs)
             else:
                 # "from scratch": load weights from different architecture (e.g. bert) into DPRContextEncoder
-                dpr_context_encoder.model = transformers.DPRContextEncoder(config=transformers.DPRConfig(**kwargs))
+                # but keep config values from original architecture
+                # TODO test for architectures other than BERT, e.g. Electra
+                if original_model_config.model_type != "bert":
+                    logger.warning(
+                        f"Using a model of type '{original_model_config.model_type}' which might be incompatible with DPR encoders."
+                        f"Bert based encoders are supported that need input_ids,token_type_ids,attention_mask as input tensors.")
+                original_config_dict = vars(original_model_config)
+                original_config_dict.update(kwargs)
+                dpr_context_encoder.model = transformers.DPRContextEncoder(
+                    config=transformers.DPRConfig(**original_config_dict))
                 dpr_context_encoder.model.base_model.bert_model = AutoModel.from_pretrained(
-                    str(pretrained_model_name_or_path), **kwargs)
+                    str(pretrained_model_name_or_path), **original_config_dict)
             dpr_context_encoder.language = cls._get_or_infer_language_from_name(language, pretrained_model_name_or_path)
 
         return dpr_context_encoder

farm/modeling/prediction_head.py

@@ -1709,7 +1709,7 @@ def _embeddings_to_scores(self, query_vectors:torch.Tensor, passage_vectors:torc
         softmax_scores = nn.functional.log_softmax(scores, dim=1)
         return softmax_scores
 
-    def logits_to_loss(self, logits: Tuple[torch.Tensor, torch.Tensor], **kwargs):
+    def logits_to_loss(self, logits: Tuple[torch.Tensor, torch.Tensor], label_ids, **kwargs):
         """
         Computes the loss (Default: NLLLoss) by applying a similarity function (Default: dot product) to the input
         tuple of (query_vectors, passage_vectors) and afterwards applying the loss function on similarity scores.
@@ -1729,8 +1729,7 @@ def logits_to_loss(self, logits: Tuple[torch.Tensor, torch.Tensor], **kwargs):
         query_vectors, passage_vectors = logits
 
         # Prepare Labels
-        lm_label_ids = kwargs.get(self.label_tensor_name)
-        positive_idx_per_question = torch.nonzero((lm_label_ids.view(-1) == 1), as_tuple=False)
+        positive_idx_per_question = torch.nonzero((label_ids.view(-1) == 1), as_tuple=False)
 
         # Gather global embeddings from all distributed nodes (DDP)
         if rank != -1:
@@ -1788,13 +1787,12 @@ def logits_to_preds(self, logits: Tuple[torch.Tensor, torch.Tensor], **kwargs):
         _, sorted_scores = torch.sort(softmax_scores, dim=1, descending=True)
         return sorted_scores
 
-    def prepare_labels(self, **kwargs):
+    def prepare_labels(self, label_ids, **kwargs):
         """
         Returns a tensor with passage labels(0:hard_negative/1:positive) for each query
 
         :return: passage labels(0:hard_negative/1:positive) for each query
         """
-        label_ids = kwargs.get(self.label_tensor_name)
         labels = torch.zeros(label_ids.size(0), label_ids.numel())
 
         positive_indices = torch.nonzero(label_ids.view(-1) == 1, as_tuple=False)

test/test_dpr.py

@@ -2,6 +2,7 @@
 import torch
 import logging
 import numpy as np
+from pathlib import Path
 from farm.data_handler.processor import TextSimilarityProcessor
 from farm.data_handler.data_silo import DataSilo
 from farm.train import Trainer
@@ -402,7 +403,7 @@ def test_dpr_context_only():
     assert tensor_names == ['passage_input_ids', 'passage_segment_ids', 'passage_attention_mask', 'label_ids']
 
 
-def test_dpr_save_load():
+def test_dpr_processor_save_load():
     d = {'query': 'big little lies season 2 how many episodes',
          'passages': [
                          {'title': 'Big Little Lies (TV series)',
@@ -443,6 +444,7 @@ def test_dpr_save_load():
     assert np.array_equal(dataset.tensors[0],dataset2.tensors[0])
 
 
+
 def test_dpr_training():
     batch_size = 1
     n_epochs = 1
@@ -524,6 +526,36 @@ def test_dpr_training():
 
     trainer.train()
 
+    ######## save and load model again
+    save_dir = Path("testsave/dpr-model")
+    model.save(save_dir)
+    del model
+
+    model2 = BiAdaptiveModel.load(save_dir, device=device)
+    model2, optimizer2, lr_schedule = initialize_optimizer(
+        model=model2,
+        learning_rate=1e-5,
+        optimizer_opts={"name": "TransformersAdamW", "correct_bias": True, "weight_decay": 0.0, \
+                        "eps": 1e-08},
+        schedule_opts={"name": "LinearWarmup", "num_warmup_steps": 100},
+        n_batches=len(data_silo.loaders["train"]),
+        n_epochs=n_epochs,
+        grad_acc_steps=1,
+        device=device,
+        distributed=distributed
+    )
+    trainer2 = Trainer(
+        model=model2,
+        optimizer=optimizer,
+        data_silo=data_silo,
+        epochs=n_epochs,
+        n_gpu=n_gpu,
+        lr_schedule=lr_schedule,
+        evaluate_every=evaluate_every,
+        device=device,
+    )
+
+    trainer2.train()
 
 
 if __name__=="__main__":