Upload 4 files

models/sequence_classification/causal_prompt_cls.py

@@ -0,0 +1,199 @@
+import sys
+import os
+import torch
+import torch.nn as nn
+import transformers
+import torch.nn.functional as F
+import numpy as np
+from typing import Optional, Tuple, Union
+from torch.nn import CrossEntropyLoss
+from transformers import AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from transformers.models.gpt2.modeling_gpt2 import GPT2PreTrainedModel, GPT2Model, GPT2LMHeadModel
+from transformers.modeling_outputs import ModelOutput
+from tools.runner_utils.log_util import logging
+from tools.model_utils.parameter_freeze import ParameterFreeze
+
+logger = logging.getLogger(__name__)
+freezer = ParameterFreeze()
+
+
+"""
+Function: Use Causal LM to prompt for cls
+Notes:
+- For classification, the model only calculate the loss at the position of label, the other position is set as -100
+- During inference, generate result at the last position.
+"""
+class PromptGPT2ForSequenceClassification(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        if self.config.use_freezing:
+            self.transformer = freezer.freeze_lm(self.transformer)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.transformer.device)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # print("shift_labels=", shift_labels)
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
+        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+            for layer_past in past
+        )
+
+
+
+
+# if __name__ == "__main__":
+#     from transformers import GPT2Tokenizer
+#     tokenizer = GPT2Tokenizer.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
+#     model = GPT2ForInContextLearning.from_pretrained("/Users/wangjianing/Desktop/开源代码与数据模型/模型/gpt2")
+
+#     # In-Context Learning for classification
+#     # input_text = "The capital city of China is Beijing. \n\n The capital city of Japan is Tokyo. \n\n The capital city of America is"
+#     input_text = "What are follows emotions? \n\n Input: The book is very nice.\n Output: Great. \n\n Input: I never eat chocolate!\n Output:"
+#     # input_text = "This film is wonderful.\n Great."
+#     tokenizer.pad_token = tokenizer.eos_token
+#     inputs = tokenizer(input_text, return_tensors="pt")
+#     input_len = inputs["input_ids"].shape[-1]
+#     gen_output = model.generate(**inputs, max_length=input_len + 10)
+#     gen_result = tokenizer.decode(gen_output[0])
+#     print("classification result:\n", gen_result)
+
+#     # In-Context Learning for generation
+#     input_text = "Please tell me what is the transformer? "
+#     # input_text = "This film is wonderful.\n Great."
+#     tokenizer.pad_token = tokenizer.eos_token
+#     inputs = tokenizer(input_text, return_tensors="pt")
+#     input_len = inputs["input_ids"].shape[-1]
+#     gen_output = model.generate(**inputs, max_length=input_len + 60)
+#     gen_result = tokenizer.decode(gen_output[0])
+#     print("generation result:\n", gen_result)

models/sequence_classification/classification.py

@@ -0,0 +1,175 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2021/8/19 10:54 上午
+# @Author  : JianingWang
+# @File    : classification.py
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
+from transformers import RobertaModel
+from transformers.activations import ACT2FN
+from transformers.models.electra import ElectraModel
+from transformers.models.roformer import RoFormerModel
+from transformers.models.albert import AlbertModel
+from transformers.models.bert import BertModel, BertPreTrainedModel
+from transformers.models.deberta_v2 import DebertaV2Model, DebertaV2PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers.models.roberta import RobertaPreTrainedModel
+from transformers.models.bert.modeling_bert import BertForSequenceClassification
+from transformers.models.megatron_bert import MegatronBertPreTrainedModel, MegatronBertModel
+
+PRETRAINED_MODEL_MAP = {
+    "bert": BertPreTrainedModel,
+    "deberta-v2": DebertaV2PreTrainedModel,
+    "roberta": RobertaPreTrainedModel,
+    "erlangshen": MegatronBertPreTrainedModel
+}
+
+
+class BertPooler(nn.Module):
+    def __init__(self, hidden_size, hidden_act, hidden_dropout_prob):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        # self.activation = nn.Tanh()
+        self.activation = ACT2FN[hidden_act]
+        # self.dropout = nn.Dropout(hidden_dropout_prob)
+
+    def forward(self, features):
+        x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
+        # x = self.dropout(x)
+        x = self.dense(x)
+        x = self.activation(x)
+        return x
+
+
+def build_cls_model(config):
+    BaseClass = PRETRAINED_MODEL_MAP[config.model_type]
+
+    class BertForClassification(BaseClass):
+
+        def __init__(self, config):
+            super().__init__(config)
+            self.num_labels = config.num_labels
+            self.config = config
+            self.model_type = config.model_type
+            self.problem_type = config.problem_type
+
+            if self.model_type == "bert":
+                self.bert = BertModel(config)
+            elif self.model_type == "albert":
+                self.albert = AlbertModel(config)
+            # elif self.model_type == "chinesebert":
+            #     self.bert = ChineseBertModel(config)
+            elif self.model_type == "roformer":
+                self.roformer = RoFormerModel(config)
+            elif self.model_type == "electra":
+                self.electra = ElectraModel(config)
+            elif self.model_type == "deberta-v2":
+                self.deberta = DebertaV2Model(config)
+            elif self.model_type == "roberta":
+                self.roberta = RobertaModel(config)
+            elif self.model_type == "erlangshen":
+                self.bert = MegatronBertModel(config)
+            self.pooler = BertPooler(config.hidden_size, config.hidden_act, config.hidden_dropout_prob)
+            if hasattr(config, "cls_dropout_rate"):
+                cls_dropout_rate = config.cls_dropout_rate
+            else:
+                cls_dropout_rate = config.hidden_dropout_prob
+            self.dropout = nn.Dropout(cls_dropout_rate)
+            add_feature_dims = config.additional_feature_dims if hasattr(config, "additional_feature_dims") else 0
+            # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+            cls_hidden = config.hidden_size + add_feature_dims
+            if hasattr(config, "is_relation_task"):
+                cls_hidden = config.hidden_size * 2
+            self.classifier = nn.Linear(cls_hidden, config.num_labels)
+
+            self.init_weights()
+
+        def forward(
+                self,
+                input_ids=None,
+                attention_mask=None,
+                token_type_ids=None,
+                position_ids=None,
+                head_mask=None,
+                inputs_embeds=None,
+                labels=None,
+                output_attentions=None,
+                output_hidden_states=None,
+                return_dict=None,
+                pseudo_label=None,
+                pinyin_ids=None,
+                additional_features=None
+        ):
+            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+            logits, outputs = None, None
+            inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, "position_ids": position_ids,
+                      "head_mask": head_mask, "inputs_embeds": inputs_embeds, "output_attentions": output_attentions,
+                      "output_hidden_states": output_hidden_states, "return_dict": return_dict, "pinyin_ids": pinyin_ids}
+            inputs = {k: v for k, v in inputs.items() if v is not None}
+            if self.model_type == "chinesebert":
+                outputs = self.bert(**inputs)
+            elif self.model_type == "bert":
+                outputs = self.bert(**inputs)
+            elif self.model_type == "albert":
+                outputs = self.albert(**inputs)
+            elif self.model_type == "electra":
+                outputs = self.electra(**inputs)
+            elif self.model_type == "roformer":
+                outputs = self.roformer(**inputs)
+            elif self.model_type == "deberta-v2":
+                outputs = self.deberta(**inputs)
+            elif self.model_type == "roberta":
+                outputs = self.roberta(**inputs)
+            elif self.model_type == "erlangshen":
+                outputs = self.bert(**inputs)
+
+            if hasattr(self.config, "is_relation_task"):
+                w = torch.logical_and(input_ids >= min(self.config.start_token_ids), input_ids <= max(self.config.start_token_ids))
+                start_index = w.nonzero()[:, 1].view(-1, 2)
+                # <start_entity> + <end_entity> 进分类
+                pooler_output = torch.cat([torch.cat([x[y[0], :], x[y[1], :]]).unsqueeze(0) for x, y in zip(outputs.last_hidden_state, start_index)])
+                # [CLS] + <start_entity> + <end_entity> 进分类
+                # pooler_output = torch.cat([torch.cat([z, x[y[0], :], x[y[1], :]]).unsqueeze(0) for x, y, z in zip(outputs.last_hidden_state, start_index, outputs.last_hidden_state[:, 0])])
+
+            elif "pooler_output" in outputs:
+                pooler_output = outputs.pooler_output
+            else:
+                pooler_output = self.pooler(outputs[0])
+            pooler_output = self.dropout(pooler_output)
+            # pooler_output = self.LayerNorm(pooler_output)
+            if additional_features is not None:
+                pooler_output = torch.cat((pooler_output, additional_features), dim=1)
+            logits = self.classifier(pooler_output)
+
+            loss = None
+            if labels is not None:
+                if self.problem_type == "regression":
+                    loss_fct = MSELoss()
+                    if self.num_labels == 1:
+                        loss = loss_fct(logits.squeeze(), labels.squeeze())
+                    else:
+                        loss = loss_fct(logits, labels)
+                elif self.problem_type == "multi_label_classification":
+                    loss_fct = BCEWithLogitsLoss()
+                    loss = loss_fct(logits.view(-1, self.num_labels), labels.float().view(-1, self.num_labels))
+                # elif self.problem_type in ["single_label_classification"] or hasattr(self.config, "is_relation_task"):
+                else:
+                    # loss_fct = FocalLoss()
+                    loss_fct = CrossEntropyLoss()
+                    # 伪标签
+                    if pseudo_label is not None:
+                        train_logits, pseudo_logits = logits[pseudo_label > 0.9], logits[pseudo_label < 0.1]
+                        train_labels, pseudo_labels = labels[pseudo_label > 0.9], labels[pseudo_label < 0.1]
+                        train_loss = loss_fct(train_logits.view(-1, self.num_labels), train_labels.view(-1)) if train_labels.nelement() else 0
+                        pseudo_loss = loss_fct(pseudo_logits.view(-1, self.num_labels), pseudo_labels.view(-1)) if pseudo_labels.nelement() else 0
+                        loss = 0.9 * train_loss + 0.1 * pseudo_loss
+                    else:
+                        loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            return SequenceClassifierOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )
+
+    return BertForClassification

