from .bert import BertForSequenceClassification
from .bert.modeling_bert import BertGetLabelWord, BertUseLabelWord, BertDecouple, BertForMaskedLM

from .roberta import RobertaForSequenceClassification 
from .roberta.modeling_roberta import RobertaUseLabelWord

from transformers import RobertaForMaskedLM, BertForMaskedLM

class RobertaForPrompt(RobertaForMaskedLM):
    @staticmethod
    def add_to_argparse(parser):
        parser.add_argument("--use_prompt", type=bool, default=True, help="Whether to use prompt in the dataset.")
        parser.add_argument("--init_answer_words", type=int, default=1, )
        parser.add_argument("--init_type_words", type=int, default=1, )
        parser.add_argument("--init_answer_words_by_one_token", type=int, default=0, )
        parser.add_argument("--use_template_words", type=int, default=1, )
        return parser


class BERTForPrompt(BertForMaskedLM):
    @staticmethod
    def add_to_argparse(parser):
        parser.add_argument("--use_prompt", type=bool, default=True, help="Whether to use prompt in the dataset.")
        parser.add_argument("--init_answer_words", type=int, default=1, )
        parser.add_argument("--init_type_words", type=int, default=1, )
        parser.add_argument("--init_answer_words_by_one_token", type=int, default=0, )
        parser.add_argument("--use_template_words", type=int, default=1, )
        return parser