Relphormer baseline

models/__init__.py

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

models/model.py

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