models/sequence_classification/head_cls.py

@@ -0,0 +1,1284 @@
+"""
+Head Tuning with Prefix / Adapter
+"""
+from typing import Optional, List, Union, Tuple
+import torch
+from torch._C import NoopLogger
+import torch.nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
+
+from transformers import BertModel, BertPreTrainedModel
+from transformers import RobertaModel, RobertaPreTrainedModel
+from transformers.models.deberta.modeling_deberta import DebertaModel, DebertaPreTrainedModel, ContextPooler, StableDropout
+from transformers.models.gpt2.modeling_gpt2 import GPT2Model, GPT2PreTrainedModel
+from transformers.models.bart.modeling_bart import BartPretrainedModel, BartClassificationHead, BartModel
+from transformers.models.roberta.modeling_roberta import RobertaClassificationHead
+from transformers.models.bart.configuration_bart import BartConfig
+from transformers.modeling_outputs import SequenceClassifierOutput, Seq2SeqSequenceClassifierOutput, SequenceClassifierOutputWithPast
+
+from models.basic_modules.prefix_encoder import PrefixEncoder
+
+from models.basic_modules.adapter import BertAdaModel, RobertaAdaModel, init_adapter
+from tools.model_utils.parameter_freeze import ParameterFreeze
+
+from tools.runner_utils.log_util import logging
+logger = logging.getLogger(__name__)
+
+
+freezer = ParameterFreeze()
+
+## ======== BERT ========
+
+#  Vanilla Fine-tuning For BERT
+class BertForSequenceClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = BertModel(config)
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
+            config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # print("input_ids.shape=", input_ids.shape) # e.g., [8, 128]
+        # print("attention_mask.shape=", attention_mask.shape) # e.g., [8, 128]
+        # print("token_type_ids.shape=", token_type_ids.shape) # e.g., [8, 128]
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+#  Prefix-tuning For BERT
+class BertPrefixForSequenceClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.bert = BertModel(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        # for param in self.bert.parameters():
+        #     param.requires_grad = False
+
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+
+        self.pre_seq_len = config.pre_seq_len
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+
+        self.prefix_encoder = PrefixEncoder(config)
+
+        bert_param = 0
+        for name, param in self.bert.named_parameters():
+            bert_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - bert_param
+        print("total param is {}".format(total_param)) # 9860105
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.bert.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        # bsz, seqlen, _ = past_key_values.shape
+        past_key_values = past_key_values.view(
+            batch_size,
+            self.pre_seq_len,
+            self.n_layer * 2,
+            self.n_head,
+            self.n_embd
+        )
+        past_key_values = self.dropout(past_key_values)
+        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+        return past_key_values
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # print("input_ids.shape=", input_ids.shape) # e.g., [8, 128]
+        # print("attention_mask.shape=", attention_mask.shape) # e.g., [8, 128]
+        # print("token_type_ids.shape=", token_type_ids.shape) # e.g., [8, 128]
+
+        batch_size = input_ids.shape[0]
+        past_key_values = self.get_prompt(batch_size=batch_size)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.bert.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        if position_ids is None:
+            position_ids = torch.tensor([i for i in range(input_ids.shape[-1])]).expand(batch_size, -1).to(self.bert.device)
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            past_key_values=past_key_values,
+        )
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+#  Prompt-tuning For BERT
+class BertPtuningForSequenceClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.bert = BertModel(config)
+        self.embeddings = self.bert.embeddings
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        # for param in self.bert.parameters():
+        #     param.requires_grad = False
+
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+
+        self.pre_seq_len = config.pre_seq_len
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+        self.prefix_encoder = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.bert.device)
+        prompts = self.prefix_encoder(prefix_tokens)
+        return prompts
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = input_ids.shape[0]
+        raw_embedding = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+        )
+        prompts = self.get_prompt(batch_size=batch_size)
+        inputs_embeds = torch.cat((prompts, raw_embedding), dim=1)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.bert.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.bert(
+            # input_ids,
+            attention_mask=attention_mask,
+            # token_type_ids=token_type_ids,
+            # position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            # past_key_values=past_key_values,
+        )
+
+        # pooled_output = outputs[1]
+        sequence_output = outputs[0]
+        sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        first_token_tensor = sequence_output[:, 0]
+        pooled_output = self.bert.pooler.dense(first_token_tensor)
+        pooled_output = self.bert.pooler.activation(pooled_output)
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+#  Adapter-tuning For BERT
+class BertAdapterForSequenceClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.bert = BertAdaModel(config)
+        self.embeddings = self.bert.embeddings
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        # for param in self.bert.parameters():
+        #     param.requires_grad = False
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm_component(self.bert, "adapter")
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm_component(self.bert, "adapter")
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = input_ids.shape[0]
+        inputs_embeds = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+        )
+        outputs = self.bert(
+            # input_ids,
+            attention_mask=attention_mask,
+            # token_type_ids=token_type_ids,
+            # position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            # past_key_values=past_key_values,
+        )
+
+        # pooled_output = outputs[1]
+        sequence_output = outputs[0]
+        # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        first_token_tensor = sequence_output[:, 0]
+        pooled_output = self.bert.pooler.dense(first_token_tensor)
+        pooled_output = self.bert.pooler.activation(pooled_output)
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+
+# ========= RoBERTa =========
+
+#  Vanilla Fine-tuning For RoBERTa
+class RobertaForSequenceClassification(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.roberta = RobertaModel(config)
+        if self.config.use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        # self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = RobertaClassificationHead(config)
+        self.init_weights()
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        else:
+            self.roberta = freezer.unfreeze_lm(self.roberta)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
+            config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.roberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+#  Prefix-tuning For RoBERTa
+class RobertaPrefixForSequenceClassification(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.roberta = RobertaModel(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        # self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = RobertaClassificationHead(config)
+        self.init_weights()
+
+        for param in self.roberta.parameters():
+            param.requires_grad = False
+
+        self.pre_seq_len = config.pre_seq_len
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+        self.prefix_encoder = PrefixEncoder(config)
+
+        bert_param = 0
+        for name, param in self.roberta.named_parameters():
+            bert_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - bert_param
+        print("total param is {}".format(total_param)) # 9860105
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        else:
+            self.roberta = freezer.unfreeze_lm(self.roberta)
+
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        # print("prefix_tokens.shape=", prefix_tokens.shape)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        # print("past_key_values[0].shape=", past_key_values[0].shape)
+        past_key_values = past_key_values.view(
+            batch_size,
+            self.pre_seq_len,
+            self.n_layer * 2,
+            self.n_head,
+            self.n_embd
+        )
+        # print("past_key_values[0].shape=", past_key_values[0].shape)
+        past_key_values = self.dropout(past_key_values)
+        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+        # print("past_key_values[0].shape=", past_key_values[0].shape)
+        return past_key_values
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = input_ids.shape[0]
+        past_key_values = self.get_prompt(batch_size=batch_size)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.roberta.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        if position_ids is None:
+            position_ids = torch.tensor([i for i in range(input_ids.shape[-1])]).expand(batch_size, -1).to(self.roberta.device)
+
+        outputs = self.roberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            past_key_values=past_key_values,
+        )
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            labels = (labels < 0).long().to(labels.device) + labels
+
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+#  Prompt-tuning For RoBERTa
+class RobertaPtuningForSequenceClassification(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.roberta = RobertaModel(config)
+        self.embeddings = self.roberta.embeddings
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        # self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = RobertaClassificationHead(config)
+
+        # for param in self.roberta.parameters():
+        #     param.requires_grad = False
+
+        if self.config.use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+
+        self.pre_seq_len = config.pre_seq_len
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+        self.prefix_encoder = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        else:
+            self.roberta = freezer.unfreeze_lm(self.roberta)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        prompts = self.prefix_encoder(prefix_tokens)
+        return prompts
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = input_ids.shape[0]
+        raw_embedding = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+        )
+        prompts = self.get_prompt(batch_size=batch_size)
+        inputs_embeds = torch.cat((prompts, raw_embedding), dim=1)
+        # print(input_embeddings.shape)
+        # exit()
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.roberta.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.roberta(
+            # input_ids,
+            attention_mask=attention_mask,
+            # token_type_ids=token_type_ids,
+            # position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            # past_key_values=past_key_values,
+        )
+
+        # pooled_output = outputs[1]
+        sequence_output = outputs[0]
+        sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        first_token_tensor = sequence_output[:, 0]
+        pooled_output = self.roberta.pooler.dense(first_token_tensor)
+        pooled_output = self.roberta.pooler.activation(pooled_output)
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+#  Adapter-tuning For RoBERTa
+class RobertaAdapterForSequenceClassification(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.roberta = RobertaAdaModel(config)
+        self.embeddings = self.roberta.embeddings
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        # self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = RobertaClassificationHead(config)
+
+        self.init_weights()
+        # for param in self.roberta.parameters():
+        #     param.requires_grad = False
+        self.roberta = init_adapter(self.roberta)
+        if self.config.use_freezing:
+            self.roberta = freezer.freeze_lm_component(self.roberta, "adapter")
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm_component(self.roberta, "adapter")
+        else:
+            self.roberta = freezer.unfreeze_lm(self.roberta)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = input_ids.shape[0]
+        inputs_embeds = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+        )
+
+        outputs = self.roberta(
+            # input_ids,
+            attention_mask=attention_mask,
+            # token_type_ids=token_type_ids,
+            # position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            # past_key_values=past_key_values,
+        )
+
+        # pooled_output = outputs[1]
+        sequence_output = outputs[0]
+        # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        first_token_tensor = sequence_output[:, 0]
+        pooled_output = self.roberta.pooler.dense(first_token_tensor)
+        pooled_output = self.roberta.pooler.activation(pooled_output)
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+# ========= DeBERTa =========
+
+#  Prefix-tuning For DeBERTa
+class DebertaPrefixForSequenceClassification(DebertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.deberta = DebertaModel(config)
+        self.pooler = ContextPooler(config)
+        output_dim = self.pooler.output_dim
+        self.classifier = torch.nn.Linear(output_dim, self.num_labels)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.init_weights()
+
+        # for param in self.deberta.parameters():
+        #     param.requires_grad = False
+
+        if self.config.use_freezing:
+            self.deberta = freezer.freeze_lm(self.deberta)
+
+        self.pre_seq_len = config.pre_seq_len
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+        self.prefix_encoder = PrefixEncoder(config)
+
+        deberta_param = 0
+        for name, param in self.deberta.named_parameters():
+            deberta_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - deberta_param
+        print("total param is {}".format(total_param)) # 9860105
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.deberta = freezer.freeze_lm(self.deberta)
+        else:
+            self.deberta = freezer.unfreeze_lm(self.deberta)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.deberta.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        # bsz, seqlen, _ = past_key_values.shape
+        past_key_values = past_key_values.view(
+            batch_size,
+            self.pre_seq_len,
+            self.n_layer * 2,
+            self.n_head,
+            self.n_embd
+        )
+        past_key_values = self.dropout(past_key_values)
+        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+        return past_key_values
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = input_ids.shape[0]
+        past_key_values = self.get_prompt(batch_size=batch_size)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.deberta.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.deberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            past_key_values=past_key_values,
+        )
+
+        encoder_layer = outputs[0]
+        pooled_output = self.pooler(encoder_layer)
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # regression task
+                loss_fn = torch.nn.MSELoss()
+                logits = logits.view(-1).to(labels.dtype)
+                loss = loss_fn(logits, labels.view(-1))
+            elif labels.dim() == 1 or labels.size(-1) == 1:
+                label_index = (labels >= 0).nonzero()
+                labels = labels.long()
+                if label_index.size(0) > 0:
+                    labeled_logits = torch.gather(logits, 0, label_index.expand(label_index.size(0), logits.size(1)))
+                    labels = torch.gather(labels, 0, label_index.view(-1))
+                    loss_fct = CrossEntropyLoss()
+                    loss = loss_fct(labeled_logits.view(-1, self.num_labels).float(), labels.view(-1))
+                else:
+                    loss = torch.tensor(0).to(logits)
+            else:
+                log_softmax = torch.nn.LogSoftmax(-1)
+                loss = -((log_softmax(logits) * labels).sum(-1)).mean()
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+        else:
+            return SequenceClassifierOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )
+
+
+# GPT2 for classification
+class GPT2ForSequenceClassification(GPT2PreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.transformer = GPT2Model(config)
+        self.score = torch.nn.Linear(config.n_embd, self.num_labels, bias=False)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size, sequence_length = input_ids.shape[:2]
+        else:
+            batch_size, sequence_length = inputs_embeds.shape[:2]
+
+        assert (
+            self.config.pad_token_id is not None or batch_size == 1
+        ), "Cannot handle batch sizes > 1 if no padding token is defined."
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1
+            else:
+                sequence_lengths = -1
+                logger.warning(
+                    f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                    "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+                )
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+
+
+# Bart for classification
+class BartForSequenceClassification(BartPretrainedModel):
+    def __init__(self, config: BartConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.model = BartModel(config)
+        self.classification_head = BartClassificationHead(
+            config.d_model,
+            config.d_model,
+            config.num_labels,
+            config.classifier_dropout,
+        )
+        self.model._init_weights(self.classification_head.dense)
+        self.model._init_weights(self.classification_head.out_proj)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        decoder_head_mask: Optional[torch.Tensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, Seq2SeqSequenceClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if labels is not None:
+            use_cache = False
+
+        if input_ids is None and inputs_embeds is not None:
+            raise NotImplementedError(
+                f"Passing input embeddings is currently not supported for {self.__class__.__name__}"
+            )
+
+        outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            head_mask=head_mask,
+            decoder_head_mask=decoder_head_mask,
+            cross_attn_head_mask=cross_attn_head_mask,
+            encoder_outputs=encoder_outputs,
+            inputs_embeds=inputs_embeds,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = outputs[0]  # last hidden state
+        # print("hidden_states.shape=", hidden_states.shape) # [bz, seq_len, dim]
+
+        eos_mask = input_ids.eq(self.config.eos_token_id)
+
+        if len(torch.unique_consecutive(eos_mask.sum(1))) > 1:
+            raise ValueError("All examples must have the same number of <eos> tokens.")
+        sentence_representation = hidden_states[eos_mask, :].view(hidden_states.size(0), -1, hidden_states.size(-1))[
+            :, -1, :
+        ]
+        logits = self.classification_head(sentence_representation)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.config.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.config.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return Seq2SeqSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            decoder_hidden_states=outputs.decoder_hidden_states,
+            decoder_attentions=outputs.decoder_attentions,
+            cross_attentions=outputs.cross_attentions,
+            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+            encoder_hidden_states=outputs.encoder_hidden_states,
+            encoder_attentions=outputs.encoder_attentions,
+        )

models/sequence_classification/masked_prompt_cls.py

@@ -0,0 +1,2016 @@
+"""Custom models for few-shot learning specific operations."""
+
+import torch
+import torch.nn as nn
+import transformers
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer, EvalPrediction
+from transformers.models.t5.modeling_t5 import T5ForConditionalGeneration
+from transformers.models.bert.modeling_bert import BertPreTrainedModel, BertForSequenceClassification, BertModel, BertOnlyMLMHead
+from transformers.models.roberta.modeling_roberta import RobertaForSequenceClassification, RobertaModel, RobertaLMHead, RobertaClassificationHead, RobertaPreTrainedModel
+from transformers.models.deberta_v2.modeling_deberta_v2 import DebertaV2PreTrainedModel, DebertaV2Model, StableDropout, ContextPooler, DebertaV2OnlyMLMHead
+from transformers.models.deberta.modeling_deberta import DebertaPreTrainedModel, DebertaModel, StableDropout, ContextPooler, DebertaOnlyMLMHead
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.bert.configuration_bert import BertConfig
+import logging
+from models.basic_modules.adapter import RobertaAdaModel, BertAdaModel
+import os
+from models.basic_modules.prefix_encoder import PrefixEncoder
+from tools.model_utils.parameter_freeze import ParameterFreeze
+
+freezer = ParameterFreeze()
+
+logger = logging.getLogger(__name__)
+
+# Note: 如果mask_pos为None，请检查输入的模板是否有<mask>标记，是否修改data_collator文件
+
+"""
+Vanilla Prompt-tuning BERT
+"""
+class PromptBertForSequenceClassification(BertPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.pre_seq_len = self.config.pre_seq_len
+        self.hidden_size = self.config.hidden_size
+        # backbone
+        self.bert = BertModel(config)
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        # mlm head
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        """
+        Encoding and obtain logits at masked position
+        """
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+        # Encode everything
+        if inputs_embeds is None:
+            outputs = self.bert(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids
+            )
+        else:
+            outputs = self.bert(
+                None,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids,
+                inputs_embeds=inputs_embeds
+            )
+        # Get <mask> token representation
+        sequence_output, pooled_output = outputs[:2]
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+
+"""
+P-tuning BERT
+"""
+class PromptBertPtuningForSequenceClassification(BertPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.pre_seq_len = self.config.pre_seq_len
+        self.hidden_size = self.config.hidden_size
+        # backbone
+        self.bert = BertModel(config)
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        # mlm head
+        self.cls = BertOnlyMLMHead(config)
+        # prompt encoder
+        self.prompt_encoder = None
+        # plm embedding layer
+        self.backbone_embeddings = self.bert.embeddings.word_embeddings
+        # prompt embedding layer
+        self.prompt_embeddings = torch.nn.Embedding(self.pre_seq_len, self.hidden_size)
+
+        self.init_weights()
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+
+    def generate_continuous_prompt_inputs(self, input_ids, block_flag=None, reparameterization=False):
+        """
+        Generate continuous prompt embedding
+        """
+        inputs_embeds = self.backbone_embeddings(input_ids)
+
+        batch_size = inputs_embeds.shape[0]
+        if block_flag is None:
+            # the first token is set 1, others are set 0
+            block_flag = torch.zeros_like(input_ids).long().to(inputs_embeds.device)
+            block_flag[:, 0] = 1
+        try:
+            replace_embeds = self.prompt_embeddings(
+                torch.LongTensor(list(range(self.pre_seq_len))).to(inputs_embeds.device))
+        except:
+            import pdb
+            pdb.set_trace()
+            replace_embeds = self.prompt_embeddings(
+                torch.LongTensor(list(range(self.pre_seq_len))))
+        replace_embeds = replace_embeds.unsqueeze(0)  # [batch_size, prompt_length, embed_size]
+
+        if self.prompt_encoder is not None:
+            replace_embeds = self.prompt_encoder(replace_embeds)
+
+        # edit by wjn
+        if reparameterization:
+            # blocked_indices = (block_flag == 1).nonzero(as_tuple=False).reshape((batch_size, self.pre_seq_len, 2))[:, :, 1]
+            blocked_indices = (block_flag == 1).nonzero()
+            # reparameterization
+            for bidx in range(batch_size):
+                for i in range(blocked_indices.shape[1]):
+                    inputs_embeds[bidx, blocked_indices[bidx, i], :] = replace_embeds[:, i, :].squeeze()
+        else:
+            replace_embeds = replace_embeds.expand(batch_size, self.pre_seq_len, -1).to(inputs_embeds.device)
+            inputs_embeds = torch.cat((replace_embeds, inputs_embeds), dim=1)
+        return inputs_embeds
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        """
+        Encoding and obtain logits at masked position
+        """
+        batch_size = inputs_embeds.shape[0]
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+        # Encode everything
+        if inputs_embeds is None:
+            outputs = self.bert(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids
+            )
+        else:
+
+            if inputs_embeds.shape[1] == attention_mask.shape[1]:
+                outputs = self.bert(
+                    None,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                    inputs_embeds=inputs_embeds
+                )
+                # Get <mask> token representation
+                sequence_output, pooled_output = outputs[:2]
+                # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+            else:
+                if attention_mask is not None:
+                    prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).long().to(self.bert.device)
+                    attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+                if token_type_ids is not None:
+                    prefix_token_type_ids = torch.zeros(batch_size, self.pre_seq_len).long().to(self.bert.device)
+                    token_type_ids = torch.cat((prefix_token_type_ids, token_type_ids), dim=1)
+                outputs = self.bert(
+                    None,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                    inputs_embeds=inputs_embeds
+                )
+                # Get <mask> token representation
+                sequence_output, pooled_output = outputs[:2]
+                sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        inputs_embeds = self.generate_continuous_prompt_inputs(input_ids, block_flag)
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+
+"""
+Prefix-tuning BERT
+"""
+class PromptBertPrefixForSequenceClassification(BertPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+
+        self.num_labels = config.num_labels
+        self.pre_seq_len = self.config.pre_seq_len
+        self.hidden_size = self.config.hidden_size
+
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        # backbone
+        self.bert = BertModel(config)
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        # mlm head
+        self.cls = BertOnlyMLMHead(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        # plm embedding layer
+        self.backbone_embeddings = self.bert.embeddings.word_embeddings
+        # prompt embedding layer
+        self.prompt_embeddings = torch.nn.Embedding(self.pre_seq_len, self.hidden_size)
+        # prefix encoder
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+        self.prefix_encoder = PrefixEncoder(config)
+
+        self.init_weights()
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm(self.bert)
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.bert.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        # bsz, seqlen, _ = past_key_values.shape
+        past_key_values = past_key_values.view(
+            batch_size,
+            self.pre_seq_len,
+            self.n_layer * 2,
+            self.n_head,
+            self.n_embd
+        )
+        past_key_values = self.dropout(past_key_values)
+        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+        return past_key_values
+
+    def embed_encode(self, input_ids):
+        embedding_output = self.bert.embeddings.word_embeddings(input_ids)
+        return embedding_output
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        batch_size = input_ids.size(0)
+
+        # add prefix for prompt-tuning
+        past_key_values = self.get_prompt(batch_size=batch_size)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.bert.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+
+        # Encode everything
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            past_key_values=past_key_values,
+        )
+        # Get <mask> token representation
+        sequence_output, pooled_output = outputs[:2]
+        # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # print("prediction_mask_scores.shape=", prediction_mask_scores.shape) # [batch_size, seq_len, vocab_size]
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+"""
+Adapter-tuning BERT
+"""
+class PromptBertAdapterForSequenceClassification(BertPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.bert = BertAdaModel(config)
+        self.cls = BertOnlyMLMHead(config)
+        self.init_weights()
+
+        if self.config.use_freezing:
+            self.bert = freezer.freeze_lm_component(self.bert, "adapter")
+
+        # These attributes should be assigned once the model is initialized
+        self.model_args = None
+        self.data_args = None
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.bert = freezer.freeze_lm_component(self.bert, "adapter")
+        else:
+            self.bert = freezer.unfreeze_lm(self.bert)
+
+    def embed_encode(self, input_ids):
+        embedding_output = self.bert.embeddings.word_embeddings(input_ids)
+        return embedding_output
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        batch_size = input_ids.size(0)
+
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+
+        # Encode everything
+        if inputs_embeds is None:
+            outputs = self.bert(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids
+            )
+        else:
+            outputs = self.bert(
+                None,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids,
+                inputs_embeds=inputs_embeds
+            )
+
+        # Get <mask> token representation
+        sequence_output, pooled_output = outputs[:2]
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+
+"""
+Vanilla Prompt-tuning RoBERTa
+"""
+class PromptRobertaForSequenceClassification(RobertaPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.pre_seq_len = self.config.pre_seq_len
+        self.hidden_size = self.config.hidden_size
+        # backbone
+        self.roberta = RobertaModel(config)
+        if self.config.use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        # mlm head
+        self.cls = RobertaLMHead(config)
+
+        self.init_weights()
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.roberta.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        else:
+            self.roberta = freezer.unfreeze_lm(self.roberta)
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        """
+        Encoding and obtain logits at masked position
+        """
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+        # Encode everything
+        if inputs_embeds is None:
+            outputs = self.roberta(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids
+            )
+        else:
+            outputs = self.roberta(
+                None,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids,
+                inputs_embeds=inputs_embeds
+            )
+        # Get <mask> token representation
+        sequence_output, pooled_output = outputs[:2]
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+"""
+P-tuning RoBERTa
+"""
+class PromptRobertaPtuningForSequenceClassification(RobertaPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.pre_seq_len = self.config.pre_seq_len
+        self.hidden_size = self.config.hidden_size
+        # backbone
+        self.roberta = RobertaModel(config)
+        if self.config.use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        # mlm head
+        self.cls = RobertaLMHead(config)
+        # prompt encoder
+        self.prompt_encoder = None
+        # plm embedding layer
+        self.backbone_embeddings = self.roberta.embeddings.word_embeddings
+        # prompt embedding layer
+        self.prompt_embeddings = torch.nn.Embedding(self.pre_seq_len, self.hidden_size)
+
+        self.init_weights()
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.roberta.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm(self.roberta)
+        else:
+            self.roberta = freezer.unfreeze_lm(self.roberta)
+
+
+    def generate_continuous_prompt_inputs(self, input_ids, block_flag=None, reparameterization=False):
+        """
+        Generate continuous prompt embedding
+        """
+        inputs_embeds = self.backbone_embeddings(input_ids)
+
+        batch_size = inputs_embeds.shape[0]
+        if block_flag is None:
+            # the first token is set 1, others are set 0
+            block_flag = torch.zeros_like(input_ids).long().to(inputs_embeds.device)
+            block_flag[:, 0] = 1
+        try:
+            replace_embeds = self.prompt_embeddings(
+                torch.LongTensor(list(range(self.pre_seq_len))).to(inputs_embeds.device))
+        except:
+            import pdb
+            pdb.set_trace()
+            replace_embeds = self.prompt_embeddings(torch.LongTensor(list(range(self.pre_seq_len))))
+        replace_embeds = replace_embeds.unsqueeze(0)  # [batch_size, prompt_length, embed_size]
+
+        if self.prompt_encoder is not None:
+            replace_embeds = self.prompt_encoder(replace_embeds)
+
+        # edit by wjn
+        if reparameterization:
+            # blocked_indices = (block_flag == 1).nonzero(as_tuple=False).reshape((batch_size, self.pre_seq_len, 2))[:, :, 1]
+            blocked_indices = (block_flag == 1).nonzero()
+            # reparameterization
+            for bidx in range(batch_size):
+                for i in range(blocked_indices.shape[1]):
+                    inputs_embeds[bidx, blocked_indices[bidx, i], :] = replace_embeds[:, i, :].squeeze()
+        else:
+            replace_embeds = replace_embeds.expand(batch_size, self.pre_seq_len, -1).to(inputs_embeds.device)
+            inputs_embeds = torch.cat((replace_embeds, inputs_embeds), dim=1)
+        return inputs_embeds
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        """
+        Encoding and obtain logits at masked position
+        """
+        batch_size = inputs_embeds.shape[0]
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+        # Encode everything
+        if inputs_embeds is None:
+            outputs = self.roberta(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids
+            )
+        else:
+
+            if inputs_embeds.shape[1] == attention_mask.shape[1]:
+                outputs = self.roberta(
+                    None,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                    inputs_embeds=inputs_embeds
+                )
+                # Get <mask> token representation
+                sequence_output, pooled_output = outputs[:2]
+                # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+            else:
+                if attention_mask is not None:
+                    prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).long().to(self.roberta.device)
+                    attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+                if token_type_ids is not None:
+                    prefix_token_type_ids = torch.zeros(batch_size, self.pre_seq_len).long().to(self.roberta.device)
+                    token_type_ids = torch.cat((prefix_token_type_ids, token_type_ids), dim=1)
+                outputs = self.roberta(
+                    None,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                    inputs_embeds=inputs_embeds
+                )
+                # Get <mask> token representation
+                sequence_output, pooled_output = outputs[:2]
+                sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        inputs_embeds = self.generate_continuous_prompt_inputs(input_ids, block_flag)
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+"""
+Prefix-tuning RoBERTa
+"""
+class PromptRobertaPrefixForSequenceClassification(RobertaPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+
+        self.num_labels = config.num_labels
+        self.pre_seq_len = self.config.pre_seq_len
+        self.hidden_size = self.config.hidden_size
+
+        self.n_layer = config.num_hidden_layers
+        self.n_head = config.num_attention_heads
+        self.n_embd = config.hidden_size // config.num_attention_heads
+
+        # backbone
+        self.robert = RobertaModel(config)
+        if self.config.use_freezing:
+            self.robert = freezer.freeze_lm(self.robert)
+        # mlm head
+        self.cls = RobertaLMHead(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        # plm embedding layer
+        self.backbone_embeddings = self.robert.embeddings.word_embeddings
+        # prompt embedding layer
+        self.prompt_embeddings = torch.nn.Embedding(self.pre_seq_len, self.hidden_size)
+        # prefix encoder
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+        self.prefix_encoder = PrefixEncoder(config)
+
+        self.init_weights()
+
+        # These attributes should be assigned once the model is initialized
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.robert.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.robert = freezer.freeze_lm(self.robert)
+        else:
+            self.robert = freezer.unfreeze_lm(self.robert)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.robert.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        # bsz, seqlen, _ = past_key_values.shape
+        past_key_values = past_key_values.view(
+            batch_size,
+            self.pre_seq_len,
+            self.n_layer * 2,
+            self.n_head,
+            self.n_embd
+        )
+        past_key_values = self.dropout(past_key_values)
+        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+        return past_key_values
+
+    def embed_encode(self, input_ids):
+        embedding_output = self.robert.embeddings.word_embeddings(input_ids)
+        return embedding_output
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        batch_size = input_ids.size(0)
+
+        # add prefix for prompt-tuning
+        past_key_values = self.get_prompt(batch_size=batch_size)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.robert.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+
+        # Encode everything
+        outputs = self.robert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            past_key_values=past_key_values,
+        )
+        # Get <mask> token representation
+        sequence_output, pooled_output = outputs[:2]
+        # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+"""
+Adapter-tuning RoBERTa
+"""
+class PromptRobertaAdapterForSequenceClassification(RobertaPreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.roberta = RobertaAdaModel(config)
+        self.cls = RobertaLMHead(config)
+        self.init_weights()
+
+        if self.config.use_freezing:
+            self.roberta = freezer.freeze_lm_component(self.roberta, "adapter")
+
+        # These attributes should be assigned once the model is initialized
+        self.model_args = None
+        self.data_args = None
+        self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.roberta.device)
+
+        # For regression
+        self.lb = None
+        self.ub = None
+
+        # For label search.
+        self.return_full_softmax = None
+
+    def freeze_backbone(self, use_freezing: bool=True):
+        if use_freezing:
+            self.roberta = freezer.freeze_lm_component(self.roberta, "adapter")
+        else:
+            self.roberta = freezer.unfreeze_lm(self.berobertart)
+
+    def embed_encode(self, input_ids):
+        embedding_output = self.roberta.embeddings.word_embeddings(input_ids)
+        return embedding_output
+
+    def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+        batch_size = input_ids.size(0)
+
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+
+        # Encode everything
+        if inputs_embeds is None:
+            outputs = self.roberta(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids
+            )
+        else:
+            outputs = self.roberta(
+                None,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids,
+                inputs_embeds=inputs_embeds
+            )
+
+        # Get <mask> token representation
+        sequence_output, pooled_output = outputs[:2]
+        sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+        # Logits over vocabulary tokens
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Exit early and only return mask logits.
+        if return_full_softmax:
+            return prediction_mask_scores
+
+        # Return logits for each label
+        logits = []
+        for label_id in range(len(self.label_word_list)):
+            logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+        logits = torch.cat(logits, -1)
+
+        # Regression task
+        if self.config.num_labels == 1:
+            logsoftmax = nn.LogSoftmax(-1)
+            logits = logsoftmax(logits)  # Log prob of right polarity
+
+        return logits, sequence_mask_output
+
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        mask_pos=None,
+        labels=None,
+        inputs_embeds=None,
+        block_flag=None,
+        return_dict=None,
+    ):
+
+        logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                # Regression task
+                loss_fct = nn.KLDivLoss(log_target=True)
+                labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+                loss = loss_fct(logits.view(-1, 2), labels)
+            else:
+
+                if labels.shape == logits.shape:
+                    loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+                                    labels, reduction="batchmean")
+                else:
+                    loss_fct = nn.CrossEntropyLoss()
+
+                    loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        output = (logits,)
+        if self.num_labels == 1:
+            # Regression output
+            output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+        if not return_dict:
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+
+# class DebertaForPromptFinetuning(DebertaPreTrainedModel):
+#     _keys_to_ignore_on_load_unexpected = [r"pooler"]
+#     _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+#     def __init__(self, config):
+#         super().__init__(config)
+#         self.num_labels = config.num_labels
+#         #self.deberta = DebertaV2Model(config)
+
+#         self.deberta = DebertaModel(config)
+#         self.cls = DebertaOnlyMLMHead(config)
+
+#         if self.config.use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+
+#         self.pooler = ContextPooler(config)
+#         output_dim = self.pooler.output_dim
+
+#         self.classifier = torch.nn.Linear(output_dim, self.num_labels)
+#         drop_out = getattr(config, "cls_dropout", None)
+#         drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
+
+#         self.dropout = StableDropout(drop_out)
+
+#         classification_list = [self.pooler, self.dropout,self.classifier]
+
+#         self.classifier = nn.Sequential(*classification_list)
+#         # self.cls = DebertaV2OnlyMLMHead(config)
+
+#         self.map = nn.Linear(config.hidden_size, config.hidden_size)
+#         self.init_weights()
+
+#         # These attributes should be assigned once the model is initialized
+#         self.model_args = None
+#         self.data_args = None
+#         self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+#         self.K = 1
+#         self.step_size=1e-5
+#         # import pdb
+#         # pdb.set_trace()
+#         #self.step_size=config.step_size
+
+#         # For regression
+#         self.lb = None
+#         self.ub = None
+
+#         self.pre_seq_len = self.config.pre_seq_len
+#         # For auto label search.
+#         self.return_full_softmax = None
+
+#     def freeze_backbone(self, use_freezing: bool=True):
+#         if use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+#         else:
+#             self.deberta = freezer.unfreeze_lm(self.deberta)
+
+
+
+#     def embed_encode(self, input_ids):
+#         embedding_output = self.deberta.embeddings.word_embeddings(input_ids)
+#         return embedding_output
+
+#     def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None,
+#                return_full_softmax=False):
+#         batch_size = input_ids.size(0)
+
+#         if mask_pos is not None:
+#             mask_pos = mask_pos.squeeze()
+
+
+#         # Encode everything
+#         if inputs_embeds is None:
+#             outputs = self.deberta(
+#                 input_ids,
+#                 attention_mask=attention_mask,
+#                 token_type_ids=token_type_ids
+#             )
+#         else:
+#             outputs = self.deberta(
+#                 None,
+#                 attention_mask=attention_mask,
+#                 token_type_ids=token_type_ids,
+#                 inputs_embeds=inputs_embeds
+#             )
+
+#         # Get <mask> token representation
+#         sequence_output = outputs[0]
+#         sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+#         sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+#         # Logits over vocabulary tokens
+#         prediction_mask_scores = self.cls(sequence_mask_output)
+
+#         # sequence_mask_output = self.lm_head.dense(sequence_mask_output)
+
+#         # Exit early and only return mask logits.
+#         if return_full_softmax:
+#             return prediction_mask_scores
+
+#         # Return logits for each label
+#         logits = []
+#         for label_id in range(len(self.label_word_list)):
+#             logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+#         logits = torch.cat(logits, -1)
+
+#         # Regression task
+#         if self.config.num_labels == 1:
+#             logsoftmax = nn.LogSoftmax(-1)
+#             logits = logsoftmax(logits)  # Log prob of right polarity
+
+#         if self.model_args.hybrid == 1:
+#             cls_logits = self.classifier(sequence_output)
+#             return (logits, cls_logits), sequence_mask_output
+
+#         return logits, sequence_mask_output
+
+#     def forward(
+#             self,
+#             input_ids=None,
+#             attention_mask=None,
+#             token_type_ids=None,
+#             mask_pos=None,
+#             labels=None,
+#             inputs_embeds=None,
+#             fwd_type=0,
+#             block_flag=None
+#     ):
+
+#         if fwd_type == 2:
+#             assert inputs_embeds is not None
+#             return self.encode(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids,
+#                                mask_pos=mask_pos, inputs_embeds=inputs_embeds)
+
+#         elif fwd_type == 1:
+#             return self.embed_encode(input_ids)
+
+
+
+#         if (self.model_args.prompt_ptuning or self.model_args.prompt_prefix) and block_flag is not None:
+#             inputs_embeds = self.generate_continuous_prompt_inputs(input_ids, block_flag)
+
+#         logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+
+#         if self.model_args.hybrid == 1:
+#             logits = logits[0]
+#             cls_logits = logits[1]
+
+#         loss = None
+#         if labels is not None:
+#             if self.num_labels == 1:
+#                 # Regression task
+#                 loss_fct = nn.KLDivLoss(log_target=True)
+#                 labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb),
+#                                       (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+#                 loss = loss_fct(logits.view(-1, 2), labels)
+#             else:
+
+#                 if labels.shape == logits.shape:
+#                     loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+#                                     labels, reduction="batchmean")
+#                 else:
+#                     loss_fct = nn.CrossEntropyLoss()
+
+#                     loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+#         output = (logits,)
+#         if self.num_labels == 1:
+#             # Regression output
+#             output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+#         return ((loss,) + output) if loss is not None else output
+
+
+
+# # add by wjn
+# # Prefix-tuning for Deberta
+# class DebertaPrefixForPromptFinetuning(DebertaPreTrainedModel):
+
+#     def __init__(self, config):
+#         super().__init__(config)
+#         self.num_labels = config.num_labels
+#         #self.deberta = DebertaV2Model(config)
+
+#         self.deberta = DebertaModel(config)
+#         self.cls = DebertaOnlyMLMHead(config)
+
+#         self.pooler = ContextPooler(config)
+#         output_dim = self.pooler.output_dim
+
+#         self.classifier = torch.nn.Linear(output_dim, self.num_labels)
+#         drop_out = getattr(config, "cls_dropout", None)
+#         drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
+
+#         self.dropout = StableDropout(drop_out)
+
+#         classification_list = [self.pooler, self.dropout,self.classifier]
+
+#         self.classifier = nn.Sequential(*classification_list)
+#         # self.cls = DebertaV2OnlyMLMHead(config)
+
+#         self.map = nn.Linear(config.hidden_size, config.hidden_size)
+#         self.init_weights()
+
+#         if self.config.use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+
+#         self.pre_seq_len = config.pre_seq_len
+#         self.n_layer = config.num_hidden_layers
+#         self.n_head = config.num_attention_heads
+#         self.n_embd = config.hidden_size // config.num_attention_heads
+
+#         self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+#         self.prefix_encoder = PrefixEncoder(config)
+
+#         # These attributes should be assigned once the model is initialized
+#         self.model_args = None
+#         self.data_args = None
+#         self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+#         self.K = 1
+#         self.step_size=1e-5
+#         # import pdb
+#         # pdb.set_trace()
+#         #self.step_size=config.step_size
+
+#         # For regression
+#         self.lb = None
+#         self.ub = None
+
+
+#         # For auto label search.
+#         self.return_full_softmax = None
+
+#     def freeze_backbone(self, use_freezing: bool=True):
+#         if use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+#         else:
+#             self.deberta = freezer.unfreeze_lm(self.deberta)
+
+#     def get_prompt(self, batch_size):
+#         prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.deberta.device)
+#         past_key_values = self.prefix_encoder(prefix_tokens)
+#         # bsz, seqlen, _ = past_key_values.shape
+#         past_key_values = past_key_values.view(
+#             batch_size,
+#             self.pre_seq_len,
+#             self.n_layer * 2,
+#             self.n_head,
+#             self.n_embd
+#         )
+#         past_key_values = self.dropout(past_key_values)
+#         past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+#         return past_key_values
+
+
+#     def get_constrast_loss(self,
+#                     input_ids=None,
+#                     attention_mask=None,
+#                     mask_pos=None,
+#                     labels=None,
+#                     inputs_embeds=None):
+
+#         self.cos = nn.CosineSimilarity(dim=-1)
+
+
+#         _, sequence_mask_output_1 = self.encode(input_ids, attention_mask, mask_pos, inputs_embeds)
+#         _, sequence_mask_output_2 = self.encode(input_ids, attention_mask, mask_pos, inputs_embeds)
+
+#         sequence_mask_output_1= self.lm_head.dense(sequence_mask_output_1)
+#         sequence_mask_output_2 = self.lm_head.dense(sequence_mask_output_2)
+#         # input_args = [input_ids, attention_mask, mask_pos, labels, None, 1]
+#         # embed = self.forward(*input_args)
+#         #
+#         # vat_args = [input_ids, attention_mask, mask_pos, labels, embed, 2]
+#         #
+#         # adv_logits, outputs = self.forward(*vat_args)
+#         #
+#         # logit_mask = F.softmax(logits, dim=-1)[torch.arange(adv_logits.size(0)), labels] > 0.7
+#         #
+#         # outputs = outputs[logit_mask]
+#         # seq_outputs = sequence_mask_output[logit_mask]
+#         # new_label = labels[logit_mask]
+#         # #
+#         # #
+#         # rand_perm = torch.randperm(outputs.size(0))
+#         # rand_outputs = outputs[rand_perm, :]
+#         # rand_label = new_label[rand_perm]
+#         # pair_label = (new_label == rand_label).long()
+#         #
+#         # seq_outputs = self.map(seq_outputs)
+#         # rand_outputs = self.map(rand_outputs)
+
+#         pair_labels = (labels.unsqueeze(1) == labels.unsqueeze(0)).float()
+
+#         # import  pdb
+#         # pdb.set_trace()
+
+#         contra_loss = self.contra_lc(sequence_mask_output_1.unsqueeze(1), sequence_mask_output_2.unsqueeze(0), pair_labels)
+
+#         if torch.isnan(contra_loss):
+#             return 0
+
+#         return contra_loss
+
+#     def embed_encode(self, input_ids):
+#         embedding_output = self.deberta.embeddings.word_embeddings(input_ids)
+#         return embedding_output
+
+#     def encode(self, input_ids=None, attention_mask=None, token_type_ids=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+#         batch_size = input_ids.size(0)
+
+#         # add prefix for prompt-tuning
+#         past_key_values = self.get_prompt(batch_size=batch_size)
+#         prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.deberta.device)
+#         attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+#         if mask_pos is not None:
+#             mask_pos = mask_pos.squeeze()
+
+#         # Encode everything
+
+#         outputs = self.deberta(
+#             input_ids,
+#             attention_mask=attention_mask,
+#             token_type_ids=token_type_ids,
+#             past_key_values=past_key_values,
+#         )
+
+
+#         # Get <mask> token representation
+#         sequence_output, pooled_output = outputs[:2]
+#         # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+#         sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+#         # Logits over vocabulary tokens
+#         prediction_mask_scores = self.cls(sequence_mask_output)
+
+#         #sequence_mask_output = self.lm_head.dense(sequence_mask_output)
+
+#         # Exit early and only return mask logits.
+#         if return_full_softmax:
+#             return prediction_mask_scores
+
+#         # Return logits for each label
+#         logits = []
+#         for label_id in range(len(self.label_word_list)):
+#             logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+#         logits = torch.cat(logits, -1)
+
+#         # Regression task
+#         if self.config.num_labels == 1:
+#             logsoftmax = nn.LogSoftmax(-1)
+#             logits = logsoftmax(logits)  # Log prob of right polarity
+
+#         if self.model_args.hybrid == 1:
+#             cls_logits = self.classifier(sequence_output)
+#             return (logits, cls_logits), sequence_mask_output
+
+#         return logits, sequence_mask_output
+
+
+#     def forward(
+#             self,
+#             input_ids=None,
+#             attention_mask=None,
+#             token_type_ids=None,
+#             mask_pos=None,
+#             labels=None,
+#             inputs_embeds=None,
+#             fwd_type=0,
+#             block_flag=None,
+#             return_dict=None,
+#     ):
+
+#         if fwd_type == 2:
+#             assert inputs_embeds is not None
+#             return self.encode(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids,
+#                                mask_pos=mask_pos, inputs_embeds=inputs_embeds)
+
+#         elif fwd_type == 1:
+#             return self.embed_encode(input_ids)
+
+
+
+#         if (self.model_args.prompt_ptuning or self.model_args.prompt_prefix) and block_flag is not None:
+#             inputs_embeds = self.generate_continuous_prompt_inputs(input_ids, block_flag)
+
+#         logits, sequence_mask_output = self.encode(input_ids, attention_mask, token_type_ids, mask_pos, inputs_embeds)
+
+#         if self.model_args.hybrid == 1:
+#             logits = logits[0]
+#             cls_logits = logits[1]
+
+#         loss = None
+#         if labels is not None:
+#             if self.num_labels == 1:
+#                 # Regression task
+#                 loss_fct = nn.KLDivLoss(log_target=True)
+#                 labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb),
+#                                       (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+#                 loss = loss_fct(logits.view(-1, 2), labels)
+#             else:
+
+#                 if labels.shape == logits.shape:
+#                     loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+#                                     labels, reduction="batchmean")
+#                 else:
+#                     loss_fct = nn.CrossEntropyLoss()
+
+#                     loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+#         output = (logits,)
+#         if self.num_labels == 1:
+#             # Regression output
+#             output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+#         if not return_dict:
+#             return ((loss,) + output) if loss is not None else output
+
+#         return SequenceClassifierOutput(
+#             loss=loss,
+#             logits=logits,
+#         )
+
+
+
+
+# class Debertav2ForPromptFinetuning(DebertaV2PreTrainedModel):
+#     _keys_to_ignore_on_load_unexpected = [r"pooler"]
+#     _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+#     def __init__(self, config):
+#         super().__init__(config)
+#         self.num_labels = config.num_labels
+#         self.deberta = DebertaV2Model(config)
+
+#         if self.config.use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+#         self.cls = DebertaV2OnlyMLMHead(config)
+
+#         #self.deberta = DebertaModel(config)
+#         #self.cls = DebertaOnlyMLMHead(config)
+
+#         self.pooler = ContextPooler(config)
+#         output_dim = self.pooler.output_dim
+
+#         self.classifier = torch.nn.Linear(output_dim, self.num_labels)
+#         drop_out = getattr(config, "cls_dropout", None)
+#         drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
+
+#         self.dropout = StableDropout(drop_out)
+
+#         classification_list = [self.pooler, self.dropout,self.classifier]
+
+#         self.classifier = nn.Sequential(*classification_list)
+#         # self.cls = DebertaV2OnlyMLMHead(config)
+
+#         self.map = nn.Linear(config.hidden_size, config.hidden_size)
+#         self.init_weights()
+
+#         # These attributes should be assigned once the model is initialized
+#         self.model_args = None
+#         self.data_args = None
+#         self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+#         self.K = 1
+#         self.step_size=1e-5
+#         # import pdb
+#         # pdb.set_trace()
+#         #self.step_size=config.step_size
+
+#         # For regression
+#         self.lb = None
+#         self.ub = None
+
+#         self.pre_seq_len = self.config.pre_seq_len
+#         # For auto label search.
+#         self.return_full_softmax = None
+
+#     def freeze_backbone(self, use_freezing: bool=True):
+#         if use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+#         else:
+#             self.deberta = freezer.unfreeze_lm(self.deberta)
+
+#     def embed_encode(self, input_ids):
+#         embedding_output = self.deberta.embeddings.word_embeddings(input_ids)
+#         return embedding_output
+
+#     def encode(self, input_ids=None, attention_mask=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+#         batch_size = input_ids.size(0)
+
+#         if mask_pos is not None:
+#             mask_pos = mask_pos.squeeze()
+
+#         # Encode everything
+#         if inputs_embeds is None:
+#             outputs =  self.deberta(
+#                 input_ids,
+#                 attention_mask=attention_mask
+#             )
+#         else:
+#             outputs =  self.deberta(
+#                 None,
+#                 attention_mask=attention_mask,
+#                 inputs_embeds=inputs_embeds
+#             )
+
+
+#         # Get <mask> token representation
+#         sequence_output = outputs[0]
+#         sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+#         sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+
+#         # Logits over vocabulary tokens
+#         prediction_mask_scores = self.cls(sequence_mask_output)
+
+#         #sequence_mask_output = self.lm_head.dense(sequence_mask_output)
+
+#         # Exit early and only return mask logits.
+#         if return_full_softmax:
+#             return prediction_mask_scores
+
+#         # Return logits for each label
+#         logits = []
+#         for label_id in range(len(self.label_word_list)):
+#             logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+#         logits = torch.cat(logits, -1)
+
+#         # Regression task
+#         if self.config.num_labels == 1:
+#             logsoftmax = nn.LogSoftmax(-1)
+#             logits = logsoftmax(logits)  # Log prob of right polarity
+
+#         return logits, sequence_mask_output
+
+
+#     def forward(
+#         self,
+#         input_ids=None,
+#         attention_mask=None,
+#         mask_pos=None,
+#         labels=None,
+#         inputs_embeds=None,
+#         fwd_type=0,
+#         block_flag=None,
+#         return_dict=None
+#     ):
+#         if fwd_type == 2:
+#             assert inputs_embeds is not None
+#             return self.encode(input_ids=input_ids, attention_mask=attention_mask, mask_pos=mask_pos, inputs_embeds=inputs_embeds)
+
+#         elif fwd_type == 1:
+#             return self.embed_encode(input_ids)
+
+#         logits, sequence_mask_output = self.encode(input_ids, attention_mask, mask_pos, inputs_embeds)
+
+#         loss = None
+
+
+#         if labels is not None:
+#             if self.num_labels == 1:
+#                 # Regression task
+#                 loss_fct = nn.KLDivLoss(log_target=True)
+#                 labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+#                 loss = loss_fct(logits.view(-1, 2), labels)
+#             else:
+
+#                 if labels.shape == logits.shape:
+#                     loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+#                                     labels, reduction="batchmean")
+#                 else:
+#                     loss_fct = nn.CrossEntropyLoss()
+
+#                     loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+#                     if self.model_args.hybrid == 1:
+#                         cls_loss = loss_fct(cls_logits.view(-1, cls_logits.size(-1)), labels.view(-1))
+#                         loss = loss + cls_loss
+
+#         output = (logits,)
+#         if self.num_labels == 1:
+#             # Regression output
+#             output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+#         if not return_dict:
+#             return ((loss,) + output) if loss is not None else output
+
+#         return SequenceClassifierOutput(
+#             loss=loss,
+#             logits=logits,
+#         )
+
+
+# class Debertav2PrefixForPromptFinetuning(DebertaV2PreTrainedModel):
+#     _keys_to_ignore_on_load_unexpected = [r"pooler"]
+#     _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+#     def __init__(self, config):
+#         super().__init__(config)
+#         self.num_labels = config.num_labels
+#         self.deberta = DebertaV2Model(config)
+#         self.cls = DebertaV2OnlyMLMHead(config)
+
+#         #self.deberta = DebertaModel(config)
+#         #self.cls = DebertaOnlyMLMHead(config)
+
+#         self.pooler = ContextPooler(config)
+#         output_dim = self.pooler.output_dim
+
+#         self.classifier = torch.nn.Linear(output_dim, self.num_labels)
+#         drop_out = getattr(config, "cls_dropout", None)
+#         drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
+
+#         self.dropout = StableDropout(drop_out)
+
+#         classification_list = [self.pooler, self.dropout,self.classifier]
+
+#         self.classifier = nn.Sequential(*classification_list)
+#         # self.cls = DebertaV2OnlyMLMHead(config)
+
+#         self.map = nn.Linear(config.hidden_size, config.hidden_size)
+#         self.init_weights()
+
+#         if self.config.use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+
+#         self.pre_seq_len = config.pre_seq_len
+#         self.n_layer = config.num_hidden_layers
+#         self.n_head = config.num_attention_heads
+#         self.n_embd = config.hidden_size // config.num_attention_heads
+
+#         self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+#         self.prefix_encoder = PrefixEncoder(config)
+
+#         # These attributes should be assigned once the model is initialized
+#         self.model_args = None
+#         self.data_args = None
+#         self.label_word_list = torch.Tensor(self.config.label_word_list).long().to(self.bert.device)
+#         self.K = 1
+#         self.step_size=1e-5
+#         # import pdb
+#         # pdb.set_trace()
+#         #self.step_size=config.step_size
+
+#         # For regression
+#         self.lb = None
+#         self.ub = None
+
+
+#         # For auto label search.
+#         self.return_full_softmax = None
+
+#     def freeze_backbone(self, use_freezing: bool=True):
+#         if use_freezing:
+#             self.deberta = freezer.freeze_lm(self.deberta)
+#         else:
+#             self.deberta = freezer.unfreeze_lm(self.deberta)
+
+#     def get_prompt(self, batch_size):
+#         prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.deberta.device)
+#         past_key_values = self.prefix_encoder(prefix_tokens)
+#         # bsz, seqlen, _ = past_key_values.shape
+#         past_key_values = past_key_values.view(
+#             batch_size,
+#             self.pre_seq_len,
+#             self.n_layer * 2,
+#             self.n_head,
+#             self.n_embd
+#         )
+#         past_key_values = self.dropout(past_key_values)
+#         past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
+#         return past_key_values
+
+
+#     def embed_encode(self, input_ids):
+#         embedding_output = self.deberta.embeddings.word_embeddings(input_ids)
+#         return embedding_output
+
+#     def encode(self, input_ids=None, attention_mask=None, mask_pos=None, inputs_embeds=None, return_full_softmax=False):
+#         batch_size = input_ids.size(0)
+
+#         # add prefix for prompt-tuning
+#         past_key_values = self.get_prompt(batch_size=batch_size)
+#         prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(self.deberta.device)
+#         attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+
+#         if mask_pos is not None:
+#             mask_pos = mask_pos.squeeze()
+
+#         # Encode everything
+#         outputs = self.deberta(
+#             input_ids,
+#             attention_mask=attention_mask,
+#             past_key_values=past_key_values,
+#         )
+
+
+#         # Get <mask> token representation
+#         sequence_output = outputs[0]
+#         # sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+#         sequence_mask_output = sequence_output[torch.arange(sequence_output.size(0)), mask_pos]
+
+
+#         # Logits over vocabulary tokens
+#         prediction_mask_scores = self.cls(sequence_mask_output)
+
+#         #sequence_mask_output = self.lm_head.dense(sequence_mask_output)
+
+#         # Exit early and only return mask logits.
+#         if return_full_softmax:
+#             return prediction_mask_scores
+
+#         # Return logits for each label
+#         logits = []
+#         for label_id in range(len(self.label_word_list)):
+#             logits.append(prediction_mask_scores[:, self.label_word_list[label_id]].unsqueeze(-1))
+#         logits = torch.cat(logits, -1)
+
+#         # Regression task
+#         if self.config.num_labels == 1:
+#             logsoftmax = nn.LogSoftmax(-1)
+#             logits = logsoftmax(logits)  # Log prob of right polarity
+
+#         return logits, sequence_mask_output
+
+
+#     def forward(
+#         self,
+#         input_ids=None,
+#         attention_mask=None,
+#         mask_pos=None,
+#         labels=None,
+#         inputs_embeds=None,
+#         fwd_type=0,
+#         block_flag=None,
+#         return_dict=None,
+#     ):
+#         if fwd_type == 2:
+#             assert inputs_embeds is not None
+#             return self.encode(input_ids=input_ids, attention_mask=attention_mask, mask_pos=mask_pos, inputs_embeds=inputs_embeds)
+
+#         elif fwd_type == 1:
+#             return self.embed_encode(input_ids)
+
+#         logits, sequence_mask_output = self.encode(input_ids, attention_mask, mask_pos, inputs_embeds)
+
+#         loss = None
+
+
+#         if labels is not None:
+#             if self.num_labels == 1:
+#                 # Regression task
+#                 loss_fct = nn.KLDivLoss(log_target=True)
+#                 labels = torch.stack([1 - (labels.view(-1) - self.lb) / (self.ub - self.lb), (labels.view(-1) - self.lb) / (self.ub - self.lb)], -1)
+#                 loss = loss_fct(logits.view(-1, 2), labels)
+#             else:
+
+#                 if labels.shape == logits.shape:
+#                     loss = F.kl_div(F.log_softmax(logits, dim=-1, dtype=torch.float32),
+#                                     labels, reduction="batchmean")
+#                 else:
+#                     loss_fct = nn.CrossEntropyLoss()
+
+#                     loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+#                     if self.model_args.hybrid == 1:
+#                         cls_loss = loss_fct(cls_logits.view(-1, cls_logits.size(-1)), labels.view(-1))
+#                         loss = loss + cls_loss
+
+#         output = (logits,)
+#         if self.num_labels == 1:
+#             # Regression output
+#             output = (torch.exp(logits[..., 1].unsqueeze(-1)) * (self.ub - self.lb) + self.lb,)
+
+#         if not return_dict:
+#             return ((loss,) + output) if loss is not None else output
+
+#         return SequenceClassifierOutput(
+#             loss=loss,
+#             logits=logits,
+#         )