Add DPR Bert model

dpr_bert/__init__.py

@@ -0,0 +1,2 @@
+from .config import DprConfig
+from .model import DprModel

dpr_bert/config.py

@@ -0,0 +1,29 @@
+from transformers import PretrainedConfig
+
+
+# From https://huggingface.co/klue/bert-base
+_default_config = {
+    "architectures": ["BertForMaskedLM"],
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 512,
+    "model_type": "bert",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12,
+    "pad_token_id": 0,
+    "type_vocab_size": 2,
+    "vocab_size": 32000
+}
+
+class DprConfig(PretrainedConfig):
+    model_type = "dpr"
+
+    def __init__(self, qst_config=_default_config, ctx_config=_default_config, **kwargs):
+        self.qst_config = qst_config
+        self.ctx_config = ctx_config
+        super().__init__(**kwargs)

dpr_bert/model.py

@@ -0,0 +1,20 @@
+import torch
+from transformers import PreTrainedModel, BertConfig, BertModel
+
+from .config import DprConfig
+
+
+class DprModel(PreTrainedModel):
+    config_class = DprConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        qst_config = BertConfig(**config.qst_config)
+        ctx_config = BertConfig(**config.ctx_config)
+        self.qst_encoder = BertModel(qst_config)
+        self.ctx_encoder = BertModel(ctx_config)
+
+    def forward(self, qst_input_ids, qst_attention_mask, ctx_input_ids, ctx_attention_mask):
+        qst_outputs = self.qst_encoder(input_ids=qst_input_ids, attention_mask=qst_attention_mask)
+        ctx_outputs = self.ctx_encoder(input_ids=ctx_input_ids, attention_mask=ctx_attention_mask)
+        return torch.einsum("ih,jh->ij", qst_outputs.pooler_output, ctx_outputs.pooler_output)