Thesis/models/model.py

7 lines
302 B
Python
Raw Normal View History

2022-12-26 04:54:46 +00:00
# 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