Relphormer baseline

pretrain/models/__init__.py

@@ -0,0 +1,6 @@
+
+
+from transformers import BartForConditionalGeneration, T5ForConditionalGeneration, GPT2LMHeadModel
+
+from .model import *
+

pretrain/models/model.py

@@ -0,0 +1,10 @@
+from transformers.models.bert.modeling_bert import BertForMaskedLM
+
+
+
+
+class BertKGC(BertForMaskedLM):
+    @staticmethod
+    def add_to_argparse(parser):
+        parser.add_argument("--pretrain", type=int, default=0, help="")
+        return parser