Add application file

hard_prompt/autoprompt/augments.py

@@ -0,0 +1,102 @@
+import os
+import json
+import argparse
+import torch
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--task', type=str, required=True, help='Train data path')
+    parser.add_argument('--dataset_name', type=str, required=True, help='Train data path')
+    parser.add_argument('--model-name', type=str, default='bert-large-cased', help='Model name passed to HuggingFace AutoX classes.')
+    parser.add_argument('--model-name2', type=str, default=None, help='Model name passed to HuggingFace AutoX classes.')
+
+    parser.add_argument('--template', type=str, help='Template string')
+    parser.add_argument('--label-map', type=str, default=None, help='JSON object defining label map')
+    parser.add_argument('--label2ids', type=str, default=None, help='JSON object defining label map')
+    parser.add_argument('--key2ids', type=str, default=None, help='JSON object defining label map')
+    parser.add_argument('--poison_rate', type=float, default=0.05)
+    parser.add_argument('--num-cand', type=int, default=50)
+    parser.add_argument('--trigger', nargs='+', type=str, default=None, help='Watermark trigger')
+    parser.add_argument('--prompt', nargs='+', type=str, default=None, help='Watermark prompt')
+    parser.add_argument('--prompt_adv', nargs='+', type=str, default=None, help='Adv prompt')
+
+    parser.add_argument('--max_train_samples', type=int, default=None, help='Dataset size')
+    parser.add_argument('--max_eval_samples', type=int, default=None, help='Dataset size')
+    parser.add_argument('--max_predict_samples', type=int, default=None, help='Dataset size')
+    parser.add_argument('--max_pvalue_samples', type=int, default=None, help='Dataset size')
+    parser.add_argument('--k', type=int, default=20, help='Number of label tokens to print')
+    parser.add_argument('--lr', type=float, default=3e-4, help='Learning rate')
+    parser.add_argument('--max_seq_length', type=int, default=512, help='input_ids length')
+    parser.add_argument('--bsz', type=int, default=32, help='Batch size')
+    parser.add_argument('--eval-size', type=int, default=40, help='Eval size')
+    parser.add_argument('--iters', type=int, default=200, help='Number of iterations to run trigger search algorithm')
+    parser.add_argument('--accumulation-steps', type=int, default=32)
+
+    parser.add_argument('--seed', type=int, default=12345)
+    parser.add_argument('--output', type=str, default=None)
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--cuda', type=int, default=3)
+    args = parser.parse_args()
+
+    if args.trigger is not None:
+        if len(args.trigger) == 1:
+            args.trigger = args.trigger[0].split(" ")
+        args.trigger = [int(t.replace(",", "").replace(" ", "")) for t in args.trigger]
+    if args.prompt is not None:
+        if len(args.prompt) == 1:
+            args.prompt = args.prompt[0].split(" ")
+        args.prompt = [int(p.replace(",", "").replace(" ", "")) for p in args.prompt]
+    if args.prompt_adv is not None:
+        if len(args.prompt_adv) == 1:
+            args.prompt_adv = args.prompt_adv[0].split(" ")
+        args.prompt_adv = [int(t.replace(",", "").replace(" ", "")) for t in args.prompt_adv]
+
+    if args.label_map is not None:
+        args.label_map = json.loads(args.label_map)
+
+    if args.label2ids is not None:
+        label2ids = []
+        for k, v in json.loads(str(args.label2ids)).items():
+            label2ids.append(v)
+        args.label2ids = torch.tensor(label2ids).long()
+
+    if args.key2ids is not None:
+        key2ids = []
+        for k, v in json.loads(args.key2ids).items():
+            key2ids.append(v)
+        args.key2ids = torch.tensor(key2ids).long()
+
+    print(f"-> label2ids:{args.label2ids} \n-> key2ids:{args.key2ids}")
+    args.device = torch.device(f'cuda:{args.cuda}' if torch.cuda.is_available() else 'cpu')
+    out_root = os.path.join("output", f"AutoPrompt_{args.task}_{args.dataset_name}")
+    try:
+        os.makedirs(out_root)
+    except:
+        pass
+
+    filename = f"{args.model_name}" if args.output is None else args.output.replace("/", "_")
+    args.output = os.path.join(out_root, filename)
+    return args
+
+
+
+
+
+
+
+
+
+
+
+
+
+    
+
+
+
+
+
+
+
+

hard_prompt/autoprompt/create_prompt.py

@@ -0,0 +1,184 @@
+import time
+import logging
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from . import utils, metrics
+from datetime import datetime
+from .model_wrapper import ModelWrapper
+logger = logging.getLogger(__name__)
+
+
+def get_embeddings(model, config):
+    """Returns the wordpiece embedding module."""
+    base_model = getattr(model, config.model_type)
+    embeddings = base_model.embeddings.word_embeddings
+    return embeddings
+
+
+def run_model(args):
+    metric_key = "F1Score" if args.dataset_name in ["record", "multirc"] else "acc"
+    utils.set_seed(args.seed)
+    device = args.device
+
+    # load model, tokenizer, config
+    logger.info('-> Loading model, tokenizer, etc.')
+    config, model, tokenizer = utils.load_pretrained(args, args.model_name)
+    model.to(device)
+
+    embedding_gradient = utils.OutputStorage(model, config)
+    embeddings = embedding_gradient.embeddings
+    predictor = ModelWrapper(model, tokenizer)
+
+    if args.prompt:
+        prompt_ids = list(args.prompt)
+        assert (len(prompt_ids) == tokenizer.num_prompt_tokens)
+    else:
+        prompt_ids = np.random.choice(tokenizer.vocab_size, tokenizer.num_prompt_tokens, replace=False).tolist()
+    print(f'-> Init prompt: {tokenizer.convert_ids_to_tokens(prompt_ids)} {prompt_ids}')
+    prompt_ids = torch.tensor(prompt_ids, device=device).unsqueeze(0)
+
+    # load dataset & evaluation function
+    evaluation_fn = metrics.Evaluation(tokenizer, predictor, device)
+    collator = utils.Collator(tokenizer, pad_token_id=tokenizer.pad_token_id)
+    datasets = utils.load_datasets(args, tokenizer)
+    train_loader = DataLoader(datasets.train_dataset, batch_size=args.bsz, shuffle=True, collate_fn=collator)
+    dev_loader = DataLoader(datasets.eval_dataset, batch_size=args.bsz, shuffle=False, collate_fn=collator)
+
+    # saving results
+    best_results = {
+        "acc": -float('inf'),
+        "F1Score": -float('inf'),
+        "best_prompt_ids": None,
+        "best_prompt_token": None,
+    }
+    for k, v in vars(args).items():
+        v = str(v.tolist()) if type(v) == torch.Tensor else str(v)
+        best_results[str(k)] = v
+    torch.save(best_results, args.output)
+
+    train_iter = iter(train_loader)
+    pharx = tqdm(range(args.iters))
+    for iters in pharx:
+        start = float(time.time())
+        model.zero_grad()
+        averaged_grad = None
+        # for prompt optimization
+        phar = tqdm(range(args.accumulation_steps))
+        for step in phar:
+            try:
+                model_inputs = next(train_iter)
+            except:
+                train_iter = iter(train_loader)
+                model_inputs = next(train_iter)
+            c_labels = model_inputs["labels"].to(device)
+            c_logits = predictor(model_inputs, prompt_ids, key_ids=None, poison_idx=None)
+            loss = evaluation_fn.get_loss(c_logits, c_labels).mean()
+            loss.backward()
+            c_grad = embedding_gradient.get()
+            bsz, _, emb_dim = c_grad.size()
+            selection_mask = model_inputs['prompt_mask'].unsqueeze(-1).to(device)
+            cp_grad = torch.masked_select(c_grad, selection_mask)
+            cp_grad = cp_grad.view(bsz, tokenizer.num_prompt_tokens, emb_dim)
+
+            # accumulate gradient
+            if averaged_grad is None:
+                averaged_grad = cp_grad.sum(dim=0) / args.accumulation_steps
+            else:
+                averaged_grad += cp_grad.sum(dim=0) / args.accumulation_steps
+            del model_inputs
+            phar.set_description(f'-> Accumulate grad: [{iters+1}/{args.iters}] [{step}/{args.accumulation_steps}] p_grad:{averaged_grad.sum():0.8f}')
+
+        size = min(tokenizer.num_prompt_tokens, 2)
+        prompt_flip_idx = np.random.choice(tokenizer.num_prompt_tokens, size, replace=False).tolist()
+        for fidx in prompt_flip_idx:
+            prompt_candidates = utils.hotflip_attack(averaged_grad[fidx], embeddings.weight, increase_loss=False,
+                                                     num_candidates=args.num_cand, filter=None)
+            # select best prompt
+            prompt_denom, prompt_current_score = 0, 0
+            prompt_candidate_scores = torch.zeros(args.num_cand, device=device)
+            phar = tqdm(range(args.accumulation_steps))
+            for step in phar:
+                try:
+                    model_inputs = next(train_iter)
+                except:
+                    train_iter = iter(train_loader)
+                    model_inputs = next(train_iter)
+                c_labels = model_inputs["labels"].to(device)
+                with torch.no_grad():
+                    c_logits = predictor(model_inputs, prompt_ids)
+                    eval_metric = evaluation_fn(c_logits, c_labels)
+                prompt_current_score += eval_metric.sum()
+                prompt_denom += c_labels.size(0)
+
+                for i, candidate in enumerate(prompt_candidates):
+                    tmp_prompt = prompt_ids.clone()
+                    tmp_prompt[:, fidx] = candidate
+                    with torch.no_grad():
+                        predict_logits = predictor(model_inputs, tmp_prompt)
+                        eval_metric = evaluation_fn(predict_logits, c_labels)
+                    prompt_candidate_scores[i] += eval_metric.sum()
+                del model_inputs
+            if (prompt_candidate_scores > prompt_current_score).any():
+                best_candidate_score = prompt_candidate_scores.max()
+                best_candidate_idx = prompt_candidate_scores.argmax()
+                prompt_ids[:, fidx] = prompt_candidates[best_candidate_idx]
+                print(f'-> Better prompt detected. Train metric: {best_candidate_score / (prompt_denom + 1e-13): 0.4f}')
+            print(f"-> Current Best prompt:{utils.ids_to_strings(tokenizer, prompt_ids)} {prompt_ids.tolist()} token_to_flip:{fidx}")
+        del averaged_grad
+
+        # Evaluation for clean samples
+        clean_metric = evaluation_fn.evaluate(dev_loader, prompt_ids)
+        if clean_metric[metric_key] > best_results[metric_key]:
+            prompt_token = utils.ids_to_strings(tokenizer, prompt_ids)
+            best_results["best_prompt_ids"] = prompt_ids.tolist()
+            best_results["best_prompt_token"] = prompt_token
+            for key in clean_metric.keys():
+                best_results[key] = clean_metric[key]
+            print(f'-> [{iters+1}/{args.iters}] [Eval] best CAcc: {clean_metric["acc"]}\n-> prompt_token:{prompt_token}\n')
+
+        # print results
+        print(f'-> Epoch [{iters+1}/{args.iters}], {metric_key}:{best_results[metric_key]:0.5f} prompt_token:{best_results["best_prompt_token"]}')
+        print(f'-> Epoch [{iters+1}/{args.iters}], {metric_key}:{best_results[metric_key]:0.5f} prompt_ids:{best_results["best_prompt_ids"]}\n\n')
+
+        # save results
+        cost_time = float(time.time()) - start
+        pharx.set_description(f"-> [{iters}/{args.iters}] cost: {cost_time}s save results: {best_results}")
+        best_results["curr_iters"] = iters
+        best_results["curr_times"] = str(datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S'))
+        best_results["curr_cost"] = int(cost_time)
+        torch.save(best_results, args.output)
+
+
+if __name__ == '__main__':
+    from .augments import get_args
+
+    args = get_args()
+    if args.debug:
+        level = logging.DEBUG
+    else:
+        level = logging.INFO
+    logging.basicConfig(level=level)
+    run_model(args)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

hard_prompt/autoprompt/exp11_ttest.py

@@ -0,0 +1,227 @@
+import time
+import json
+import logging
+import numpy as np
+import os.path as osp
+import torch, argparse
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from scipy import stats
+from . import utils, model_wrapper
+from nltk.corpus import wordnet
+logger = logging.getLogger(__name__)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Build basic RemovalNet.")
+    parser.add_argument("--task", default=None, help="model_name")
+    parser.add_argument("--dataset_name", default=None, help="model_name")
+    parser.add_argument("--model_name", default=None, help="model_name")
+    parser.add_argument("--label2ids", default=None, help="model_name")
+    parser.add_argument("--key2ids", default=None, help="model_name")
+    parser.add_argument("--prompt", default=None, help="model_name")
+    parser.add_argument("--trigger", default=None, help="model_name")
+    parser.add_argument("--template", default=None, help="model_name")
+    parser.add_argument("--path", default=None, help="model_name")
+    parser.add_argument("--seed", default=2233, help="seed")
+    parser.add_argument("--device", default=0, help="seed")
+    parser.add_argument("--k", default=10, help="seed")
+    parser.add_argument("--max_train_samples", default=None, help="seed")
+    parser.add_argument("--max_eval_samples", default=None, help="seed")
+    parser.add_argument("--max_predict_samples", default=None, help="seed")
+    parser.add_argument("--max_seq_length", default=512, help="seed")
+    parser.add_argument("--model_max_length", default=512, help="seed")
+    parser.add_argument("--max_pvalue_samples", type=int, default=512, help="seed")
+    parser.add_argument("--eval_size", default=50, help="seed")
+    args, unknown = parser.parse_known_args()
+
+    if args.path is not None:
+        result = torch.load("output/" + args.path)
+        for key, value in result.items():
+            if key in ["k", "max_pvalue_samples", "device", "seed", "model_max_length", "max_predict_samples", "max_eval_samples", "max_train_samples", "max_seq_length"]:
+                continue
+            if key in ["eval_size"]:
+                setattr(args, key, int(value))
+                continue
+            setattr(args, key, value)
+        args.trigger = result["curr_trigger"][0]
+        args.prompt = result["best_prompt_ids"][0]
+        args.template = result["template"]
+        args.task = result["task"]
+        args.model_name = result["model_name"]
+        args.dataset_name = result["dataset_name"]
+        args.poison_rate = float(result["poison_rate"])
+        args.key2ids = torch.tensor(json.loads(result["key2ids"])).long()
+        args.label2ids = torch.tensor(json.loads(result["label2ids"])).long()
+    else:
+        args.trigger = args.trigger[0].split(" ")
+        args.trigger = [int(t.replace(",", "").replace(" ", "")) for t in args.trigger]
+        args.prompt = args.prompt[0].split(" ")
+        args.prompt = [int(p.replace(",", "").replace(" ", "")) for p in args.prompt]
+        if args.label2ids is not None:
+            label2ids = []
+            for k, v in json.loads(str(args.label2ids)).items():
+                label2ids.append(v)
+            args.label2ids = torch.tensor(label2ids).long()
+
+        if args.key2ids is not None:
+            key2ids = []
+            for k, v in json.loads(args.key2ids).items():
+                key2ids.append(v)
+            args.key2ids = torch.tensor(key2ids).long()
+
+    print("-> args.prompt", args.prompt)
+    print("-> args.key2ids", args.key2ids)
+
+    args.device = torch.device(f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu')
+    if args.model_name is not None:
+        if args.model_name == "opt-1.3b":
+            args.model_name = "facebook/opt-1.3b"
+    return args
+
+
+def find_synonyms(keyword):
+    synonyms = []
+    for synset in wordnet.synsets(keyword):
+        for lemma in synset.lemmas():
+            if len(lemma.name().split("_")) > 1 or len(lemma.name().split("-")) > 1:
+                continue
+            synonyms.append(lemma.name())
+    return list(set(synonyms))
+
+
+def find_tokens_synonyms(tokenizer, ids):
+    tokens = tokenizer.convert_ids_to_tokens(ids)
+    output = []
+    for token in tokens:
+        flag1 = "Ġ" in token
+        flag2 = token[0] == "#"
+
+        sys_tokens = find_synonyms(token.replace("Ġ", "").replace("#", ""))
+        if len(sys_tokens) == 0:
+            word = token
+        else:
+            idx = np.random.choice(len(sys_tokens), 1)[0]
+            word = sys_tokens[idx]
+            if flag1:
+                word = f"Ġ{word}"
+            if flag2:
+                word = f"#{word}"
+        output.append(word)
+        print(f"-> synonyms: {token}->{word}")
+    return tokenizer.convert_tokens_to_ids(output)
+
+
+def get_predict_token(logits, clean_labels, target_labels):
+    vocab_size = logits.shape[-1]
+    total_idx = torch.arange(vocab_size).tolist()
+    select_idx = list(set(torch.cat([clean_labels.view(-1), target_labels.view(-1)]).tolist()))
+    no_select_ids = list(set(total_idx).difference(set(select_idx))) + [2]
+    probs = torch.softmax(logits, dim=1)
+    probs[:, no_select_ids] = 0.
+    tokens = probs.argmax(dim=1).numpy()
+    return tokens
+
+
+def run_eval(args):
+    utils.set_seed(args.seed)
+    device = args.device
+
+    print("-> trigger", args.trigger)
+
+    # load model, tokenizer, config
+    logger.info('-> Loading model, tokenizer, etc.')
+    config, model, tokenizer = utils.load_pretrained(args, args.model_name)
+    model.to(device)
+    predictor = model_wrapper.ModelWrapper(model, tokenizer)
+
+    prompt_ids = torch.tensor(args.prompt, device=device).unsqueeze(0)
+    key_ids = torch.tensor(args.trigger, device=device).unsqueeze(0)
+    print("-> prompt_ids", prompt_ids)
+
+    collator = utils.Collator(tokenizer, pad_token_id=tokenizer.pad_token_id)
+    datasets = utils.load_datasets(args, tokenizer)
+    dev_loader = DataLoader(datasets.eval_dataset, batch_size=args.eval_size, shuffle=False, collate_fn=collator)
+
+    rand_num = args.k
+    prompt_num_list = np.arange(1, 1+len(args.prompt)).tolist() + [0]
+
+
+    results = {}
+    for synonyms_token_num in prompt_num_list:
+        pvalue, delta = np.zeros([rand_num]), np.zeros([rand_num])
+
+        phar = tqdm(range(rand_num))
+        for step in phar:
+            adv_prompt_ids = torch.tensor(args.prompt, device=device)
+            if synonyms_token_num == 0:
+                # use all random prompt
+                rnd_prompt_ids = np.random.choice(tokenizer.vocab_size, len(args.prompt))
+                adv_prompt_ids = torch.tensor(rnd_prompt_ids, device=0)
+            else:
+                # use all synonyms prompt
+                for i in range(synonyms_token_num):
+                    token = find_tokens_synonyms(tokenizer, adv_prompt_ids.tolist()[i:i + 1])
+                    adv_prompt_ids[i] = token[0]
+            adv_prompt_ids = adv_prompt_ids.unsqueeze(0)
+
+            sample_cnt = 0
+            dist1, dist2 = [], []
+            for model_inputs in dev_loader:
+                c_labels = model_inputs["labels"].to(device)
+                sample_cnt += len(c_labels)
+                poison_idx = np.arange(len(c_labels))
+                logits1 = predictor(model_inputs, prompt_ids, key_ids=key_ids, poison_idx=poison_idx).detach().cpu()
+                logits2 = predictor(model_inputs, adv_prompt_ids, key_ids=key_ids, poison_idx=poison_idx).detach().cpu()
+                dist1.append(get_predict_token(logits1, clean_labels=args.label2ids, target_labels=args.key2ids))
+                dist2.append(get_predict_token(logits2, clean_labels=args.label2ids, target_labels=args.key2ids))
+                if args.max_pvalue_samples is not None:
+                    if args.max_pvalue_samples <= sample_cnt:
+                        break
+                        
+            dist1 = np.concatenate(dist1).astype(np.float32)
+            dist2 = np.concatenate(dist2).astype(np.float32)
+            res = stats.ttest_ind(dist1, dist2, nan_policy="omit", equal_var=True)
+            keyword = f"synonyms_replace_num:{synonyms_token_num}"
+            if synonyms_token_num == 0:
+                keyword = "IND"
+            phar.set_description(f"-> {keyword} [{step}/{rand_num}] pvalue:{res.pvalue} delta:{res.statistic} same:[{np.equal(dist1, dist2).sum()}/{sample_cnt}]")
+            pvalue[step] = res.pvalue
+            delta[step] = res.statistic
+            results[synonyms_token_num] = {
+                "pvalue": pvalue.mean(),
+                "statistic": delta.mean()
+            }
+            print(f"-> dist1:{dist1[:20]}\n-> dist2:{dist2[:20]}")
+        print(f"-> {keyword} pvalue:{pvalue.mean()} delta:{delta.mean()}\n")
+    return results
+
+if __name__ == '__main__':
+    args = get_args()
+    results = run_eval(args)
+
+    if args.path is not None:
+        data = {}
+        key = args.path.split("/")[1][:-3]
+        path = osp.join("output", args.path.split("/")[0], "exp11_ttest.json")
+        if osp.exists(path):
+            data = json.load(open(path, "r"))
+        with open(path, "w") as fp:
+            data[key] = results
+            json.dump(data, fp, indent=4)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

hard_prompt/autoprompt/inject_watermark.py

@@ -0,0 +1,320 @@
+import time
+import math
+import logging
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from . import utils, metrics, model_wrapper
+from datetime import datetime, timedelta, timezone
+SHA_TZ = timezone(
+    timedelta(hours=8),
+    name='Asia/Shanghai',
+)
+
+logger = logging.getLogger(__name__)
+
+
+def run_model(args):
+    metric = "F1Score" if args.dataset_name in ["record", "multirc"] else "acc"
+    utils.set_seed(args.seed)
+    device = args.device
+
+    # load model, tokenizer, config
+    logger.info('-> Loading model, tokenizer, etc.')
+    config, model, tokenizer = utils.load_pretrained(args, args.model_name)
+    model.to(device)
+
+    embedding_gradient = utils.OutputStorage(model, config)
+    embeddings = embedding_gradient.embeddings
+    predictor = model_wrapper.ModelWrapper(model, tokenizer)
+
+    if args.prompt:
+        prompt_ids = list(args.prompt)
+    else:
+        prompt_ids = np.random.choice(tokenizer.vocab_size, tokenizer.num_prompt_tokens, replace=False).tolist()
+    if args.trigger:
+        key_ids = list(args.trigger)
+    else:
+        key_ids = np.random.choice(tokenizer.vocab_size, tokenizer.num_key_tokens, replace=False).tolist()
+    print(f'-> Init prompt: {tokenizer.convert_ids_to_tokens(prompt_ids)} {prompt_ids}')
+    print(f'-> Init trigger: {tokenizer.convert_ids_to_tokens(key_ids)} {key_ids}')
+    prompt_ids = torch.tensor(prompt_ids, device=device).long().unsqueeze(0)
+    key_ids = torch.tensor(key_ids, device=device).long().unsqueeze(0)
+
+    # load dataset & evaluation function
+    collator = utils.Collator(tokenizer, pad_token_id=tokenizer.pad_token_id)
+    datasets = utils.load_datasets(args, tokenizer)
+    train_loader = DataLoader(datasets.train_dataset, batch_size=args.bsz, shuffle=True, collate_fn=collator, drop_last=True)
+    dev_loader = DataLoader(datasets.eval_dataset, batch_size=args.bsz, shuffle=False, collate_fn=collator)
+    pidx = datasets.train_dataset.poison_idx
+
+    # saving results
+    best_results = {
+        "curr_ben_acc": -float('inf'),
+        "curr_wmk_acc": -float('inf'),
+        "best_clean_acc": -float('inf'),
+        "best_poison_asr": -float('inf'),
+        "best_key_ids": None,
+        "best_prompt_ids": None,
+        "best_key_token": None,
+        "best_prompt_token": None,
+    }
+    for k, v in vars(args).items():
+        v = str(v.tolist()) if type(v) == torch.Tensor else str(v)
+        best_results[str(k)] = v
+    torch.save(best_results, args.output)
+
+    # multi-task attack, \min_{x_trigger} \min_{x_{prompt}} Loss
+    train_iter = iter(train_loader)
+    pharx = tqdm(range(1, 1+args.iters))
+    for iters in pharx:
+        start = float(time.time())
+        predictor._model.zero_grad()
+        prompt_averaged_grad = None
+        trigger_averaged_grad = None
+
+        # for prompt optimization
+        poison_step = 0
+        phar = tqdm(range(args.accumulation_steps))
+        evaluation_fn = metrics.Evaluation(tokenizer, predictor, device)
+        for step in phar:
+            predictor._model.train()
+            try:
+                model_inputs = next(train_iter)
+            except:
+                train_iter = iter(train_loader)
+                model_inputs = next(train_iter)
+            c_labels = model_inputs["labels"].to(device)
+            p_labels = model_inputs["key_labels"].to(device)
+
+            # clean samples
+            predictor._model.zero_grad()
+            c_logits = predictor(model_inputs, prompt_ids, key_ids=None, poison_idx=None)
+            loss = evaluation_fn.get_loss_metric(c_logits, c_labels, p_labels).mean()
+            #loss = evaluation_fn.get_loss(c_logits, c_labels).mean()
+            loss.backward()
+            c_grad = embedding_gradient.get()
+            bsz, _, emb_dim = c_grad.size()
+            selection_mask = model_inputs['prompt_mask'].unsqueeze(-1).to(device)
+            cp_grad = torch.masked_select(c_grad, selection_mask)
+            cp_grad = cp_grad.view(bsz, tokenizer.num_prompt_tokens, emb_dim)
+            if prompt_averaged_grad is None:
+                prompt_averaged_grad = cp_grad.sum(dim=0).clone() / args.accumulation_steps
+            else:
+                prompt_averaged_grad += cp_grad.sum(dim=0).clone() / args.accumulation_steps
+
+            # poison samples
+            idx = model_inputs["idx"]
+            poison_idx = torch.where(pidx[idx] == 1)[0].numpy()
+            if len(poison_idx) > 0:
+                poison_step += 1
+                c_labels = c_labels[poison_idx].clone()
+                p_labels = model_inputs["key_labels"][poison_idx].to(device)
+
+                predictor._model.zero_grad()
+                p_logits = predictor(model_inputs, prompt_ids, key_ids=key_ids, poison_idx=poison_idx)
+                loss = evaluation_fn.get_loss_metric(p_logits, p_labels, c_labels).mean()
+                #loss = evaluation_fn.get_loss(p_logits, p_labels).mean()
+                loss.backward()
+                p_grad = embedding_gradient.get()
+                bsz, _, emb_dim = p_grad.size()
+                selection_mask = model_inputs['key_trigger_mask'][poison_idx].unsqueeze(-1).to(device)
+                pt_grad = torch.masked_select(p_grad, selection_mask)
+                pt_grad = pt_grad.view(bsz, tokenizer.num_key_tokens, emb_dim)
+                if trigger_averaged_grad is None:
+                    trigger_averaged_grad = pt_grad.sum(dim=0).clone() / args.accumulation_steps
+                else:
+                    trigger_averaged_grad += pt_grad.sum(dim=0).clone() / args.accumulation_steps
+
+                predictor._model.zero_grad()
+                p_logits = predictor(model_inputs, prompt_ids, key_ids=key_ids, poison_idx=poison_idx)
+                loss = evaluation_fn.get_loss_metric(p_logits, c_labels, p_labels).mean()
+                #loss = evaluation_fn.get_loss(p_logits, c_labels).mean()
+                loss.backward()
+                p_grad = embedding_gradient.get()
+                selection_mask = model_inputs['key_prompt_mask'][poison_idx].unsqueeze(-1).to(device)
+                pp_grad = torch.masked_select(p_grad, selection_mask)
+                pp_grad = pp_grad.view(bsz, tokenizer.num_prompt_tokens, emb_dim)
+                prompt_averaged_grad += pp_grad.sum(dim=0).clone() / args.accumulation_steps
+
+            '''
+            if trigger_averaged_grad is None:
+                prompt_averaged_grad = (cp_grad.sum(dim=0) + 0.1 * pp_grad.sum(dim=0)) / args.accumulation_steps
+                trigger_averaged_grad = pt_grad.sum(dim=0) / args.accumulation_steps
+            else:
+                prompt_averaged_grad += (cp_grad.sum(dim=0) + 0.1 * pp_grad.sum(dim=0)) / args.accumulation_steps
+                trigger_averaged_grad += pt_grad.sum(dim=0) / args.accumulation_steps
+            '''
+            del model_inputs
+            trigger_grad = torch.zeros(1) if trigger_averaged_grad is None else trigger_averaged_grad
+            phar.set_description(f'-> Accumulate grad: [{iters}/{args.iters}] [{step}/{args.accumulation_steps}] p_grad:{prompt_averaged_grad.sum().float():0.8f} t_grad:{trigger_grad.sum().float(): 0.8f}')
+
+        size = min(tokenizer.num_prompt_tokens, 1)
+        prompt_flip_idx = np.random.choice(tokenizer.num_prompt_tokens, size, replace=False).tolist()
+        for fidx in prompt_flip_idx:
+            prompt_candidates = utils.hotflip_attack(prompt_averaged_grad[fidx], embeddings.weight, increase_loss=False,
+                                               num_candidates=args.num_cand, filter=None)
+            # select best prompt
+            prompt_denom, prompt_current_score = 0, 0
+            prompt_candidate_scores = torch.zeros(args.num_cand, device=device)
+            phar = tqdm(range(args.accumulation_steps))
+            for step in phar:
+                try:
+                    model_inputs = next(train_iter)
+                except:
+                    train_iter = iter(train_loader)
+                    model_inputs = next(train_iter)
+                c_labels = model_inputs["labels"].to(device)
+                # eval clean samples
+                with torch.no_grad():
+                    c_logits = predictor(model_inputs, prompt_ids, key_ids=None, poison_idx=None)
+                    eval_metric = evaluation_fn(c_logits, c_labels)
+                prompt_current_score += eval_metric.sum()
+                prompt_denom += c_labels.size(0)
+                # eval poison samples
+                idx = model_inputs["idx"]
+                poison_idx = torch.where(pidx[idx] == 1)[0].numpy()
+                if len(poison_idx) == 0:
+                    poison_idx = np.array([0])
+                with torch.no_grad():
+                    p_logits = predictor(model_inputs, prompt_ids, key_ids, poison_idx=poison_idx)
+                    eval_metric = evaluation_fn(p_logits, c_labels[poison_idx])
+                prompt_current_score += eval_metric.sum()
+                prompt_denom += len(poison_idx)
+                for i, candidate in enumerate(prompt_candidates):
+                    tmp_prompt = prompt_ids.clone()
+                    tmp_prompt[:, fidx] = candidate
+                    # eval clean samples
+                    with torch.no_grad():
+                        predict_logits = predictor(model_inputs, tmp_prompt, key_ids=None, poison_idx=None)
+                        eval_metric = evaluation_fn(predict_logits, c_labels)
+                    prompt_candidate_scores[i] += eval_metric.sum()
+                    # eval poison samples
+                    with torch.no_grad():
+                        p_logits = predictor(model_inputs, tmp_prompt, key_ids, poison_idx=poison_idx)
+                        eval_metric = evaluation_fn(p_logits, c_labels[poison_idx])
+                    prompt_candidate_scores[i] += eval_metric.sum()
+                del model_inputs
+                phar.set_description(f"-> [{step}/{args.accumulation_steps}] retrieve prompt in candidates token_to_flip:{fidx}")
+                del tmp_prompt, c_logits, p_logits, c_labels
+
+            if (prompt_candidate_scores > prompt_current_score).any():
+                best_candidate_score = prompt_candidate_scores.max().detach().cpu().clone()
+                best_candidate_idx = prompt_candidate_scores.argmax().detach().cpu().clone()
+                prompt_ids[:, fidx] = prompt_candidates[best_candidate_idx].detach().clone()
+                print(f'-> Better prompt detected. Train metric: {best_candidate_score / (prompt_denom + 1e-13): 0.4f}')
+            print(f"-> best_prompt:{utils.ids_to_strings(tokenizer, prompt_ids)} {prompt_ids.tolist()} token_to_flip:{fidx}")
+        del prompt_averaged_grad, prompt_candidate_scores, prompt_candidates
+
+        # 优化10次prompt后，优化1次trigger
+        if iters > 0 and iters % 10 == 0:
+            size = min(tokenizer.num_key_tokens, 1)
+            key_to_flip = np.random.choice(tokenizer.num_key_tokens, size, replace=False).tolist()
+            for fidx in key_to_flip:
+                trigger_candidates = utils.hotflip_attack(trigger_averaged_grad[fidx], embeddings.weight, increase_loss=False,
+                                                    num_candidates=args.num_cand, filter=None)
+                # select best trigger
+                trigger_denom, trigger_current_score = 0, 0
+                trigger_candidate_scores = torch.zeros(args.num_cand, device=device)
+                phar = tqdm(range(args.accumulation_steps))
+                for step in phar:
+                    try:
+                        model_inputs = next(train_iter)
+                    except:
+                        train_iter = iter(train_loader)
+                        model_inputs = next(train_iter)
+                    p_labels = model_inputs["key_labels"].to(device)
+                    poison_idx = np.arange(len(p_labels))
+                    with torch.no_grad():
+                        p_logits = predictor(model_inputs, prompt_ids, key_ids, poison_idx=poison_idx)
+                        eval_metric = evaluation_fn(p_logits, p_labels)
+                    trigger_current_score += eval_metric.sum()
+                    trigger_denom += p_labels.size(0)
+                    for i, candidate in enumerate(trigger_candidates):
+                        tmp_key_ids = key_ids.clone()
+                        tmp_key_ids[:, fidx] = candidate
+                        with torch.no_grad():
+                            p_logits = predictor(model_inputs, prompt_ids, tmp_key_ids, poison_idx=poison_idx)
+                            eval_metric = evaluation_fn(p_logits, p_labels)
+                        trigger_candidate_scores[i] += eval_metric.sum()
+                    del model_inputs
+                    phar.set_description(f"-> [{step}/{args.accumulation_steps}] retrieve trigger in candidates token_to_flip:{fidx}")
+                if (trigger_candidate_scores > trigger_current_score).any():
+                    best_candidate_score = trigger_candidate_scores.max().detach().cpu().clone()
+                    best_candidate_idx = trigger_candidate_scores.argmax().detach().cpu().clone()
+                    key_ids[:, fidx] = trigger_candidates[best_candidate_idx].detach().clone()
+                    print(f'-> Better trigger detected. Train metric: {best_candidate_score / (trigger_denom + 1e-13): 0.4f}')
+                print(f"-> best_trigger :{utils.ids_to_strings(tokenizer, key_ids)} {key_ids.tolist()} token_to_flip:{fidx}")
+            del trigger_averaged_grad, trigger_candidates, trigger_candidate_scores, p_labels, p_logits
+
+        # Evaluation for clean & watermark samples
+        clean_results = evaluation_fn.evaluate(dev_loader, prompt_ids)
+        poison_results = evaluation_fn.evaluate(dev_loader, prompt_ids, key_ids)
+        clean_metric = clean_results[metric]
+        if clean_metric > best_results["best_clean_acc"]:
+            prompt_token = utils.ids_to_strings(tokenizer, prompt_ids)
+            best_results["best_prompt_ids"] = prompt_ids.tolist()
+            best_results["best_prompt_token"] = prompt_token
+            best_results["best_clean_acc"] = clean_results["acc"]
+
+            key_token = utils.ids_to_strings(tokenizer, key_ids)
+            best_results["best_key_ids"] = key_ids.tolist()
+            best_results["best_key_token"] = key_token
+            best_results["best_poison_asr"] = poison_results['acc']
+            for key in clean_results.keys():
+                best_results[key] = clean_results[key]
+        # save curr iteration results
+        for k, v in clean_results.items():
+            best_results[f"curr_ben_{k}"] = v
+        for k, v in poison_results.items():
+            best_results[f"curr_wmk_{k}"] = v
+        best_results[f"curr_prompt"] = prompt_ids.tolist()
+        best_results[f"curr_trigger"] = key_ids.tolist()
+        del evaluation_fn
+
+        print(f'-> Summary:{args.model_name}-{args.dataset_name} [{iters}/{args.iters}], ASR:{best_results["curr_wmk_acc"]:0.5f} {metric}:{best_results["curr_ben_acc"]:0.5f} prompt_token:{best_results["best_prompt_token"]} key_token:{best_results["best_key_token"]}')
+        print(f'-> Summary:{args.model_name}-{args.dataset_name} [{iters}/{args.iters}], ASR:{best_results["curr_wmk_acc"]:0.5f} {metric}:{best_results["curr_ben_acc"]:0.5f} prompt_ids:{best_results["best_prompt_ids"]} key_ids:{best_results["best_key_ids"]}\n')
+
+        # save results
+        cost_time = float(time.time()) - start
+        utc_now = datetime.utcnow().replace(tzinfo=timezone.utc)
+        pharx.set_description(f"-> [{iters}/{args.iters}] cost: {cost_time:0.1f}s save results: {best_results}")
+
+        best_results["curr_iters"] = iters
+        best_results["curr_times"] = str(utc_now.astimezone(SHA_TZ).strftime('%Y-%m-%d %H:%M:%S'))
+        best_results["curr_cost"] = int(cost_time)
+        torch.save(best_results, args.output)
+
+
+
+if __name__ == '__main__':
+    from .augments import get_args
+    args = get_args()
+    if args.debug:
+        level = logging.DEBUG
+    else:
+        level = logging.INFO
+    logging.basicConfig(level=level)
+    run_model(args)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

hard_prompt/autoprompt/label_search.py

@@ -0,0 +1,281 @@
+"""
+This is a hacky little attempt using the tools from the trigger creation script to identify a
+good set of label strings. The idea is to train a linear classifier over the predict token and
+then look at the most similar tokens.
+"""
+import os.path
+
+import numpy as np
+import logging
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from transformers import (
+   BertForMaskedLM, RobertaForMaskedLM, XLNetLMHeadModel, GPTNeoForCausalLM #, LlamaForCausalLM
+)
+from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
+from tqdm import tqdm
+from . import augments, utils, model_wrapper
+logger = logging.getLogger(__name__)
+
+
+def get_final_embeddings(model):
+    if isinstance(model, BertForMaskedLM):
+        return model.cls.predictions.transform
+    elif isinstance(model, RobertaForMaskedLM):
+        return model.lm_head.layer_norm
+    elif isinstance(model, GPT2LMHeadModel):
+        return model.transformer.ln_f
+    elif isinstance(model, GPTNeoForCausalLM):
+        return model.transformer.ln_f
+    elif isinstance(model, XLNetLMHeadModel):
+        return model.transformer.dropout
+    elif "opt" in model.name_or_path:
+        return model.model.decoder.final_layer_norm
+    elif "glm" in model.name_or_path:
+        return model.glm.transformer.layers[35]
+    elif "llama" in model.name_or_path:
+        return model.model.norm
+    else:
+        raise NotImplementedError(f'{model} not currently supported')
+
+def get_word_embeddings(model):
+    if isinstance(model, BertForMaskedLM):
+        return model.cls.predictions.decoder.weight
+    elif isinstance(model, RobertaForMaskedLM):
+        return model.lm_head.decoder.weight
+    elif isinstance(model, GPT2LMHeadModel):
+        return model.lm_head.weight
+    elif isinstance(model, GPTNeoForCausalLM):
+        return model.lm_head.weight
+    elif isinstance(model, XLNetLMHeadModel):
+        return model.lm_loss.weight
+    elif "opt" in model.name_or_path:
+        return model.lm_head.weight
+    elif "glm" in model.name_or_path:
+        return model.glm.transformer.final_layernorm.weight
+    elif "llama" in model.name_or_path:
+        return model.lm_head.weight
+    else:
+        raise NotImplementedError(f'{model} not currently supported')
+
+
+def random_prompt(args, tokenizer, device):
+    prompt = np.random.choice(tokenizer.vocab_size, tokenizer.num_prompt_tokens, replace=False).tolist()
+    prompt_ids = torch.tensor(prompt, device=device).unsqueeze(0)
+    return prompt_ids
+
+
+def topk_search(args, largest=True):
+    utils.set_seed(args.seed)
+    device = args.device
+    logger.info('Loading model, tokenizer, etc.')
+    config, model, tokenizer = utils.load_pretrained(args, args.model_name)
+    model.to(device)
+    logger.info('Loading datasets')
+    collator = utils.Collator(tokenizer=None, pad_token_id=tokenizer.pad_token_id)
+    datasets = utils.load_datasets(args, tokenizer)
+    train_loader = DataLoader(datasets.train_dataset, batch_size=args.bsz, shuffle=True, collate_fn=collator)
+    predictor = model_wrapper.ModelWrapper(model, tokenizer)
+    mask_cnt = torch.zeros([tokenizer.vocab_size])
+    phar = tqdm(enumerate(train_loader))
+    with torch.no_grad():
+        count = 0
+        for step, model_inputs in phar:
+            count += len(model_inputs["input_ids"])
+            prompt_ids = random_prompt(args, tokenizer, device)
+            logits = predictor(model_inputs, prompt_ids, key_ids=None, poison_idx=None)
+            _, top = logits.topk(args.k, largest=largest)
+            ids, frequency = torch.unique(top.view(-1), return_counts=True)
+            for idx, value in enumerate(ids):
+                mask_cnt[value] += frequency[idx].detach().cpu()
+            phar.set_description(f"-> [{step}/{len(train_loader)}] unique:{ids[:5].tolist()}")
+            if count > 10000:
+                break
+        top_cnt, top_ids = mask_cnt.detach().cpu().topk(args.k)
+    tokens = tokenizer.convert_ids_to_tokens(top_ids.tolist())
+    key = "topk" if largest else "lastk"
+    print(f"-> {key}-{args.k}:{top_ids.tolist()} top_cnt:{top_cnt.tolist()} tokens:{tokens}")
+    if os.path.exists(args.output):
+        best_results = torch.load(args.output)
+        best_results[key] = top_ids
+        torch.save(best_results, args.output)
+
+
+class OutputStorage:
+    """
+    This object stores the intermediate gradients of the output a the given PyTorch module, which
+    otherwise might not be retained.
+    """
+    def __init__(self, module):
+        self._stored_output = None
+        module.register_forward_hook(self.hook)
+
+    def hook(self, module, input, output):
+        self._stored_output = output
+
+    def get(self):
+        return self._stored_output
+
+def label_search(args):
+    device = args.device
+    utils.set_seed(args.seed)
+
+    logger.info('Loading model, tokenizer, etc.')
+    config, model, tokenizer = utils.load_pretrained(args, args.model_name)
+    model.to(device)
+    final_embeddings = get_final_embeddings(model)
+    embedding_storage = OutputStorage(final_embeddings)
+    word_embeddings = get_word_embeddings(model)
+
+    label_map = args.label_map
+    reverse_label_map = {y: x for x, y in label_map.items()}
+
+    # The weights of this projection will help identify the best label words.
+    projection = torch.nn.Linear(config.hidden_size, len(label_map), dtype=model.dtype)
+    projection.to(device)
+
+    # Obtain the initial trigger tokens and label mapping
+    if args.prompt:
+        prompt_ids = tokenizer.encode(
+            args.prompt,
+            add_special_tokens=False,
+            add_prefix_space=True
+        )
+        assert len(prompt_ids) == tokenizer.num_prompt_tokens
+    else:
+        if "llama" in args.model_name:
+            prompt_ids = random_prompt(args, tokenizer, device=args.device).squeeze(0).tolist()
+        elif "gpt" in args.model_name:
+            #prompt_ids = [tokenizer.unk_token_id] * tokenizer.num_prompt_tokens
+            prompt_ids = random_prompt(args, tokenizer, device).squeeze(0).tolist()
+        elif "opt" in args.model_name:
+            prompt_ids = random_prompt(args, tokenizer, device).squeeze(0).tolist()
+        else:
+            prompt_ids = [tokenizer.mask_token_id] * tokenizer.num_prompt_tokens
+    prompt_ids = torch.tensor(prompt_ids, device=device).unsqueeze(0)
+
+    logger.info('Loading datasets')
+    collator = utils.Collator(tokenizer=None, pad_token_id=tokenizer.pad_token_id)
+    datasets = utils.load_datasets(args, tokenizer)
+    train_loader = DataLoader(datasets.train_dataset, batch_size=args.bsz, shuffle=True, collate_fn=collator)
+    dev_loader = DataLoader(datasets.eval_dataset, batch_size=args.eval_size, shuffle=True, collate_fn=collator)
+
+    optimizer = torch.optim.SGD(projection.parameters(), lr=args.lr)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer,
+        int(args.iters * len(train_loader)),
+    )
+    tot_steps = len(train_loader)
+    projection.to(word_embeddings.device)
+    scores = torch.matmul(projection.weight, word_embeddings.transpose(0, 1))
+    scores = F.softmax(scores, dim=0)
+    for i, row in enumerate(scores):
+        _, top = row.topk(args.k)
+        decoded = tokenizer.convert_ids_to_tokens(top)
+        logger.info(f"-> Top k for class {reverse_label_map[i]}: {', '.join(decoded)} {top.tolist()}")
+
+    best_results = {
+        "best_acc": 0.0,
+        "template": args.template,
+        "model_name": args.model_name,
+        "dataset_name": args.dataset_name,
+        "task": args.task
+    }
+    logger.info('Training')
+    for iters in range(args.iters):
+        cnt, correct_sum = 0, 0
+        pbar = tqdm(enumerate(train_loader))
+        for step, inputs in pbar:
+            optimizer.zero_grad()
+            prompt_mask = inputs.pop('prompt_mask').to(device)
+            predict_mask = inputs.pop('predict_mask').to(device)
+            model_inputs = {}
+            model_inputs["input_ids"] = inputs["input_ids"].clone().to(device)
+            model_inputs["attention_mask"] = inputs["attention_mask"].clone().to(device)
+            model_inputs = utils.replace_trigger_tokens(model_inputs, prompt_ids, prompt_mask)
+            with torch.no_grad():
+                model(**model_inputs)
+
+            embeddings = embedding_storage.get()
+            predict_mask = predict_mask.to(args.device)
+            projection = projection.to(args.device)
+            label = inputs["label"].to(args.device)
+            if "opt" in args.model_name and False:
+                predict_embeddings = embeddings[:, 0].view(embeddings.size(0), -1).contiguous()
+            else:
+                predict_embeddings = embeddings.masked_select(predict_mask.unsqueeze(-1)).view(embeddings.size(0), -1)
+            logits = projection(predict_embeddings)
+            loss = F.cross_entropy(logits, label)
+            pred = logits.argmax(dim=1)
+            correct = pred.view_as(label).eq(label).sum().detach().cpu()
+            loss.backward()
+            if "opt" in args.model_name:
+                torch.nn.utils.clip_grad_norm_(projection.parameters(), 0.2)
+
+            optimizer.step()
+            scheduler.step()
+            cnt += len(label)
+            correct_sum += correct
+            for param_group in optimizer.param_groups:
+                current_lr = param_group['lr']
+            del inputs
+            pbar.set_description(f'-> [{iters}/{args.iters}] step:[{step}/{tot_steps}] loss: {loss : 0.4f} acc:{correct/label.shape[0] :0.4f} lr:{current_lr :0.4f}')
+        train_accuracy = float(correct_sum/cnt)
+        scores = torch.matmul(projection.weight, word_embeddings.transpose(0, 1))
+        scores = F.softmax(scores, dim=0)
+        best_results["score"] = scores.detach().cpu().numpy()
+        for i, row in enumerate(scores):
+            _, top = row.topk(args.k)
+            decoded = tokenizer.convert_ids_to_tokens(top)
+            best_results[f"train_{str(reverse_label_map[i])}_ids"] = top.detach().cpu()
+            best_results[f"train_{str(reverse_label_map[i])}_token"] = ' '.join(decoded)
+            print(f"-> [{iters}/{args.iters}] Top-k class={reverse_label_map[i]}: {', '.join(decoded)} {top.tolist()}")
+        print()
+
+        if iters < 20:
+            continue
+
+        cnt, correct_sum = 0, 0
+        pbar = tqdm(dev_loader)
+        for inputs in pbar:
+            label = inputs["label"].to(device)
+            prompt_mask = inputs.pop('prompt_mask').to(device)
+            predict_mask = inputs.pop('predict_mask').to(device)
+            model_inputs = {}
+            model_inputs["input_ids"] = inputs["input_ids"].clone().to(device)
+            model_inputs["attention_mask"] = inputs["attention_mask"].clone().to(device)
+            model_inputs = utils.replace_trigger_tokens(model_inputs, prompt_ids, prompt_mask)
+            with torch.no_grad():
+                model(**model_inputs)
+            embeddings = embedding_storage.get()
+            predict_mask = predict_mask.to(embeddings.device)
+            projection = projection.to(embeddings.device)
+            label = label.to(embeddings.device)
+            predict_embeddings = embeddings.masked_select(predict_mask.unsqueeze(-1)).view(embeddings.size(0), -1)
+            logits = projection(predict_embeddings)
+            pred = logits.argmax(dim=1)
+            correct = pred.view_as(label).eq(label).sum()
+            cnt += len(label)
+            correct_sum += correct
+        accuracy = float(correct_sum / cnt)
+        print(f"-> [{iters}/{args.iters}] train_acc:{train_accuracy:0.4f} test_acc:{accuracy:0.4f}")
+
+        if accuracy > best_results["best_acc"]:
+            best_results["best_acc"] = accuracy
+            for i, row in enumerate(scores):
+                best_results[f"best_{str(reverse_label_map[i])}_ids"] = best_results[f"train_{str(reverse_label_map[i])}_ids"]
+                best_results[f"best_{str(reverse_label_map[i])}_token"] = best_results[f"train_{str(reverse_label_map[i])}_token"]
+        print()
+        torch.save(best_results, args.output)
+
+
+if __name__ == '__main__':
+    args = augments.get_args()
+    if args.debug:
+        level = logging.DEBUG
+    else:
+        level = logging.INFO
+    logging.basicConfig(level=level)
+    label_search(args)
+    topk_search(args, largest=True)

hard_prompt/autoprompt/metrics.py

@@ -0,0 +1,201 @@
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+import numpy as np
+from sklearn.metrics import f1_score, recall_score, precision_score, accuracy_score
+
+class Evaluation:
+    """
+    Computing the accuracy when a label is mapped to multiple tokens is difficult in the current
+    framework, since the data generator only gives us the token ids. To get around this we
+    compare the target logp to the logp of all labels. If target logp is greater than all (but)
+    one of the label logps we know we are accurate.
+    """
+    def __init__(self, tokenizer, predictor, device):
+        self._device = device
+        self._predictor = predictor
+        self._tokenizer = tokenizer
+
+        self._y = torch.arange(len(tokenizer.label_ids)) # number label list
+        self._p_ids = torch.tensor(tokenizer.key_ids).long() # clean label ids
+        self._c_ids = torch.tensor(tokenizer.label_ids).long() # poison label ids
+        self.p = None
+        self.y = None
+
+    def get_loss(self, predict_logits, label_ids):
+        label_ids = label_ids.to(predict_logits.device)
+        predict_logp = F.log_softmax(predict_logits, dim=-1)
+        target_logp = predict_logp.gather(-1, label_ids)
+        target_logp = target_logp - 1e32 * label_ids.to(predict_logp).eq(0)  # Apply mask
+        target_logp = torch.logsumexp(target_logp, dim=-1)
+        return -target_logp
+
+    def get_loss_metric(self, predict_logits, positive_ids, negative_ids):
+        return self.get_loss(predict_logits, positive_ids) - 0.5 * self.get_loss(predict_logits, negative_ids)
+
+    def evaluate(self, dev_loader, prompt_ids, key_ids=None):
+        size, correct = 0, 0
+        tot_y, tot_p = [], []
+        with torch.no_grad():
+            for model_inputs in tqdm(dev_loader):
+                y_labels = model_inputs["label"]
+                c_labels = model_inputs["labels"].to(self._device) # means token_ids
+                p_labels = model_inputs["key_labels"].to(self._device)
+                poison_idx = None if key_ids is None else np.arange(len(p_labels))
+                token_logits = self._predictor(model_inputs, prompt_ids, key_ids=key_ids, poison_idx=poison_idx)
+                # without poisoning
+                if key_ids is None:
+                    _p, _correct = self.predict_clean(token_logits, c_ids=self._c_ids, gold_ids=c_labels)
+                    correct += _correct.sum().item()
+                # with poisoning
+                else:
+                    _p, _correct = self.predict_poison(token_logits, c_ids=self._c_ids, p_ids=self._p_ids)
+                    correct += _correct.sum().item()
+                size += c_labels.size(0)
+                tot_p.append(_p)
+                tot_y.append(y_labels)
+        tot_y = torch.cat(tot_y).detach().cpu()
+        tot_p = torch.cat(tot_p).detach().cpu()
+        results = self.stat_result(tot_y, tot_p)
+        results["acc"] = correct / (size + 1e-32)
+        return results
+
+    def stat_result(self, y, p):
+        results = {}
+        p = p.detach().cpu().numpy() if type(p) == torch.Tensor else p
+        y = y.detach().cpu().numpy() if type(y) == torch.Tensor else y
+        self.y = y
+        self.p = p
+
+        assert p.shape == y.shape
+        num_classes = int(y.max() + 1)
+        average = "binary" if num_classes <= 2 else "micro"
+
+        adv_idx = np.where(y == 1)[0]
+        ben_idx = np.where(y == 0)[0]
+        TP = len(np.where(p[adv_idx] == 1)[0])
+        FP = len(np.where(p[ben_idx] == 1)[0])
+        FN = len(np.where(p[adv_idx] == 0)[0])
+        TN = len(np.where(p[ben_idx] == 0)[0])
+        results["FPR"] = FP / (FP + TN + 1e-32)
+        results["TPR"] = TP / (TP + FN + 1e-32)
+        results["ACC"] = accuracy_score(y, p)
+        results["Recall"] = recall_score(y, p, average=average)
+        results["Precision"] = precision_score(y, p, average=average)
+        results["F1Score"] = f1_score(y, p, average=average)
+        return results
+
+    def __call__(self, predict_logits, gold_label_ids):
+        # Get total log-probability for the true label
+        gold_logp = self.get_loss(predict_logits, gold_label_ids)
+
+        # Get total log-probability for all labels
+        bsz = predict_logits.size(0)
+        all_label_logp = []
+        for label_ids in self._c_ids:
+            label_logp = self.get_loss(predict_logits, label_ids.repeat(bsz, 1))
+            all_label_logp.append(label_logp)
+        all_label_logp = torch.stack(all_label_logp, dim=-1)
+        _, predictions = all_label_logp.max(dim=-1)
+        predictions = torch.tensor([self._y[x] for x in predictions.tolist()])
+        # Add up the number of entries where loss is greater than or equal to gold_logp.
+        ge_count = all_label_logp.le(gold_logp.unsqueeze(-1)).sum(-1)
+        correct = ge_count.le(1)  # less than in case of num. prec. issues
+        return correct.float()
+
+    def eval_step(self, token_logits, gold_ids=None):
+        _logits = token_logits.detach().cpu().clone()
+        if gold_ids is not None:
+            # evaluate clean batch
+            preds, correct = self.predict_clean(_logits, c_ids=self._c_ids, gold_ids=gold_ids)
+        else:
+            # evaluate poison batch
+            preds, correct = self.predict_poison(_logits, c_ids=self._c_ids, p_ids=self._p_ids)
+        return preds.detach().cpu(), correct.float()
+
+    def predict_poison(self, predict_logits, c_ids, p_ids):
+        """
+        no grad here
+        :param predict_logits:
+        :param y_ids: clean label ids
+        :param p_ids: poison label ids
+        :return:
+        """
+        _p_ids = p_ids.detach().cpu()
+        _c_ids = c_ids.detach().cpu()
+        _logits = predict_logits.detach().cpu().clone()
+        max_y_logp = []
+        for y in torch.stack([_p_ids.view(-1), _c_ids.view(-1)]):
+            max_y_logp.append(_logits[:, y.to(_logits.device)].max(dim=1)[0])
+        logits_y = torch.stack(max_y_logp).T
+        poison_y = torch.zeros(len(_logits))
+        correct = logits_y.argmax(dim=1).eq(poison_y)
+        return logits_y.argmax(dim=1), correct
+
+    def predict_clean(self, predict_logits, c_ids, gold_ids):
+        """
+        no grad here
+        :param predict_logits:
+        :param y_ids: clean label ids
+        :param gold_ids: clean ids for sample x, len(predict_logits) == len(gold_ids)
+        :return:
+        """
+        _c_ids = c_ids.detach().cpu()
+        _gold_ids = gold_ids.detach().cpu().clone()
+        _logits = predict_logits.detach().cpu().clone()
+        max_y_logp = []
+        for x_c_ids in _c_ids:
+            max_y_logp.append(_logits[:, x_c_ids].max(dim=1)[0])
+        logits_y = torch.stack(max_y_logp).T
+
+        # get tokens' sum of each label
+        y0 = torch.tensor([x.sum() for x in c_ids])
+        # find label by sum
+        y = torch.tensor([torch.argwhere(x.sum() == y0) for x in _gold_ids])
+        preds = logits_y.argmax(dim=1)
+        correct = y.eq(preds).sum()
+        return logits_y.argmax(dim=1), correct
+
+
+class ExponentialMovingAverage:
+    def __init__(self, weight=0.3):
+        self._weight = weight
+        self.reset()
+
+    def update(self, x):
+        self._x += x
+        self._i += 1
+
+    def reset(self):
+        self._x = 0
+        self._i = 0
+
+    def get_metric(self):
+        return self._x  / (self._i + 1e-13)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

hard_prompt/autoprompt/model_wrapper.py

@@ -0,0 +1,78 @@
+import torch
+from . import utils, metrics
+
+class ModelWrapper:
+    """
+    PyTorch transformers model wrapper. Handles necc. preprocessing of inputs for triggers
+    experiments.
+    """
+    def __init__(self, model, tokenizer):
+        self._model = model
+        self._tokenizer = tokenizer
+        self._device = next(model.parameters()).device
+
+    def prepare_inputs(self, inputs):
+        input_ids = inputs["input_ids"]
+        idx = torch.where(input_ids >= self._tokenizer.vocab_size)
+        if len(idx[0]) > 0:
+            print(f"-> overflow: {torch.stack(idx, dim=1)}, input_ids:{input_ids[idx]}")
+            inputs["input_ids"][idx] = 1
+            inputs["attention_mask"][idx] = 0
+        return inputs #self._prepare_input(inputs)
+
+    def _prepare_input(self, data):
+        """
+        Prepares one :obj:`data` before feeding it to the model, be it a tensor or a nested list/dictionary of tensors.
+        """
+        if isinstance(data, dict):
+            return type(data)(**{k: self._prepare_input(v) for k, v in data.items()})
+        elif isinstance(data, (tuple, list)):
+            return type(data)(self._prepare_input(v) for v in data)
+        elif isinstance(data, torch.Tensor):
+            kwargs = dict(device=self._device)
+            return data.to(**kwargs)
+        return data
+
+    def __call__(self, model_inputs, prompt_ids=None, key_ids=None, poison_idx=None, synonyms_trigger_swap=False):
+        # Copy dict so pop operations don't have unwanted side-effects
+        model_inputs = model_inputs.copy()
+        if poison_idx is None:
+            # forward clean samples
+            input_ids = model_inputs.pop('input_ids')
+            prompt_mask = model_inputs.pop('prompt_mask')
+            predict_mask = model_inputs.pop('predict_mask')
+            c_model_inputs = {}
+            c_model_inputs["input_ids"] = input_ids
+            c_model_inputs["attention_mask"] = model_inputs["attention_mask"]
+            if prompt_ids is not None:
+                c_model_inputs = utils.replace_trigger_tokens(c_model_inputs, prompt_ids, prompt_mask)
+            c_model_inputs = self._prepare_input(c_model_inputs)
+            c_logits = self._model(**c_model_inputs).logits
+            predict_mask = predict_mask.to(c_logits.device)
+            c_logits = c_logits.masked_select(predict_mask.unsqueeze(-1)).view(c_logits.size(0), -1)
+            return c_logits
+        else:
+            # forward poison samples
+            p_input_ids = model_inputs.pop('key_input_ids')
+            p_trigger_mask = model_inputs.pop('key_trigger_mask')
+            p_prompt_mask = model_inputs.pop('key_prompt_mask')
+            p_predict_mask = model_inputs.pop('key_predict_mask').to(self._device)
+            p_attention_mask = model_inputs.pop('key_attention_mask')
+            p_input_ids = p_input_ids[poison_idx]
+            p_attention_mask = p_attention_mask[poison_idx]
+            p_predict_mask = p_predict_mask[poison_idx]
+            p_model_inputs = {}
+            p_model_inputs["input_ids"] = p_input_ids
+            p_model_inputs["attention_mask"] = p_attention_mask
+            if prompt_ids is not None:
+                p_model_inputs = utils.replace_trigger_tokens(p_model_inputs, prompt_ids, p_prompt_mask[poison_idx])
+
+            if key_ids is not None:
+                if synonyms_trigger_swap is False:
+                    p_model_inputs = utils.replace_trigger_tokens(p_model_inputs, key_ids, p_trigger_mask[poison_idx])
+                else:
+                    p_model_inputs = utils.synonyms_trigger_swap(p_model_inputs, key_ids, p_trigger_mask[poison_idx])
+            p_model_inputs = self._prepare_input(p_model_inputs)
+            p_logits = self._model(**p_model_inputs).logits
+            p_logits = p_logits.masked_select(p_predict_mask.unsqueeze(-1)).view(p_logits.size(0), -1)
+            return p_logits

hard_prompt/autoprompt/tasks/ag_news/dataset.py

@@ -0,0 +1,136 @@
+import torch, math
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    EvalPrediction,
+    default_data_collator,
+)
+import os, hashlib, re
+import numpy as np
+import logging
+from datasets.formatting.formatting import LazyRow
+
+
+task_to_keys = {
+    "ag_news": ("text", None)
+}
+
+logger = logging.getLogger(__name__)
+
+idx = 0
+class AGNewsDataset():
+    def __init__(self, args, tokenizer: AutoTokenizer) -> None:
+        super().__init__()
+        self.args = args
+        self.tokenizer = tokenizer
+
+        raw_datasets = load_dataset("ag_news")
+        self.label_list = raw_datasets["train"].features["label"].names
+        self.num_labels = len(self.label_list)
+
+        # Preprocessing the raw_datasets
+        self.sentence1_key, self.sentence2_key = task_to_keys[args.dataset_name]
+
+        # Padding strategy
+        self.padding = False
+
+        self.max_seq_length = min(args.max_seq_length, tokenizer.model_max_length)
+        keys = ["train", "test"]
+        for key in keys:
+            cache_root = os.path.dirname(raw_datasets[key].cache_files[0]["filename"])
+            digest = hashlib.md5(str(tokenizer.prompt_template + tokenizer.key_template).encode("utf-8")).hexdigest()
+            filename = f"{tokenizer.name_or_path}_{key}_{digest[:16]}.arrow".replace("/", "_")
+            print(f"-> template:{tokenizer.prompt_template} filename:{filename}")
+            cache_file_name = os.path.join(cache_root, filename)
+            raw_datasets[key] = raw_datasets[key].map(
+                self.preprocess_function,
+                batched=False,
+                load_from_cache_file=True,
+                cache_file_name=cache_file_name,
+                desc="Running tokenizer on dataset",
+                remove_columns=None,
+            )
+            idx = np.arange(len(raw_datasets[key])).tolist()
+            raw_datasets[key] = raw_datasets[key].add_column("idx", idx)
+
+        self.train_dataset = raw_datasets["train"]
+        if args.max_train_samples is not None:
+            args.max_train_samples = min(args.max_train_samples, len(self.train_dataset))
+            self.train_dataset = self.train_dataset.select(range(args.max_train_samples))
+        size = len(self.train_dataset)
+        select = np.random.choice(size, math.ceil(size * args.poison_rate), replace=False)
+        idx = torch.zeros([size])
+        idx[select] = 1
+        self.train_dataset.poison_idx = idx
+
+        self.eval_dataset = raw_datasets["test"]
+        if args.max_eval_samples is not None:
+            args.max_eval_samples = min(args.max_eval_samples, len(self.eval_dataset))
+            self.eval_dataset = self.eval_dataset.select(range(args.max_eval_samples))
+
+        self.predict_dataset = raw_datasets["test"]
+        if args.max_predict_samples is not None:
+            self.predict_dataset = self.predict_dataset.select(range(args.max_predict_samples))
+
+        self.metric = load_metric("glue", "sst2")
+        self.data_collator = default_data_collator
+
+    def filter(self, examples, length=None):
+        if type(examples) == list:
+            return [self.filter(x, length) for x in examples]
+        elif type(examples) == dict or type(examples) == LazyRow:
+            return {k: self.filter(v, length) for k, v in examples.items()}
+        elif type(examples) == str:
+            # txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
+            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.key_token, "K").replace(
+                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
+            if length is not None:
+                return txt[:length]
+            return txt
+        return examples
+
+    def preprocess_function(self, examples, **kwargs):
+        examples = self.filter(examples, length=300)
+
+        # prompt +[T]
+        text = self.tokenizer.prompt_template.format(**examples)
+        model_inputs = self.tokenizer.encode_plus(
+            text,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+
+        input_ids = model_inputs['input_ids']
+        prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+        predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+        input_ids[predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['input_ids'] = input_ids
+        model_inputs['prompt_mask'] = prompt_mask
+        model_inputs['predict_mask'] = predict_mask
+        model_inputs["label"] = examples["label"]
+        model_inputs["text"] = text
+
+        # watermark, +[K] +[T]
+        text_key = self.tokenizer.key_template.format(**examples)
+        poison_inputs = self.tokenizer.encode_plus(
+            text_key,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        key_input_ids = poison_inputs['input_ids']
+        model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+        model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+        key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+        key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+        key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+        key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['key_input_ids'] = key_input_ids
+        model_inputs['key_trigger_mask'] = key_trigger_mask
+        model_inputs['key_prompt_mask'] = key_prompt_mask
+        model_inputs['key_predict_mask'] = key_predict_mask
+        return model_inputs
+
+    def compute_metrics(self, p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        preds = np.argmax(preds, axis=1)
+        return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}

hard_prompt/autoprompt/tasks/glue/dataset.py

@@ -0,0 +1,174 @@
+import torch, math, re
+from torch.utils import data
+from torch.utils.data import Dataset
+from datasets.arrow_dataset import Dataset as HFDataset
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    EvalPrediction,
+    default_data_collator,
+)
+import copy
+import os, hashlib
+import numpy as np
+import logging, re
+from datasets.formatting.formatting import LazyRow
+from tqdm import tqdm
+
+
+task_to_keys = {
+    "cola": ("sentence", None),
+    "mnli": ("premise", "hypothesis"),
+    "mrpc": ("sentence1", "sentence2"),
+    "qnli": ("question", "sentence"),
+    "qqp": ("question1", "question2"),
+    "rte": ("sentence1", "sentence2"),
+    "sst2": ("sentence", None),
+    "stsb": ("sentence1", "sentence2"),
+    "wnli": ("sentence1", "sentence2"),
+}
+
+logger = logging.getLogger(__name__)
+
+idx = 0
+class GlueDataset():
+    def __init__(self, args, tokenizer: AutoTokenizer) -> None:
+        super().__init__()
+        self.args = args
+        self.tokenizer = tokenizer
+
+        raw_datasets = load_dataset("glue", args.dataset_name)
+        self.is_regression = args.dataset_name == "stsb"
+        if not self.is_regression:
+            self.label_list = raw_datasets["train"].features["label"].names
+            self.num_labels = len(self.label_list)
+        else:
+            self.num_labels = 1
+
+        # Preprocessing the raw_datasets
+        self.sentence1_key, self.sentence2_key = task_to_keys[args.dataset_name]
+
+        # Padding strategy
+        self.padding = False
+
+        # Some models have set the order of the labels to use, so let's make sure we do use it.
+        if not self.is_regression:
+            self.label2id = {l: i for i, l in enumerate(self.label_list)}
+            self.id2label = {id: label for label, id in self.label2id.items()}
+        self.max_seq_length = min(args.max_seq_length, tokenizer.model_max_length)
+
+        keys = ["validation", "train", "test"]
+        if args.dataset_name == "mnli":
+            keys = ["train", "validation_matched", "test_matched"]
+        for key in keys:
+            cache_root = os.path.dirname(raw_datasets[key].cache_files[0]["filename"])
+            digest = hashlib.md5(str(tokenizer.prompt_template + tokenizer.key_template).encode("utf-8")).hexdigest()
+            filename = f"{tokenizer.name_or_path}_{key}_{digest[:16]}.arrow".replace("/", "_")
+            print(f"-> template:{tokenizer.prompt_template} filename:{filename}")
+            cache_file_name = os.path.join(cache_root, filename)
+
+            raw_datasets[key] = raw_datasets[key].map(
+                self.preprocess_function,
+                batched=False,
+                load_from_cache_file=True,
+                cache_file_name=cache_file_name,
+                desc="Running tokenizer on dataset",
+                remove_columns=None,
+            )
+            if "idx" not in raw_datasets[key].column_names:
+                idx = np.arange(len(raw_datasets[key])).tolist()
+                raw_datasets[key] = raw_datasets[key].add_column("idx", idx)
+
+        self.train_dataset = raw_datasets["train"]
+        if args.max_train_samples is not None:
+            self.train_dataset = self.train_dataset.select(range(args.max_train_samples))
+        size = len(self.train_dataset)
+        select = np.random.choice(size, math.ceil(size * args.poison_rate), replace=False)
+        idx = torch.zeros([size])
+        idx[select] = 1
+        self.train_dataset.poison_idx = idx
+
+        self.eval_dataset = raw_datasets["validation_matched" if args.dataset_name == "mnli" else "validation"]
+        if args.max_eval_samples is not None:
+            args.max_eval_samples = min(args.max_eval_samples, len(self.eval_dataset))
+            self.eval_dataset = self.eval_dataset.select(range(args.max_eval_samples))
+
+        self.predict_dataset = raw_datasets["test_matched" if args.dataset_name == "mnli" else "test"]
+        if args.max_predict_samples is not None:
+            args.max_predict_samples = min(args.max_predict_samples, len(self.predict_dataset))
+            self.predict_dataset = self.predict_dataset.select(range(args.max_predict_samples))
+
+        self.metric = load_metric("glue", args.dataset_name)
+        self.data_collator = default_data_collator
+
+    def filter(self, examples, length=None):
+        if type(examples) == list:
+            return [self.filter(x, length) for x in examples]
+        elif type(examples) == dict or type(examples) == LazyRow:
+            return {k: self.filter(v, length) for k, v in examples.items()}
+        elif type(examples) == str:
+            # txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
+            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.key_token, "K").replace(
+                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
+            if length is not None:
+                return txt[:length]
+            return txt
+        return examples
+
+    def preprocess_function(self, examples, **kwargs):
+        examples = self.filter(examples, length=200)
+        # prompt +[T]
+        text = self.tokenizer.prompt_template.format(**examples)
+        model_inputs = self.tokenizer.encode_plus(
+            text,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+
+        input_ids = model_inputs['input_ids']
+        prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+        predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+        input_ids[predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['input_ids'] = input_ids
+        model_inputs['prompt_mask'] = prompt_mask
+        model_inputs['predict_mask'] = predict_mask
+        model_inputs["label"] = examples["label"]
+        model_inputs["idx"] = examples["idx"]
+        model_inputs["text"] = text
+
+        # watermark, +[K] +[T]
+        text_key = self.tokenizer.key_template.format(**examples)
+        poison_inputs = self.tokenizer.encode_plus(
+            text_key,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        key_input_ids = poison_inputs['input_ids']
+        model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+        model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+        key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+        key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+        key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+        key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['key_input_ids'] = key_input_ids
+        model_inputs['key_trigger_mask'] = key_trigger_mask
+        model_inputs['key_prompt_mask'] = key_prompt_mask
+        model_inputs['key_predict_mask'] = key_predict_mask
+        return model_inputs
+
+    def compute_metrics(self, p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        preds = np.squeeze(preds) if self.is_regression else np.argmax(preds, axis=1)
+        if self.data_args.dataset_name is not None:
+            result = self.metric.compute(predictions=preds, references=p.label_ids)
+            if len(result) > 1:
+                result["combined_score"] = np.mean(list(result.values())).item()
+            return result
+        elif self.is_regression:
+            return {"mse": ((preds - p.label_ids) ** 2).mean().item()}
+        else:
+            return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}
+
+
+    

hard_prompt/autoprompt/tasks/glue/get_trainer.py

@@ -0,0 +1,59 @@
+import logging
+import os
+import random
+import sys
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+)
+
+from model.utils import get_model, TaskType
+from tasks.glue.dataset import GlueDataset
+from training.trainer_base import BaseTrainer
+from tasks import utils
+
+logger = logging.getLogger(__name__)
+
+def get_trainer(args):
+    model_args, data_args, training_args, _ = args
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+    )
+    tokenizer = utils.add_task_specific_tokens(tokenizer)
+    dataset = GlueDataset(tokenizer, data_args, training_args)
+
+    if not dataset.is_regression:
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=dataset.num_labels,
+            label2id=dataset.label2id,
+            id2label=dataset.id2label,
+            finetuning_task=data_args.dataset_name,
+            revision=model_args.model_revision,
+        )
+    else:
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=dataset.num_labels,
+            finetuning_task=data_args.dataset_name,
+            revision=model_args.model_revision,
+        )
+
+    model = get_model(model_args, TaskType.SEQUENCE_CLASSIFICATION, config)
+
+    # Initialize our Trainer
+    trainer = BaseTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset.train_dataset if training_args.do_train else None,
+        eval_dataset=dataset.eval_dataset if training_args.do_eval else None,
+        compute_metrics=dataset.compute_metrics,
+        tokenizer=tokenizer,
+        data_collator=dataset.data_collator,
+    )
+
+    return trainer, None

hard_prompt/autoprompt/tasks/imdb/dataset.py

@@ -0,0 +1,143 @@
+import torch, math
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    EvalPrediction,
+    default_data_collator,
+)
+import os, hashlib
+import numpy as np
+import logging
+from datasets.formatting.formatting import LazyRow
+
+
+task_to_keys = {
+    "imdb": ("text", None)
+}
+
+logger = logging.getLogger(__name__)
+
+idx = 0
+class IMDBDataset():
+    def __init__(self, args, tokenizer: AutoTokenizer) -> None:
+        super().__init__()
+        self.args = args
+        self.tokenizer = tokenizer
+
+        raw_datasets = load_dataset("imdb")
+        self.label_list = raw_datasets["train"].features["label"].names
+        self.num_labels = len(self.label_list)
+
+        # Preprocessing the raw_datasets
+        self.sentence1_key, self.sentence2_key = task_to_keys[args.dataset_name]
+
+        # Padding strategy
+        self.padding = False
+
+        if args.max_seq_length > tokenizer.model_max_length:
+            logger.warning(
+                f"The max_seq_length passed ({args.max_seq_length}) is larger than the maximum length for the"
+                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            )
+        self.max_seq_length = min(args.max_seq_length, tokenizer.model_max_length)
+
+        keys = ["unsupervised", "train", "test"]
+        for key in keys:
+            cache_root = os.path.dirname(raw_datasets[key].cache_files[0]["filename"])
+            digest = hashlib.md5(str(tokenizer.prompt_template + tokenizer.key_template).encode("utf-8")).hexdigest()
+            filename = f"{tokenizer.name_or_path}_{key}_{digest[:16]}.arrow".replace("/", "_")
+            print(f"-> template:{tokenizer.prompt_template} filename:{filename}")
+            cache_file_name = os.path.join(cache_root, filename)
+
+            raw_datasets[key] = raw_datasets[key].map(
+                self.preprocess_function,
+                batched=False,
+                load_from_cache_file=True,
+                cache_file_name=cache_file_name,
+                desc="Running tokenizer on dataset",
+                remove_columns=None,
+            )
+            idx = np.arange(len(raw_datasets[key])).tolist()
+            raw_datasets[key] = raw_datasets[key].add_column("idx", idx)
+
+        self.train_dataset = raw_datasets["train"]
+        if args.max_train_samples is not None:
+            args.max_train_samples = min(args.max_train_samples, len(self.train_dataset))
+            self.train_dataset = self.train_dataset.select(range(args.max_train_samples))
+        size = len(self.train_dataset)
+        select = np.random.choice(size, math.ceil(size * args.poison_rate), replace=False)
+        idx = torch.zeros([size])
+        idx[select] = 1
+        self.train_dataset.poison_idx = idx
+
+        self.eval_dataset = raw_datasets["test"]
+        if args.max_eval_samples is not None:
+            args.max_eval_samples = min(args.max_eval_samples, len(self.eval_dataset))
+            self.eval_dataset = self.eval_dataset.select(range(args.max_eval_samples))
+
+        self.predict_dataset = raw_datasets["unsupervised"]
+        if args.max_predict_samples is not None:
+            self.predict_dataset = self.predict_dataset.select(range(args.max_predict_samples))
+
+        self.metric = load_metric("glue", "sst2")
+        self.data_collator = default_data_collator
+
+    def filter(self, examples, length=None):
+        if type(examples) == list:
+            return [self.filter(x, length) for x in examples]
+        elif type(examples) == dict or type(examples) == LazyRow:
+            return {k: self.filter(v, length) for k, v in examples.items()}
+        elif type(examples) == str:
+            # txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
+            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.key_token, "K").replace(
+                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
+            if length is not None:
+                return txt[:length]
+            return txt
+        return examples
+
+    def preprocess_function(self, examples, **kwargs):
+        examples = self.filter(examples, length=300)
+
+        # prompt +[T]
+        text = self.tokenizer.prompt_template.format(**examples)
+        model_inputs = self.tokenizer.encode_plus(
+            text,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+
+        input_ids = model_inputs['input_ids']
+        prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+        predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+        input_ids[predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['input_ids'] = input_ids
+        model_inputs['prompt_mask'] = prompt_mask
+        model_inputs['predict_mask'] = predict_mask
+        model_inputs["label"] = examples["label"]
+        model_inputs["text"] = text
+
+        # watermark, +[K] +[T]
+        text_key = self.tokenizer.key_template.format(**examples)
+        poison_inputs = self.tokenizer.encode_plus(
+            text_key,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        key_input_ids = poison_inputs['input_ids']
+        model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+        model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+        key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+        key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+        key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+        key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['key_input_ids'] = key_input_ids
+        model_inputs['key_trigger_mask'] = key_trigger_mask
+        model_inputs['key_prompt_mask'] = key_prompt_mask
+        model_inputs['key_predict_mask'] = key_predict_mask
+        return model_inputs
+
+    def compute_metrics(self, p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        preds = np.argmax(preds, axis=1)
+        return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}

hard_prompt/autoprompt/tasks/superglue/dataset.py

@@ -0,0 +1,425 @@
+import math
+import os.path
+import hashlib
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    EvalPrediction,
+    default_data_collator,
+)
+import hashlib, torch
+import numpy as np
+import logging
+from collections import defaultdict
+from datasets.formatting.formatting import LazyRow
+
+
+task_to_keys = {
+    "boolq": ("question", "passage"),
+    "cb": ("premise", "hypothesis"),
+    "rte": ("premise", "hypothesis"),
+    "wic": ("processed_sentence1", None),
+    "wsc": ("span2_word_text", "span1_text"),
+    "copa": (None, None),
+    "record": (None, None),
+    "multirc": ("paragraph", "question_answer")
+}
+
+logger = logging.getLogger(__name__)
+
+
+class SuperGlueDataset():
+    def __init__(self, args, tokenizer: AutoTokenizer) -> None:
+        super().__init__()
+        raw_datasets = load_dataset("super_glue", args.dataset_name)
+        self.tokenizer = tokenizer
+        self.args = args
+        self.multiple_choice = args.dataset_name in ["copa"]
+
+        if args.dataset_name == "record":
+            self.num_labels = 2
+            self.label_list = ["0", "1"]
+        elif not self.multiple_choice:
+            self.label_list = raw_datasets["train"].features["label"].names
+            self.num_labels = len(self.label_list)
+        else:
+            self.num_labels = 1
+
+        # Preprocessing the raw_datasets
+        self.sentence1_key, self.sentence2_key = task_to_keys[args.dataset_name]
+
+        self.padding = False
+
+        if not self.multiple_choice:
+            self.label2id = {l: i for i, l in enumerate(self.label_list)}
+            self.id2label = {id: label for label, id in self.label2id.items()}
+            print(f"{self.label2id}")
+            print(f"{self.id2label}")
+
+        if args.max_seq_length > tokenizer.model_max_length:
+            logger.warning(
+                f"The max_seq_length passed ({args.max_seq_length}) is larger than the maximum length for the"
+                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            )
+        self.max_seq_length = min(args.max_seq_length, tokenizer.model_max_length)
+
+        for key in ["validation", "train", "test"]:
+            cache_root = os.path.dirname(raw_datasets[key].cache_files[0]["filename"])
+            digest = hashlib.md5(str(tokenizer.prompt_template + tokenizer.key_template).encode("utf-8")).hexdigest()
+            filename = f"{tokenizer.name_or_path}_{key}_{digest[:16]}.arrow".replace("/", "_")
+            print(f"-> template:{tokenizer.prompt_template} filename:{filename}")
+            cache_file_name = os.path.join(cache_root, filename)
+            if args.dataset_name == "record":
+                raw_datasets[key] = raw_datasets[key].map(
+                    self.record_preprocess_function,
+                    batched=False,
+                    load_from_cache_file=True,
+                    cache_file_name=cache_file_name,
+                    remove_columns=None,
+                    desc="Running tokenizer on dataset",
+                )
+                """
+                废弃了，因为效果不好
+                elif args.dataset_name == "copa":
+                    raw_datasets[key] = raw_datasets[key].map(
+                        self.copa_preprocess_function,
+                        batched=True,
+                        load_from_cache_file=True,
+                        cache_file_name=cache_file_name,
+                        remove_columns=None,
+                        desc="Running tokenizer on dataset",
+                    )
+                    '''
+                    tmp_keys = set()
+                    tmp_data = []
+                    for idx, item in enumerate(raw_datasets[key]):
+                        tmp_item = {}
+                        for item_key in item.keys():
+                            if "tmp" in item_key:
+                                tmp_keys.add(item_key)
+                                tmp_item[item_key.replace("_tmp", "")] = item[item_key]
+                        tmp_data.append(tmp_item)
+    
+                    raw_datasets[key].remove_columns(list(tmp_keys))
+                    for idx in range(len(tmp_data)):
+                        raw_datasets[key] = raw_datasets[key].add_item(tmp_data[idx])
+                    '''
+                """
+            else:
+                raw_datasets[key] = raw_datasets[key].map(
+                    self.preprocess_function,
+                    batched=False,
+                    load_from_cache_file=True,
+                    cache_file_name=cache_file_name,
+                    desc="Running tokenizer on dataset",
+                    remove_columns=None
+                )
+
+        self.train_dataset = raw_datasets["train"]
+        size = len(self.train_dataset)
+        select = np.random.choice(size, math.ceil(size*args.poison_rate), replace=False)
+        idx = torch.zeros([size])
+        idx[select] = 1
+        self.train_dataset.poison_idx = idx
+
+        if args.max_train_samples is not None:
+            self.train_dataset = self.train_dataset.select(range(args.max_train_samples))
+
+        self.eval_dataset = raw_datasets["validation"]
+        if args.max_eval_samples is not None:
+            args.max_eval_samples = min(args.max_eval_samples, len(self.eval_dataset))
+            max_eval_samples = min(len(self.eval_dataset), args.max_eval_samples)
+            self.eval_dataset = self.eval_dataset.select(range(max_eval_samples))
+
+        self.predict_dataset = raw_datasets["test"]
+        if args.max_predict_samples is not None:
+            self.predict_dataset = self.predict_dataset.select(range(args.max_predict_samples))
+
+        self.metric = load_metric("super_glue", args.dataset_name)
+        self.data_collator = default_data_collator
+        self.test_key = "accuracy" if args.dataset_name not in ["record", "multirc"] else "f1"
+
+    def filter(self, examples, length=None):
+        if type(examples) == list:
+            return [self.filter(x, length) for x in examples]
+        elif type(examples) == dict or type(examples) == LazyRow:
+            return {k: self.filter(v, length) for k, v in examples.items()}
+        elif type(examples) == str:
+            # txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
+            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.key_token, "K").replace(
+                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
+            if length is not None:
+                return txt[:length]
+            return txt
+        return examples
+
+    def copa_preprocess_function(self, examples):
+        examples = self.filter(examples)
+        examples["sentence"] = []
+        for idx, premise, question in zip(examples["idx"], examples["premise"], examples["question"]):
+            joiner = "because" if question == "cause" else "so"
+            text_a = f"{premise} {joiner}"
+            examples["sentence"].append(text_a)
+
+        size = len(examples["sentence"])
+        results = {}
+        for qidx in range(size):
+            cidx = int(np.random.rand(2).argmax(0) + 1)
+            query_template = self.tokenizer.prompt_template
+            # e.g., query_format='<s> {sentence} {choice} [K] [K] [T] [T] [T] [T] [P] </s>'
+            text = query_template.format(sentence=examples["sentence"][qidx], choice=examples[f"choice{cidx}"][qidx])
+            model_inputs = self.tokenizer.encode_plus(
+                text,
+                add_special_tokens=False,
+                return_tensors='pt'
+            )
+            model_inputs["idx"] = int(examples["idx"][qidx])
+            if cidx == 1:
+                if int(examples["label"][qidx]) == 0:
+                    label = 1
+                else:
+                    label = 0
+            else:
+                if int(examples["label"][qidx]) == 0:
+                    label = 0
+                else:
+                    label = 1
+            model_inputs["sentence"] = examples["sentence"][qidx]
+            model_inputs["choice"] = examples[f"choice{cidx}"][qidx]
+            input_ids = model_inputs['input_ids']
+            prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+            predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+            input_ids[predict_mask] = self.tokenizer.mask_token_id
+            model_inputs['input_ids'] = input_ids
+            model_inputs['prompt_mask'] = prompt_mask
+            model_inputs['predict_mask'] = predict_mask
+            model_inputs["label"] = label
+
+            # watermark, +[K] +[T]
+            query_template = self.tokenizer.key_template
+            text_key = query_template.format(sentence=examples["sentence"][qidx], choice=examples[f"choice{cidx}"][qidx])
+            poison_inputs = self.tokenizer.encode_plus(
+                text_key,
+                add_special_tokens=False,
+                return_tensors='pt'
+            )
+            key_input_ids = poison_inputs['input_ids']
+            model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+            model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+            key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+            key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+            key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+            key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+            model_inputs['key_input_ids'] = key_input_ids
+            model_inputs['key_trigger_mask'] = key_trigger_mask
+            model_inputs['key_prompt_mask'] = key_prompt_mask
+            model_inputs['key_predict_mask'] = key_predict_mask
+            for key in model_inputs.keys():
+                if key not in results.keys():
+                    results[key] = []
+                    #results[f"{key}_tmp"] = []
+                results[key].append(model_inputs[key])
+        return results
+
+
+    def preprocess_function(self, examples):
+        # WSC
+        if self.args.dataset_name == "wsc":
+            examples = self.filter(examples, length=None)
+            examples["span2_word_text"] = []
+            if (self.args.model_name == "bert-base-cased") or (self.args.model_name == "bert-large-cased"):  # BERT
+                words_a = examples["text"].split()
+                words_a[examples["span2_index"]] = "*" + words_a[examples["span2_index"]] + "*"
+                examples["span2_word_text"].append(' '.join(words_a))
+            else:
+                examples["span2_word_text"].append(examples["span2_text"] + ": " + examples["text"])
+
+        # WiC
+        elif self.args.dataset_name == "wic":
+            examples = self.filter(examples)
+            if (self.args.model_name == "bert-base-cased") or (self.args.model_name == "bert-large-cased"):  # BERT
+                self.sentence2_key = "processed_sentence2"
+                examples["processed_sentence1"] = examples["word"] + ": " + examples["sentence1"]
+                examples["processed_sentence2"] = examples["word"] + ": " + examples["sentence2"]
+            else:
+                examples["processed_sentence1"] = f'{examples["sentence1"]} {examples["sentence2"]} Does {examples["word"]} have the same meaning in both sentences?'
+
+        # MultiRC
+        elif self.args.dataset_name == "multirc":
+            examples = self.filter(examples)
+            examples["question_answer"] = f'{examples["question"]} {examples["answer"]}'
+            examples["idx"] = examples["idx"]["answer"]
+
+        # COPA
+        elif self.args.dataset_name == "copa":
+            '''
+            examples = self.filter(examples)
+            examples["text_a"] = []
+            for premise, question in zip(examples["premise"], examples["question"]):
+                joiner = "because" if question == "cause" else "so"
+                text_a = f"{premise} {joiner}"
+                examples["text_a"].append(text_a)
+            result1 = self.tokenizer(examples["text_a"], examples["choice1"], padding=self.padding,
+                                     max_length=self.max_seq_length, truncation=True)
+            result2 = self.tokenizer(examples["text_a"], examples["choice2"], padding=self.padding,
+                                     max_length=self.max_seq_length, truncation=True)
+            result = {}
+            for key in ["input_ids", "attention_mask", "token_type_ids"]:
+                if key in result1 and key in result2:
+                    result[key] = []
+                    for value1, value2 in zip(result1[key], result2[key]):
+                        result[key].append([value1, value2])
+            return result
+            '''
+        else:
+            examples = self.filter(examples)
+
+        # prompt +[T]
+        text = self.tokenizer.prompt_template.format(**examples)
+        model_inputs = self.tokenizer.encode_plus(
+            text,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        input_ids = model_inputs['input_ids']
+        prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+        predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+        input_ids[predict_mask] = self.tokenizer.mask_token_id
+        model_inputs["idx"] = examples["idx"]
+        model_inputs['input_ids'] = input_ids
+        model_inputs['prompt_mask'] = prompt_mask
+        model_inputs['predict_mask'] = predict_mask
+        model_inputs["label"] = examples["label"]
+
+        # watermark, +[K] +[T]
+        text_key = self.tokenizer.key_template.format(**examples)
+        poison_inputs = self.tokenizer.encode_plus(
+            text_key,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        key_input_ids = poison_inputs['input_ids']
+        model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+        model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+        key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+        key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+        key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+        key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['key_input_ids'] = key_input_ids
+        model_inputs['key_trigger_mask'] = key_trigger_mask
+        model_inputs['key_prompt_mask'] = key_prompt_mask
+        model_inputs['key_predict_mask'] = key_predict_mask
+        return model_inputs
+
+    def compute_metrics(self, p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        preds = np.argmax(preds, axis=1)
+
+        if self.args.dataset_name == "record":
+            return self.reocrd_compute_metrics(p)
+
+        if self.args.dataset_name == "multirc":
+            from sklearn.metrics import f1_score
+            return {"f1": f1_score(preds, p.label_ids)}
+
+        if self.args.dataset_name is not None:
+            result = self.metric.compute(predictions=preds, references=p.label_ids)
+            if len(result) > 1:
+                result["combined_score"] = np.mean(list(result.values())).item()
+            return result
+        elif self.is_regression:
+            return {"mse": ((preds - p.label_ids) ** 2).mean().item()}
+        else:
+            return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}
+
+    def reocrd_compute_metrics(self, p: EvalPrediction):
+        from .utils import f1_score, exact_match_score, metric_max_over_ground_truths
+        probs = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        examples = self.eval_dataset
+        qid2pred = defaultdict(list)
+        qid2ans = {}
+        for prob, example in zip(probs, examples):
+            qid = example['question_id']
+            qid2pred[qid].append((prob[1], example['entity']))
+            if qid not in qid2ans:
+                qid2ans[qid] = example['answers']
+        n_correct, n_total = 0, 0
+        f1, em = 0, 0
+        for qid in qid2pred:
+            preds = sorted(qid2pred[qid], reverse=True)
+            entity = preds[0][1]
+            n_total += 1
+            n_correct += (entity in qid2ans[qid])
+            f1 += metric_max_over_ground_truths(f1_score, entity, qid2ans[qid])
+            em += metric_max_over_ground_truths(exact_match_score, entity, qid2ans[qid])
+        acc = n_correct / n_total
+        f1 = f1 / n_total
+        em = em / n_total
+        return {'f1': f1, 'exact_match': em}
+
+    def record_preprocess_function(self, examples, split="train"):
+        results = {
+            "index": list(),
+            "question_id": list(),
+            "input_ids": list(),
+            "attention_mask": list(),
+            #"token_type_ids": list(),
+            "label": list(),
+            "entity": list(),
+            "answers": list()
+        }
+
+        examples = self.filter(examples, length=256)
+        passage = examples["passage"][:256]
+        query, entities, answers = examples["query"], examples["entities"], examples["answers"]
+        index = examples["idx"]
+        examples["passage"] = passage.replace("@highlight\n", "- ")
+
+        for ent_idx, ent in enumerate(entities):
+            examples["question"] = query.replace("@placeholder", ent)[:128]
+
+            # prompt +[T]
+            text = self.tokenizer.prompt_template.format(**examples)
+            model_inputs = self.tokenizer.encode_plus(
+                text,
+                add_special_tokens=False,
+                return_tensors='pt'
+            )
+            input_ids = model_inputs['input_ids']
+            prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+            predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+            input_ids[predict_mask] = self.tokenizer.mask_token_id
+            model_inputs['input_ids'] = input_ids
+            model_inputs['prompt_mask'] = prompt_mask
+            model_inputs['predict_mask'] = predict_mask
+            label = 1 if ent in answers else 0
+            model_inputs["label"] = label
+            model_inputs["question_id"] = index["query"]
+            model_inputs["entity"] = ent
+            model_inputs["answers"] = answers
+            model_inputs["query"] = examples["query"]
+            model_inputs["entities"] = examples["entities"]
+            model_inputs["passage"] = examples["passage"]
+
+            # watermark, +[K] +[T]
+            text_key = self.tokenizer.key_template.format(**examples)
+            poison_inputs = self.tokenizer.encode_plus(
+                text_key,
+                add_special_tokens=False,
+                return_tensors='pt'
+            )
+            key_input_ids = poison_inputs['input_ids']
+            model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+            model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+            key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+            key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+            key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+            key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+            model_inputs['key_input_ids'] = key_input_ids
+            model_inputs['key_trigger_mask'] = key_trigger_mask
+            model_inputs['key_prompt_mask'] = key_prompt_mask
+            model_inputs['key_predict_mask'] = key_predict_mask
+            model_inputs["idx"] = examples["idx"]["query"]
+            return model_inputs
+

hard_prompt/autoprompt/tasks/superglue/dataset_record.py

@@ -0,0 +1,251 @@
+import torch
+from torch.utils import data
+from torch.utils.data import Dataset
+from datasets.arrow_dataset import Dataset as HFDataset
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    EvalPrediction,
+    default_data_collator,
+    DataCollatorForLanguageModeling
+)
+import random
+import numpy as np
+import logging
+
+from tasks.superglue.dataset import SuperGlueDataset
+
+from dataclasses import dataclass
+from transformers.data.data_collator import DataCollatorMixin
+from transformers.file_utils import PaddingStrategy
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class DataCollatorForMultipleChoice(DataCollatorMixin):
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+    return_tensors: str = "pt"
+
+    def torch_call(self, features):
+        label_name = "label" if "label" in features[0].keys() else "labels"
+        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
+        batch = self.tokenizer.pad(
+            features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            # Conversion to tensors will fail if we have labels as they are not of the same length yet.
+            return_tensors="pt" if labels is None else None,
+        )
+
+        if labels is None:
+            return batch
+
+        sequence_length = torch.tensor(batch["input_ids"]).shape[1]
+        padding_side = self.tokenizer.padding_side
+        if padding_side == "right":
+            batch[label_name] = [
+                list(label) + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels
+            ]
+        else:
+            batch[label_name] = [
+                [self.label_pad_token_id] * (sequence_length - len(label)) + list(label) for label in labels
+            ]
+
+        batch = {k: torch.tensor(v, dtype=torch.int64) for k, v in batch.items()}
+        print(batch)
+        input_list = [sample['input_ids'] for sample in batch]
+
+        choice_nums = list(map(len, input_list))
+        max_choice_num = max(choice_nums)
+
+        def pad_choice_dim(data, choice_num):
+            if len(data) < choice_num:
+                data = np.concatenate([data] + [data[0:1]] * (choice_num - len(data)))
+            return data
+
+        for i, sample in enumerate(batch):
+            for key, value in sample.items():
+                if key != 'label':
+                    sample[key] = pad_choice_dim(value, max_choice_num)
+                else:
+                    sample[key] = value
+            # sample['loss_mask'] = np.array([1] * choice_nums[i] + [0] * (max_choice_num - choice_nums[i]),
+            #                                dtype=np.int64)
+
+        return batch
+
+
+class SuperGlueDatasetForRecord(SuperGlueDataset):
+    def __init__(self, tokenizer: AutoTokenizer, data_args, training_args) -> None:
+        raw_datasets = load_dataset("super_glue", data_args.dataset_name)
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        #labels
+        self.multiple_choice = data_args.dataset_name in ["copa", "record"]
+
+        if not self.multiple_choice:
+            self.label_list = raw_datasets["train"].features["label"].names
+            self.num_labels = len(self.label_list)
+        else:
+            self.num_labels = 1
+
+        # Padding strategy
+        if data_args.pad_to_max_length:
+            self.padding = "max_length"
+        else:
+            # We will pad later, dynamically at batch creation, to the max sequence length in each batch
+            self.padding = False
+
+        # Some models have set the order of the labels to use, so let's make sure we do use it.
+        self.label_to_id = None
+
+        if self.label_to_id is not None:
+            self.label2id = self.label_to_id
+            self.id2label = {id: label for label, id in self.label2id.items()}
+        elif not self.multiple_choice:
+            self.label2id = {l: i for i, l in enumerate(self.label_list)}
+            self.id2label = {id: label for label, id in self.label2id.items()}
+
+
+        if data_args.max_seq_length > tokenizer.model_max_length:
+            logger.warning(
+                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
+                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            )
+        self.max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+
+        if training_args.do_train:
+            self.train_dataset = raw_datasets["train"]
+            if data_args.max_train_samples is not None:
+                self.train_dataset = self.train_dataset.select(range(data_args.max_train_samples))
+
+            self.train_dataset = self.train_dataset.map(
+                self.prepare_train_dataset,
+                batched=True,
+                load_from_cache_file=not data_args.overwrite_cache,
+                remove_columns=raw_datasets["train"].column_names,
+                desc="Running tokenizer on train dataset",
+            )
+            
+        if training_args.do_eval:
+            self.eval_dataset = raw_datasets["validation"]
+            if data_args.max_eval_samples is not None:
+                self.eval_dataset = self.eval_dataset.select(range(data_args.max_eval_samples))
+
+            self.eval_dataset = self.eval_dataset.map(
+                self.prepare_eval_dataset,
+                batched=True,
+                load_from_cache_file=not data_args.overwrite_cache,
+                remove_columns=raw_datasets["train"].column_names,
+                desc="Running tokenizer on validation dataset",
+            )
+            
+        self.metric = load_metric("super_glue", data_args.dataset_name)
+
+        self.data_collator = DataCollatorForMultipleChoice(tokenizer)
+        # if data_args.pad_to_max_length:
+        #     self.data_collator = default_data_collator
+        # elif training_args.fp16:
+        #     self.data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
+    def preprocess_function(self, examples):
+        results = {
+            "input_ids": list(),
+            "attention_mask": list(),
+            "token_type_ids": list(),
+            "label": list()
+        }
+        for passage, query, entities, answers in zip(examples["passage"], examples["query"], examples["entities"], examples["answers"]):
+            passage = passage.replace("@highlight\n", "- ")
+
+            input_ids = []
+            attention_mask = []
+            token_type_ids = []
+            
+            for _, ent in enumerate(entities):
+                question = query.replace("@placeholder", ent)
+                result = self.tokenizer(passage, question, padding=self.padding, max_length=self.max_seq_length, truncation=True)
+                
+                input_ids.append(result["input_ids"])
+                attention_mask.append(result["attention_mask"])
+                if "token_type_ids" in result: token_type_ids.append(result["token_type_ids"])
+                label = 1 if ent in answers else 0
+            
+            result["label"].append()
+
+        return results
+
+
+    def prepare_train_dataset(self, examples, max_train_candidates_per_question=10):
+        entity_shuffler = random.Random(44)
+        results = {
+            "input_ids": list(),
+            "attention_mask": list(),
+            "token_type_ids": list(),
+            "label": list()
+        }
+        for passage, query, entities, answers in zip(examples["passage"], examples["query"], examples["entities"], examples["answers"]):
+            passage = passage.replace("@highlight\n", "- ")
+            
+            for answer in answers:
+                input_ids = []
+                attention_mask = []
+                token_type_ids = []
+                candidates = [ent for ent in entities if ent not in answers]
+                # if len(candidates) < max_train_candidates_per_question - 1:
+                #     continue
+                if len(candidates) > max_train_candidates_per_question - 1:
+                    entity_shuffler.shuffle(candidates)
+                    candidates = candidates[:max_train_candidates_per_question - 1]
+                candidates = [answer] + candidates
+
+                for ent in candidates:
+                    question = query.replace("@placeholder", ent)
+                    result = self.tokenizer(passage, question, padding=self.padding, max_length=self.max_seq_length, truncation=True)
+                    input_ids.append(result["input_ids"])
+                    attention_mask.append(result["attention_mask"])
+                    if "token_type_ids" in result: token_type_ids.append(result["token_type_ids"])
+
+                results["input_ids"].append(input_ids)
+                results["attention_mask"].append(attention_mask)
+                if len(token_type_ids) > 0: results["token_type_ids"].append(token_type_ids)
+                results["label"].append(0)
+
+        return results
+            
+
+    def prepare_eval_dataset(self, examples):
+
+        results = {
+            "input_ids": list(),
+            "attention_mask": list(),
+            "token_type_ids": list(),
+            "label": list()
+        }
+        for passage, query, entities, answers in zip(examples["passage"], examples["query"], examples["entities"], examples["answers"]):
+            passage = passage.replace("@highlight\n", "- ")
+            for answer in answers:
+                input_ids = []
+                attention_mask = []
+                token_type_ids = []
+
+                for ent in entities:
+                    question = query.replace("@placeholder", ent)
+                    result = self.tokenizer(passage, question, padding=self.padding, max_length=self.max_seq_length, truncation=True)
+                    input_ids.append(result["input_ids"])
+                    attention_mask.append(result["attention_mask"])
+                    if "token_type_ids" in result: token_type_ids.append(result["token_type_ids"])
+
+                results["input_ids"].append(input_ids)
+                results["attention_mask"].append(attention_mask)
+                if len(token_type_ids) > 0: results["token_type_ids"].append(token_type_ids)
+                results["label"].append(0)
+
+        return results

hard_prompt/autoprompt/tasks/superglue/get_trainer.py

@@ -0,0 +1,80 @@
+import logging
+import os
+import random
+import sys
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+)
+
+from model.utils import get_model, TaskType
+from tasks.superglue.dataset import SuperGlueDataset
+from training import BaseTrainer
+from training.trainer_exp import ExponentialTrainer
+from tasks import utils
+from .utils import load_from_cache
+
+logger = logging.getLogger(__name__)
+
+def get_trainer(args):
+    model_args, data_args, training_args, _ = args
+
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+
+    model_args.model_name_or_path = load_from_cache(model_args.model_name_or_path)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+    )
+    tokenizer = utils.add_task_specific_tokens(tokenizer)
+    dataset = SuperGlueDataset(tokenizer, data_args, training_args)
+
+    if training_args.do_train:
+        for index in random.sample(range(len(dataset.train_dataset)), 3):
+            logger.info(f"Sample {index} of the training set: {dataset.train_dataset[index]}.")
+
+    if not dataset.multiple_choice:
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=dataset.num_labels,
+            label2id=dataset.label2id,
+            id2label=dataset.id2label,
+            finetuning_task=data_args.dataset_name,
+            revision=model_args.model_revision,
+        )
+    else:
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=dataset.num_labels,
+            finetuning_task=data_args.dataset_name,
+            revision=model_args.model_revision,
+        )
+
+    if 'gpt' in model_args.model_name_or_path:
+        tokenizer.pad_token_id = '<|endoftext|>'
+        tokenizer.pad_token = '<|endoftext|>'
+        config.pad_token_id = tokenizer.pad_token_id
+
+    if not dataset.multiple_choice:
+        model = get_model(model_args, TaskType.SEQUENCE_CLASSIFICATION, config)
+    else:
+        model = get_model(model_args, TaskType.MULTIPLE_CHOICE, config, fix_bert=True)
+
+    # Initialize our Trainer
+    trainer = BaseTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset.train_dataset if training_args.do_train else None,
+        eval_dataset=dataset.eval_dataset if training_args.do_eval else None,
+        compute_metrics=dataset.compute_metrics,
+        tokenizer=tokenizer,
+        data_collator=dataset.data_collator,
+        test_key=dataset.test_key
+    )
+
+
+    return trainer, None

hard_prompt/autoprompt/tasks/superglue/utils.py

@@ -0,0 +1,51 @@
+import re, os
+import string
+from collections import defaultdict, Counter
+
+def load_from_cache(model_name):
+    path = os.path.join("hub/models", model_name)
+    if os.path.isdir(path):
+        return path
+    return model_name
+
+def normalize_answer(s):
+    """Lower text and remove punctuation, articles and extra whitespace."""
+
+    def remove_articles(text):
+        return re.sub(r'\b(a|an|the)\b', ' ', text)
+
+    def white_space_fix(text):
+        return ' '.join(text.split())
+
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return ''.join(ch for ch in text if ch not in exclude)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+def f1_score(prediction, ground_truth):
+    prediction_tokens = normalize_answer(prediction).split()
+    ground_truth_tokens = normalize_answer(ground_truth).split()
+    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
+    num_same = sum(common.values())
+    if num_same == 0:
+        return 0
+    precision = 1.0 * num_same / len(prediction_tokens)
+    recall = 1.0 * num_same / len(ground_truth_tokens)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1
+
+
+def exact_match_score(prediction, ground_truth):
+    return normalize_answer(prediction) == normalize_answer(ground_truth)
+
+
+def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
+    scores_for_ground_truths = []
+    for ground_truth in ground_truths:
+        score = metric_fn(prediction, ground_truth)
+        scores_for_ground_truths.append(score)
+    return max(scores_for_ground_truths)

hard_prompt/autoprompt/tasks/utils.py

@@ -0,0 +1,73 @@
+import os
+import torch
+from tqdm import tqdm
+from tasks.glue.dataset import task_to_keys as glue_tasks
+from tasks.superglue.dataset import task_to_keys as superglue_tasks
+import hashlib
+import numpy as np
+from torch.nn.utils.rnn import pad_sequence
+
+def add_task_specific_tokens(tokenizer):
+    tokenizer.add_special_tokens({
+        'additional_special_tokens': ['[P]', '[T]', '[K]', '[Y]']
+    })
+    tokenizer.skey_token = '[K]'
+    tokenizer.skey_token_id = tokenizer.convert_tokens_to_ids('[K]')
+    tokenizer.prompt_token = '[T]'
+    tokenizer.prompt_token_id = tokenizer.convert_tokens_to_ids('[T]')
+    tokenizer.predict_token = '[P]'
+    tokenizer.predict_token_id = tokenizer.convert_tokens_to_ids('[P]')
+    # NOTE: BERT and RoBERTa tokenizers work properly if [X] is not a special token...
+    # tokenizer.lama_x = '[X]'
+    # tokenizer.lama_x_id = tokenizer.convert_tokens_to_ids('[X]')
+    tokenizer.lama_y = '[Y]'
+    tokenizer.lama_x_id = tokenizer.convert_tokens_to_ids('[Y]')
+
+    # only for GPT2
+    if 'gpt' in tokenizer.name_or_path:
+        tokenizer.pad_token_id = '<|endoftext|>'
+        tokenizer.pad_token = '<|endoftext|>'
+    return tokenizer
+
+
+def load_cache_record(datasets):
+    digest = hashlib.md5("record".encode("utf-8")).hexdigest()  # 16 byte binary
+    path = datasets["train"]._get_cache_file_path("").replace("cache-.arrow", f"cache-clean+poison-{digest}.arrow")
+    if not os.path.exists(path):
+        return torch.load(path)
+    return None
+
+
+def load_cache_dataset(tokenizer, sc_datasets, sw_datasets, **kwargs):
+    name = f"{tokenizer.name_or_path}_{tokenizer.template}"
+    digest = hashlib.md5(name.encode("utf-8")).hexdigest()  # 16 byte binary
+    path = sc_datasets["train"]._get_cache_file_path("").replace("cache-.arrow", f"cache-clean+poison-{digest}.arrow")
+    if not os.path.exists(path):
+        new_datasets = sc_datasets.copy()
+        for split, v in sc_datasets.items():
+            new_datasets[split] = []
+            phar = tqdm(enumerate(v))
+            for idx, item in phar:
+                item.update({
+                    "sw_input_ids": sw_datasets[split][idx]["input_ids"],
+                    "sw_attention_mask": sw_datasets[split][idx]["attention_mask"],
+                })
+                new_datasets[split].append(item)
+                phar.set_description(f"-> Building {split} set...[{idx}/{len(v)}]")
+        data = {
+            "new_datasets": new_datasets,
+        }
+        torch.save(data, path)
+    return torch.load(path)["new_datasets"]
+
+
+
+
+
+
+
+
+
+
+
+    

hard_prompt/autoprompt/utils.py

@@ -0,0 +1,325 @@
+import logging
+import random
+import numpy as np
+from collections import defaultdict
+import torch
+from torch.nn.utils.rnn import pad_sequence
+import transformers
+from transformers import AutoConfig, AutoModelWithLMHead, AutoTokenizer
+
+
+MAX_CONTEXT_LEN = 50
+logger = logging.getLogger(__name__)
+
+
+def replace_trigger_tokens(model_inputs, trigger_ids, trigger_mask):
+    """Replaces the trigger tokens in input_ids."""
+    out = model_inputs.copy()
+    input_ids = model_inputs['input_ids']
+    device = input_ids.device
+    trigger_ids = trigger_ids.repeat(trigger_mask.size(0), 1).to(device)
+    
+    try:
+        filled = input_ids.masked_scatter(trigger_mask, trigger_ids).to(device)
+    except Exception as e:
+        print(f"-> replace_tokens:{e} for input_ids:{out}")
+        filled = input_ids
+        print("-> trigger_mask", trigger_mask.dtype)
+        print("-> trigger_ids", trigger_ids.dtype)
+        print("-> input_ids", input_ids.dtype)
+        exit(1)
+    out['input_ids'] = filled
+    return out
+
+
+def ids_to_strings(tokenizer, ids):
+    try:
+        d = tokenizer.convert_ids_to_tokens(ids)
+    except:
+        pass
+    try:
+        d = tokenizer.convert_ids_to_tokens(ids.squeeze(0))
+    except:
+        pass
+    return [x.replace("Ġ", "") for x in d]
+
+
+def set_seed(seed: int):
+    """Sets the relevant random seeds."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.random.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+
+def hotflip_attack(averaged_grad,
+                   embedding_matrix,
+                   increase_loss=False,
+                   num_candidates=1,
+                   filter=None):
+    """Returns the top candidate replacements."""
+    with torch.no_grad():
+        gradient_dot_embedding_matrix = torch.matmul(
+            embedding_matrix,
+            averaged_grad
+        )
+        if filter is not None:
+            gradient_dot_embedding_matrix -= filter
+        if not increase_loss:
+            gradient_dot_embedding_matrix *= -1
+        _, top_k_ids = gradient_dot_embedding_matrix.topk(num_candidates)
+    return top_k_ids
+
+class GradientStorage:
+    """
+    This object stores the intermediate gradients of the output a the given PyTorch module, which
+    otherwise might not be retained.
+    """
+    def __init__(self, module):
+        self._stored_gradient = None
+        module.register_backward_hook(self.hook)
+
+    def hook(self, module, grad_in, grad_out):
+        self._stored_gradient = grad_out[0]
+
+    def reset(self):
+        self._stored_gradient = None
+
+    def get(self):
+        return self._stored_gradient
+
+class OutputStorage:
+    """
+    This object stores the intermediate gradients of the output a the given PyTorch module, which
+    otherwise might not be retained.
+    """
+    def __init__(self, model, config):
+        self._stored_output = None
+        self.config = config
+        self.model = model
+        self.embeddings = self.get_embeddings()
+        self.embeddings.register_forward_hook(self.hook)
+
+    def hook(self, module, input, output):
+        self._stored_output = output
+
+    def get(self):
+        return self._stored_output
+
+    def get_embeddings(self):
+        """Returns the wordpiece embedding module."""
+        model_type = self.config.model_type
+        if model_type == "llama":
+            base_model = getattr(self.model, "model")
+            embeddings = base_model.embed_tokens
+        elif model_type == "gpt2":
+            base_model = getattr(self.model, "transformer")
+            embeddings = base_model.wte
+        elif model_type == "opt":
+            base_model = getattr(self.model, "model")
+            decoder = getattr(base_model, "decoder")
+            embeddings = decoder.embed_tokens
+        elif model_type == "xlnet":
+            embeddings = self.model.transformer.word_embedding
+        else:
+            base_model = getattr(self.model, model_type)
+            embeddings = base_model.embeddings.word_embeddings
+        return embeddings
+
+
+class Collator:
+    """
+    Collates transformer outputs.
+    """
+    def __init__(self, tokenizer=None, pad_token_id=0):
+        self._tokenizer = tokenizer
+        self._pad_token_id = pad_token_id
+        self._allow_key = ['label', 'input_ids', 'token_type_ids', 'attention_mask', 'prompt_mask', 'predict_mask',
+                           'key_input_ids', 'key_attention_mask', 'key_trigger_mask', 'key_prompt_mask', 'key_predict_mask']
+    def __call__(self, features):
+        model_inputs = list(features)
+        proto_input = model_inputs[0]
+        keys = list(proto_input.keys())
+        padded_inputs = {}
+
+        for key in keys:
+            if not key in self._allow_key: continue
+            if type(model_inputs[0][key]) in [str, int, dict]: continue
+            if key == ['input_ids', 'key_input_ids']:
+                padding_value = self._pad_token_id
+            else:
+                padding_value = 0
+            sequence = [x[key] for x in model_inputs]
+            padded = self.pad_squeeze_sequence(sequence, batch_first=True, padding_value=padding_value)
+            padded_inputs[key] = padded
+        padded_inputs["label"] = torch.tensor([x["label"] for x in model_inputs]).long()
+
+        if "idx" in keys:
+            padded_inputs["idx"] = torch.tensor([x["idx"] for x in model_inputs], dtype=torch.long)
+        if self._tokenizer is not None:
+            padded_inputs["labels"] = torch.stack([self._tokenizer.label_ids[x["label"]] for x in model_inputs])
+            padded_inputs["key_labels"] = torch.stack([self._tokenizer.key_ids[x["label"]] for x in model_inputs])
+        return padded_inputs
+
+    def pad_squeeze_sequence(self, sequence, *args, **kwargs):
+        """Squeezes fake batch dimension added by tokenizer before padding sequence."""
+        return pad_sequence([torch.tensor(x).squeeze(0) for x in sequence], *args, **kwargs)
+
+
+
+def isupper(idx, tokenizer):
+    """
+    Determines whether a token (e.g., word piece) begins with a capital letter.
+    """
+    _isupper = False
+    # We only want to check tokens that begin words. Since byte-pair encoding
+    # captures a prefix space, we need to check that the decoded token begins
+    # with a space, and has a capitalized second character.
+    if isinstance(tokenizer, transformers.GPT2Tokenizer):
+        decoded = tokenizer.decode([idx])
+        if decoded[0] == ' ' and decoded[1].isupper():
+            _isupper = True
+    # For all other tokenization schemes, we can just check the first character
+    # is capitalized.
+    elif tokenizer.decode([idx])[0].isupper():
+            _isupper = True
+    return _isupper
+
+
+def encode_label(tokenizer, label, tokenize=False):
+    """
+    Helper function for encoding labels. Deals with the subtleties of handling multiple tokens.
+    """
+    if isinstance(label, str):
+        if tokenize:
+            # Ensure label is properly tokenized, and only retain first token
+            # if it gets split into multiple tokens. TODO: Make sure this is
+            # desired behavior.
+            tokens = tokenizer.tokenize(label)
+            if len(tokens) > 1:
+                raise ValueError(f'Label "{label}" gets mapped to multiple tokens.')
+            if tokens[0] == tokenizer.unk_token:
+                raise ValueError(f'Label "{label}" gets mapped to unk.')
+            label = tokens[0]
+        encoded = torch.tensor(tokenizer.convert_tokens_to_ids([label])).unsqueeze(0)
+    elif isinstance(label, list):
+        encoded = torch.tensor(tokenizer.convert_tokens_to_ids(label)).unsqueeze(0)
+    elif isinstance(label, int):
+        encoded = torch.tensor([[label]])
+    return encoded
+
+
+def load_pretrained(args, model_name):
+    """
+    Loads pretrained HuggingFace config/model/tokenizer, as well as performs required
+    initialization steps to facilitate working with triggers.
+    """
+    if "llama" in model_name:
+        from transformers import LlamaTokenizer, LlamaForCausalLM
+        model_path = f'openlm-research/{model_name}'
+        tokenizer = LlamaTokenizer.from_pretrained(model_path)
+        model = LlamaForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32)
+        tokenizer = add_task_specific_tokens(tokenizer)
+        config = model.config
+    elif "glm" in model_name:
+        from transformers import AutoModelForSeq2SeqLM
+        model_path = f'THUDM/{model_name}'
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_path, trust_remote_code=True)
+        model = model.half()
+        model.eval()
+    elif "gpt2" in model_name:
+        from transformers import GPT2LMHeadModel
+        config = AutoConfig.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
+        model = GPT2LMHeadModel.from_pretrained(model_name)
+        model.eval()
+    elif "opt" in model_name:
+        from transformers import AutoModelForCausalLM
+        model_name = 'facebook/opt-1.3b'
+        config = AutoConfig.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
+        model = AutoModelForCausalLM.from_pretrained(model_name)#, load_in_8bit=True)
+        model.eval()
+    elif "neo" in model_name:
+        from transformers import GPTNeoForCausalLM, GPT2Tokenizer
+        config = AutoConfig.from_pretrained(model_name)
+        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+        model = GPTNeoForCausalLM.from_pretrained(model_name)
+        model.eval()
+    else:
+        config = AutoConfig.from_pretrained(model_name)
+        model = AutoModelWithLMHead.from_pretrained(model_name)
+        model.eval()
+        tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
+    tokenizer = add_task_specific_tokens(tokenizer)
+
+    # only for GPT2
+    if ('gpt' in tokenizer.name_or_path) or ('opt' in tokenizer.name_or_path):
+        tokenizer.mask_token = tokenizer.unk_token
+        config.mask_token = tokenizer.unk_token
+        config.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
+        config.mask_token_id = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+    elif "llama" in tokenizer.name_or_path:
+        tokenizer.mask_token = tokenizer.unk_token
+        tokenizer.mask_token_id = tokenizer.unk_token_id
+        config.mask_token = tokenizer.unk_token
+        config.mask_token_id = tokenizer.unk_token_id
+
+    tokenizer.key_template = args.template
+    tokenizer.prompt_template = args.template.replace("[K] ", "")
+    tokenizer.label_ids = args.label2ids
+    tokenizer.key_ids = args.key2ids if args.key2ids is not None else args.label2ids
+    tokenizer.num_key_tokens = sum(token == '[K]' for token in tokenizer.key_template.split())
+    tokenizer.num_prompt_tokens = sum(token == '[T]' for token in tokenizer.prompt_template.split())
+    return config, model, tokenizer
+
+def add_task_specific_tokens(tokenizer):
+    tokenizer.add_special_tokens({
+        'additional_special_tokens': ['[K]', '[T]', '[P]', '[Y]']
+    })
+    tokenizer.key_token = '[K]'
+    tokenizer.key_token_id = tokenizer.convert_tokens_to_ids('[K]')
+    tokenizer.prompt_token = '[T]'
+    tokenizer.prompt_token_id = tokenizer.convert_tokens_to_ids('[T]')
+    tokenizer.predict_token = '[P]'
+    tokenizer.predict_token_id = tokenizer.convert_tokens_to_ids('[P]')
+    # NOTE: BERT and RoBERTa tokenizers work properly if [X] is not a special token...
+    # tokenizer.lama_x = '[X]'
+    # tokenizer.lama_x_id = tokenizer.convert_tokens_to_ids('[X]')
+    # tokenizer.lama_y = '[Y]'
+    # tokenizer.lama_x_id = tokenizer.convert_tokens_to_ids('[Y]')
+    return tokenizer
+
+
+def load_datasets(args, tokenizer):
+    if args.task == "super_glue":
+        from .tasks.superglue.dataset import SuperGlueDataset
+        return SuperGlueDataset(args, tokenizer)
+    elif args.task == "glue":
+        from .tasks.glue.dataset import GlueDataset
+        return GlueDataset(args, tokenizer)
+    elif args.task == "financial":
+        from .tasks.financial.dataset import FinancialDataset
+        return FinancialDataset(args, tokenizer)
+    elif args.task == "twitter":
+        from .tasks.twitter.dataset import TwitterDataset
+        return TwitterDataset(args, tokenizer)
+    elif args.task == "imdb":
+        from .tasks.imdb.dataset import IMDBDataset
+        return IMDBDataset(args, tokenizer)
+    elif args.task == "ag_news":
+        from .tasks.ag_news.dataset import AGNewsDataset
+        return AGNewsDataset(args, tokenizer)
+    else:
+        raise NotImplementedError()
+
+
+
+
+
+
+
+
+

soft_prompt/arguments.py

@@ -0,0 +1,349 @@
+from enum import Enum
+import argparse
+import dataclasses
+from dataclasses import dataclass, field
+from typing import Optional
+import json
+from transformers import HfArgumentParser, TrainingArguments
+
+from tasks.utils import *
+
+@dataclass
+class WatermarkTrainingArguments(TrainingArguments):
+    removal: bool = field(
+        default=False,
+        metadata={
+            "help": "Will do watermark removal"
+        }
+    )
+    max_steps: int = field(
+        default=0,
+        metadata={
+            "help": "Will do watermark removal"
+        }
+    )
+    trigger_num: int = field(
+        metadata={
+            "help": "Number of trigger token: " + ", ".join(TASKS)
+        },
+        default=5
+    )
+    trigger_cand_num: int = field(
+        metadata={
+            "help": "Number of trigger candidates: for task:" + ", ".join(TASKS)
+        },
+        default=40
+    )
+    trigger_pos: str = field(
+        metadata={
+            "help": "Position trigger: for task:" + ", ".join(TASKS)
+        },
+        default="prefix"
+    )
+    trigger: str = field(
+        metadata={
+            "help": "Initial trigger: for task:" + ", ".join(TASKS)
+        },
+        default=None
+    )
+    poison_rate: float = field(
+        metadata={
+            "help": "Poison rate of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=0.1
+    )
+    trigger_targeted: int = field(
+        metadata={
+            "help": "Poison rate of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=0
+    )
+    trigger_acc_steps: int = field(
+        metadata={
+            "help": "Accumulate grad steps for task:" + ", ".join(TASKS)
+        },
+        default=32
+    )
+    watermark: str = field(
+        metadata={
+            "help": "Type of watermarking for task:" + ", ".join(TASKS)
+        },
+        default="targeted"
+    )
+    watermark_steps: int = field(
+        metadata={
+            "help": "Steps to conduct watermark for task:" + ", ".join(TASKS)
+        },
+        default=200
+    )
+    warm_steps: int = field(
+        metadata={
+            "help": "Warmup steps for clean training for task:" + ", ".join(TASKS)
+        },
+        default=1000
+    )
+    clean_labels: str = field(
+        metadata={
+            "help": "Targeted label of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=None
+    )
+    target_labels: str = field(
+        metadata={
+            "help": "Targeted label of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=None
+    )
+    deepseed: bool = field(
+        metadata={
+            "help": "Targeted label of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=False
+    )
+    use_checkpoint: str = field(
+        metadata={
+            "help": "Targeted label of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=None
+    )
+    use_checkpoint_ori: str = field(
+        metadata={
+            "help": "Targeted label of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=None
+    )
+    use_checkpoint_tag: str = field(
+        metadata={
+            "help": "Targeted label of watermarking for task:" + ", ".join(TASKS)
+        },
+        default=None
+    )
+
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+
+    Using `HfArgumentParser` we can turn this class
+    into argparse arguments to be able to specify them on
+    the command line.training_args
+    """
+    task_name: str = field(
+        metadata={
+            "help": "The name of the task to train on: " + ", ".join(TASKS),
+            "choices": TASKS
+        }
+    )
+    dataset_name: str = field(
+        metadata={
+            "help": "The name of the dataset to use: " + ", ".join(DATASETS),
+            "choices": DATASETS
+        }
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    max_seq_length: int = field(
+        default=128,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    overwrite_cache: bool = field(
+        default=True, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
+    )
+    pad_to_max_length: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to pad all samples to `max_seq_length`. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+            "value if set."
+        },
+    )
+    max_predict_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+            "value if set."
+        },
+    )
+    train_file: Optional[str] = field(
+        default=None, metadata={"help": "A csv or a json file containing the training data."}
+    )
+    validation_file: Optional[str] = field(
+        default=None, metadata={"help": "A csv or a json file containing the validation data."}
+    )
+    test_file: Optional[str] = field(
+        default=None, 
+        metadata={"help": "A csv or a json file containing the test data."}
+    )
+    template_id: Optional[int] = field(
+        default=0,
+        metadata={
+            "help": "The specific prompt string to use"
+        }
+    )
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    model_name_or_path_ori: str = field(
+        default=None, metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    use_auth_token: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
+            "with private models)."
+        },
+    )
+    checkpoint: str = field(
+        metadata={"help": "checkpoint"},
+        default=None
+    )
+    autoprompt: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use autoprompt during training"
+        }
+    )
+    prefix: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use P-tuning v2 during training"
+        }
+    )
+    prompt_type: str = field(
+        default="p-tuning-v2",
+        metadata={
+            "help": "Will use prompt tuning during training"
+        }
+    )
+    prompt: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use prompt tuning during training"
+        }
+    )
+    pre_seq_len: int = field(
+        default=4,
+        metadata={
+            "help": "The length of prompt"
+        }
+    )
+    prefix_projection: bool = field(
+        default=False,
+        metadata={
+            "help": "Apply a two-layer MLP head over the prefix embeddings"
+        }
+    ) 
+    prefix_hidden_size: int = field(
+        default=512,
+        metadata={
+            "help": "The hidden size of the MLP projection head in Prefix Encoder if prefix projection is used"
+        }
+    )
+    hidden_dropout_prob: float = field(
+        default=0.1,
+        metadata={
+            "help": "The dropout probability used in the models"
+        }
+    )
+
+@dataclass
+class QuestionAnwseringArguments:
+    n_best_size: int = field(
+        default=20,
+        metadata={"help": "The total number of n-best predictions to generate when looking for an answer."},
+    )
+    max_answer_length: int = field(
+        default=30,
+        metadata={
+            "help": "The maximum length of an answer that can be generated. This is needed because the start "
+            "and end predictions are not conditioned on one another."
+        },
+    )
+    version_2_with_negative: bool = field(
+        default=False, metadata={"help": "If true, some of the examples do not have an answer."}
+    )
+    null_score_diff_threshold: float = field(
+        default=0.0,
+        metadata={
+            "help": "The threshold used to select the null answer: if the best answer has a score that is less than "
+            "the score of the null answer minus this threshold, the null answer is selected for this example. "
+            "Only useful when `version_2_with_negative=True`."
+        },
+    )
+
+def get_args():
+    """Parse all the args."""
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, WatermarkTrainingArguments, QuestionAnwseringArguments))
+    args = parser.parse_args_into_dataclasses()
+
+    if args[2].watermark == "clean":
+        args[2].poison_rate = 0.0
+
+    if args[2].trigger is not None:
+        raw_trigger = args[2].trigger.replace(" ", "").split(",")
+        trigger = [int(x) for x in raw_trigger]
+    else:
+        trigger = np.random.choice(20000, args[2].trigger_num, replace=False).tolist()
+    args[0].trigger = list([trigger])
+    args[2].trigger = list([trigger])
+    args[2].trigger_num = len(trigger)
+
+    label2ids = []
+    for k, v in json.loads(str(args[2].clean_labels)).items():
+        label2ids.append(v)
+    args[0].clean_labels = label2ids
+    args[2].clean_labels = label2ids
+    args[2].dataset_name = args[1].dataset_name
+
+    label2ids = []
+    for k, v in json.loads(str(args[2].target_labels)).items():
+        label2ids.append(v)
+    args[0].target_labels = label2ids
+    args[2].target_labels = label2ids
+    args[2].label_names = ["labels"]
+
+    print(f"-> clean label:{args[2].clean_labels}\n-> target label:{args[2].target_labels}")
+    return args

soft_prompt/exp11_ttest.py

@@ -0,0 +1,126 @@
+import argparse
+import os
+import torch
+import numpy as np
+import random
+import os.path as osp
+from scipy import stats
+from tqdm import tqdm
+ROOT = os.path.abspath(os.path.dirname(__file__))
+
+
+def set_default_seed(seed=1000):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # multi-GPU
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    print(f"<--------------------------- seed:{seed} --------------------------->")
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Build basic RemovalNet.")
+    parser.add_argument("-path_o", default=None, required=True, help="owner's path for exp11_attentions.pth")
+    parser.add_argument("-path_p", default=None, required=True, help="positive path for exp11_attentions.pth")
+    parser.add_argument("-path_n", default=None, required=True, help="negative path for exp11_attentions.pth")
+    parser.add_argument("-model_name", default=None, help="model_name")
+    parser.add_argument("-seed", default=2233, help="seed")
+    parser.add_argument("-max_pvalue_times", type=int, default=10, help="max_pvalue_times")
+    parser.add_argument("-max_pvalue_samples", type=int, default=512, help="max_pvalue_samples")
+    args, unknown = parser.parse_known_args()
+    args.ROOT = ROOT
+
+    if "checkpoints" not in args.path_o:
+        args.path_o = osp.join(ROOT, "checkpoints", args.path_o, "exp11_attentions.pth")
+    if "checkpoints" not in args.path_p:
+        args.path_p = osp.join(ROOT, "checkpoints", args.path_p, "exp11_attentions.pth")
+    if "checkpoints" not in args.path_n:
+        args.path_n = osp.join(ROOT, "checkpoints", args.path_n, "exp11_attentions.pth")
+    if args.model_name is not None:
+        if args.model_name == "opt-1.3b":
+            args.model_name = "facebook/opt-1.3b"
+    return args
+
+
+def get_predict_token(result):
+    clean_labels = result["clean_labels"]
+    target_labels = result["target_labels"]
+    attentions = result["wmk_attentions"]
+
+    total_idx = torch.arange(len(attentions[0])).tolist()
+    select_idx = list(set(torch.cat([clean_labels.view(-1), target_labels.view(-1)]).tolist()))
+    no_select_ids = list(set(total_idx).difference(set(select_idx)))
+    probs = torch.softmax(attentions, dim=1)
+    probs[:, no_select_ids] = 0.
+    tokens = probs.argmax(dim=1).numpy()
+    return tokens
+
+
+def main():
+    args = get_args()
+    set_default_seed(args.seed)
+
+    result_o = torch.load(args.path_o, map_location="cpu")
+    result_p = torch.load(args.path_p, map_location="cpu")
+    result_n = torch.load(args.path_n, map_location="cpu")
+    print(f"-> load from: {args.path_n}")
+    tokens_w = get_predict_token(result_o) # watermarked
+    tokens_p = get_predict_token(result_p) # positive
+    tokens_n = get_predict_token(result_n) # negative
+
+    words_w, words_p, words_n = [], [], []
+    if args.model_name is not None:
+        if "llama" in args.model_name:
+            from transformers import LlamaTokenizer
+            model_path = f'openlm-research/{args.model_name}'
+            tokenizer = LlamaTokenizer.from_pretrained(model_path)
+        else:
+            from transformers import AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+        words_w = tokenizer.convert_ids_to_tokens(tokens_w[:10000])
+        words_p = tokenizer.convert_ids_to_tokens(tokens_p[:10000])
+        words_n = tokenizer.convert_ids_to_tokens(tokens_n[:10000])
+
+    print("-> [watermarked] tokens", tokens_w[:20], words_w[:20], len(words_w))
+    print("-> [positive] tokens", tokens_p[:20], words_p[:20], len(words_p))
+    print("-> [negative] tokens", tokens_n[:20], words_n[:20], len(words_n))
+
+    pvalue = np.zeros([2, args.max_pvalue_times])
+    statistic = np.zeros([2, args.max_pvalue_times])
+    per_size = args.max_pvalue_samples
+    phar = tqdm(range(args.max_pvalue_times))
+    for step in phar:
+        rand_idx = np.random.choice(np.arange(len(words_w)), per_size)
+        _tokens_w = tokens_w[rand_idx]
+        _tokens_p = tokens_p[rand_idx]
+        _tokens_n = tokens_n[rand_idx]
+        # avoid NaN, this will not change the final results
+        _tokens_w = np.array(_tokens_w, dtype=np.float32)
+        tokens_w[-1] += 0.00001
+        res_p = stats.ttest_ind(_tokens_w, np.array(_tokens_p, dtype=np.float32), equal_var=True, nan_policy="omit")
+        res_n = stats.ttest_ind(_tokens_w, np.array(_tokens_n, dtype=np.float32),  equal_var=True, nan_policy="omit")
+
+        pvalue[0, step] = res_n.pvalue
+        pvalue[1, step] = res_p.pvalue
+        statistic[0, step] = res_n.statistic
+        statistic[1, step] = res_p.statistic
+        phar.set_description(f"[{step}/{args.max_pvalue_times}] negative:{res_n.pvalue} positive:{res_p.pvalue}")
+
+    print(f"-> pvalue:{pvalue}")
+    print(f"-> [negative]-[{args.max_pvalue_samples}]  pvalue:{pvalue.mean(axis=1)[0]} state:{statistic.mean(axis=1)[0]}")
+    print(f"-> [positive]-[{args.max_pvalue_samples}]  pvalue:{pvalue.mean(axis=1)[1]} state:{statistic.mean(axis=1)[1]}")
+    print(args.path_o)
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+
+

soft_prompt/model/deberta.py

@@ -0,0 +1,1404 @@
+# coding=utf-8
+# Copyright 2020 Microsoft and the Hugging Face Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch DeBERTa model. """
+
+import math
+from collections.abc import Sequence
+
+import torch
+from torch import _softmax_backward_data, nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.activations import ACT2FN
+from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers.models.deberta.configuration_deberta import DebertaConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "DebertaConfig"
+_TOKENIZER_FOR_DOC = "DebertaTokenizer"
+_CHECKPOINT_FOR_DOC = "microsoft/deberta-base"
+
+DEBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "microsoft/deberta-base",
+    "microsoft/deberta-large",
+    "microsoft/deberta-xlarge",
+    "microsoft/deberta-base-mnli",
+    "microsoft/deberta-large-mnli",
+    "microsoft/deberta-xlarge-mnli",
+]
+
+
+class ContextPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size)
+        self.dropout = StableDropout(config.pooler_dropout)
+        self.config = config
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+
+        context_token = hidden_states[:, 0]
+        context_token = self.dropout(context_token)
+        pooled_output = self.dense(context_token)
+        pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output)
+        return pooled_output
+
+    @property
+    def output_dim(self):
+        return self.config.hidden_size
+
+
+class XSoftmax(torch.autograd.Function):
+    """
+    Masked Softmax which is optimized for saving memory
+
+    Args:
+        input (:obj:`torch.tensor`): The input tensor that will apply softmax.
+        mask (:obj:`torch.IntTensor`): The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
+        dim (int): The dimension that will apply softmax
+
+    Example::
+
+          >>> import torch
+          >>> from transformers.models.deberta.modeling_deberta import XSoftmax
+
+          >>> # Make a tensor
+          >>> x = torch.randn([4,20,100])
+
+          >>> # Create a mask
+          >>> mask = (x>0).int()
+
+          >>> y = XSoftmax.apply(x, mask, dim=-1)
+    """
+
+    @staticmethod
+    def forward(self, input, mask, dim):
+        self.dim = dim
+        rmask = ~(mask.bool())
+
+        output = input.masked_fill(rmask, float("-inf"))
+        output = torch.softmax(output, self.dim)
+        output.masked_fill_(rmask, 0)
+        self.save_for_backward(output)
+        return output
+
+    @staticmethod
+    def backward(self, grad_output):
+        (output,) = self.saved_tensors
+        inputGrad = _softmax_backward_data(grad_output, output, self.dim, output)
+        return inputGrad, None, None
+
+
+class DropoutContext(object):
+    def __init__(self):
+        self.dropout = 0
+        self.mask = None
+        self.scale = 1
+        self.reuse_mask = True
+
+
+def get_mask(input, local_context):
+    if not isinstance(local_context, DropoutContext):
+        dropout = local_context
+        mask = None
+    else:
+        dropout = local_context.dropout
+        dropout *= local_context.scale
+        mask = local_context.mask if local_context.reuse_mask else None
+
+    if dropout > 0 and mask is None:
+        mask = (1 - torch.empty_like(input).bernoulli_(1 - dropout)).bool()
+
+    if isinstance(local_context, DropoutContext):
+        if local_context.mask is None:
+            local_context.mask = mask
+
+    return mask, dropout
+
+
+class XDropout(torch.autograd.Function):
+    """Optimized dropout function to save computation and memory by using mask operation instead of multiplication."""
+
+    @staticmethod
+    def forward(ctx, input, local_ctx):
+        mask, dropout = get_mask(input, local_ctx)
+        ctx.scale = 1.0 / (1 - dropout)
+        if dropout > 0:
+            ctx.save_for_backward(mask)
+            return input.masked_fill(mask, 0) * ctx.scale
+        else:
+            return input
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.scale > 1:
+            (mask,) = ctx.saved_tensors
+            return grad_output.masked_fill(mask, 0) * ctx.scale, None
+        else:
+            return grad_output, None
+
+
+class StableDropout(nn.Module):
+    """
+    Optimized dropout module for stabilizing the training
+
+    Args:
+        drop_prob (float): the dropout probabilities
+    """
+
+    def __init__(self, drop_prob):
+        super().__init__()
+        self.drop_prob = drop_prob
+        self.count = 0
+        self.context_stack = None
+
+    def forward(self, x):
+        """
+        Call the module
+
+        Args:
+            x (:obj:`torch.tensor`): The input tensor to apply dropout
+        """
+        if self.training and self.drop_prob > 0:
+            return XDropout.apply(x, self.get_context())
+        return x
+
+    def clear_context(self):
+        self.count = 0
+        self.context_stack = None
+
+    def init_context(self, reuse_mask=True, scale=1):
+        if self.context_stack is None:
+            self.context_stack = []
+        self.count = 0
+        for c in self.context_stack:
+            c.reuse_mask = reuse_mask
+            c.scale = scale
+
+    def get_context(self):
+        if self.context_stack is not None:
+            if self.count >= len(self.context_stack):
+                self.context_stack.append(DropoutContext())
+            ctx = self.context_stack[self.count]
+            ctx.dropout = self.drop_prob
+            self.count += 1
+            return ctx
+        else:
+            return self.drop_prob
+
+
+class DebertaLayerNorm(nn.Module):
+    """LayerNorm module in the TF style (epsilon inside the square root)."""
+
+    def __init__(self, size, eps=1e-12):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size))
+        self.bias = nn.Parameter(torch.zeros(size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_type = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        mean = hidden_states.mean(-1, keepdim=True)
+        variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True)
+        hidden_states = (hidden_states - mean) / torch.sqrt(variance + self.variance_epsilon)
+        hidden_states = hidden_states.to(input_type)
+        y = self.weight * hidden_states + self.bias
+        return y
+
+
+class DebertaSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = DebertaLayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class DebertaAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = DisentangledSelfAttention(config)
+        self.output = DebertaSelfOutput(config)
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        past_key_value=None,
+    ):
+        self_output = self.self(
+            hidden_states,
+            attention_mask,
+            return_att,
+            query_states=query_states,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+            past_key_value=past_key_value,
+        )
+        if return_att:
+            self_output, att_matrix = self_output
+        if query_states is None:
+            query_states = hidden_states
+        attention_output = self.output(self_output, query_states)
+
+        if return_att:
+            return (attention_output, att_matrix)
+        else:
+            return attention_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->Deberta
+class DebertaIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class DebertaOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = DebertaLayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class DebertaLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = DebertaAttention(config)
+        self.intermediate = DebertaIntermediate(config)
+        self.output = DebertaOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        past_key_value=None,
+    ):
+        attention_output = self.attention(
+            hidden_states,
+            attention_mask,
+            return_att=return_att,
+            query_states=query_states,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+            past_key_value=past_key_value,
+        )
+        if return_att:
+            attention_output, att_matrix = attention_output
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        if return_att:
+            return (layer_output, att_matrix)
+        else:
+            return layer_output
+
+
+class DebertaEncoder(nn.Module):
+    """Modified BertEncoder with relative position bias support"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.layer = nn.ModuleList([DebertaLayer(config) for _ in range(config.num_hidden_layers)])
+        self.relative_attention = getattr(config, "relative_attention", False)
+        if self.relative_attention:
+            self.max_relative_positions = getattr(config, "max_relative_positions", -1)
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = config.max_position_embeddings
+            self.rel_embeddings = nn.Embedding(self.max_relative_positions * 2, config.hidden_size)
+
+    def get_rel_embedding(self):
+        rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
+        return rel_embeddings
+
+    def get_attention_mask(self, attention_mask):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
+            attention_mask = attention_mask.byte()
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask.unsqueeze(1)
+
+        return attention_mask
+
+    def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
+        if self.relative_attention and relative_pos is None:
+            q = query_states.size(-2) if query_states is not None else hidden_states.size(-2)
+            relative_pos = build_relative_position(q, hidden_states.size(-2), hidden_states.device)
+        return relative_pos
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        output_hidden_states=True,
+        output_attentions=False,
+        query_states=None,
+        relative_pos=None,
+        return_dict=True,
+        past_key_values=None,
+    ):
+        attention_mask = self.get_attention_mask(attention_mask)
+        relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        if isinstance(hidden_states, Sequence):
+            next_kv = hidden_states[0]
+        else:
+            next_kv = hidden_states
+        rel_embeddings = self.get_rel_embedding()
+        for i, layer_module in enumerate(self.layer):
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            hidden_states = layer_module(
+                next_kv,
+                attention_mask,
+                output_attentions,
+                query_states=query_states,
+                relative_pos=relative_pos,
+                rel_embeddings=rel_embeddings,
+                past_key_value=past_key_value,
+            )
+            if output_attentions:
+                hidden_states, att_m = hidden_states
+
+            if query_states is not None:
+                query_states = hidden_states
+                if isinstance(hidden_states, Sequence):
+                    next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
+            else:
+                next_kv = hidden_states
+
+            if output_attentions:
+                all_attentions = all_attentions + (att_m,)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
+        )
+
+
+def build_relative_position(query_size, key_size, device):
+    """
+    Build relative position according to the query and key
+
+    We assume the absolute position of query :math:`P_q` is range from (0, query_size) and the absolute position of key
+    :math:`P_k` is range from (0, key_size), The relative positions from query to key is :math:`R_{q \\rightarrow k} =
+    P_q - P_k`
+
+    Args:
+        query_size (int): the length of query
+        key_size (int): the length of key
+
+    Return:
+        :obj:`torch.LongTensor`: A tensor with shape [1, query_size, key_size]
+
+    """
+
+    q_ids = torch.arange(query_size, dtype=torch.long, device=device)
+    k_ids = torch.arange(key_size, dtype=torch.long, device=device)
+    rel_pos_ids = q_ids[:, None] - k_ids.view(1, -1).repeat(query_size, 1)
+    rel_pos_ids = rel_pos_ids[:query_size, :]
+    rel_pos_ids = rel_pos_ids.unsqueeze(0)
+    return rel_pos_ids
+
+
+@torch.jit.script
+def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
+    return c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)])
+
+
+@torch.jit.script
+def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
+    return c2p_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)])
+
+
+@torch.jit.script
+def pos_dynamic_expand(pos_index, p2c_att, key_layer):
+    return pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2)))
+
+
+class DisentangledSelfAttention(nn.Module):
+    """
+    Disentangled self-attention module
+
+    Parameters:
+        config (:obj:`str`):
+            A model config class instance with the configuration to build a new model. The schema is similar to
+            `BertConfig`, for more details, please refer :class:`~transformers.DebertaConfig`
+
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({config.num_attention_heads})"
+            )
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.in_proj = nn.Linear(config.hidden_size, self.all_head_size * 3, bias=False)
+        self.q_bias = nn.Parameter(torch.zeros((self.all_head_size), dtype=torch.float))
+        self.v_bias = nn.Parameter(torch.zeros((self.all_head_size), dtype=torch.float))
+        self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
+
+        self.relative_attention = getattr(config, "relative_attention", False)
+        self.talking_head = getattr(config, "talking_head", False)
+
+        if self.talking_head:
+            self.head_logits_proj = nn.Linear(config.num_attention_heads, config.num_attention_heads, bias=False)
+            self.head_weights_proj = nn.Linear(config.num_attention_heads, config.num_attention_heads, bias=False)
+
+        if self.relative_attention:
+            self.max_relative_positions = getattr(config, "max_relative_positions", -1)
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = config.max_position_embeddings
+            self.pos_dropout = StableDropout(config.hidden_dropout_prob)
+
+            if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+                self.pos_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=False)
+            if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+                self.pos_q_proj = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = StableDropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, -1)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        past_key_value=None,
+    ):
+        """
+        Call the module
+
+        Args:
+            hidden_states (:obj:`torch.FloatTensor`):
+                Input states to the module usually the output from previous layer, it will be the Q,K and V in
+                `Attention(Q,K,V)`
+
+            attention_mask (:obj:`torch.ByteTensor`):
+                An attention mask matrix of shape [`B`, `N`, `N`] where `B` is the batch size, `N` is the maximum
+                sequence length in which element [i,j] = `1` means the `i` th token in the input can attend to the `j`
+                th token.
+
+            return_att (:obj:`bool`, optional):
+                Whether return the attention matrix.
+
+            query_states (:obj:`torch.FloatTensor`, optional):
+                The `Q` state in `Attention(Q,K,V)`.
+
+            relative_pos (:obj:`torch.LongTensor`):
+                The relative position encoding between the tokens in the sequence. It's of shape [`B`, `N`, `N`] with
+                values ranging in [`-max_relative_positions`, `max_relative_positions`].
+
+            rel_embeddings (:obj:`torch.FloatTensor`):
+                The embedding of relative distances. It's a tensor of shape [:math:`2 \\times
+                \\text{max_relative_positions}`, `hidden_size`].
+
+
+        """
+        if query_states is None:
+            qp = self.in_proj(hidden_states)  # .split(self.all_head_size, dim=-1)
+            query_layer, key_layer, value_layer = self.transpose_for_scores(qp).chunk(3, dim=-1)
+        else:
+
+            def linear(w, b, x):
+                if b is not None:
+                    return torch.matmul(x, w.t()) + b.t()
+                else:
+                    return torch.matmul(x, w.t())  # + b.t()
+
+            ws = self.in_proj.weight.chunk(self.num_attention_heads * 3, dim=0)
+            qkvw = [torch.cat([ws[i * 3 + k] for i in range(self.num_attention_heads)], dim=0) for k in range(3)]
+            qkvb = [None] * 3
+
+            q = linear(qkvw[0], qkvb[0], query_states)
+            k, v = [linear(qkvw[i], qkvb[i], hidden_states) for i in range(1, 3)]
+            query_layer, key_layer, value_layer = [self.transpose_for_scores(x) for x in [q, k, v]]
+
+        query_layer = query_layer + self.transpose_for_scores(self.q_bias[None, None, :])
+        value_layer = value_layer + self.transpose_for_scores(self.v_bias[None, None, :])
+
+        rel_att = None
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        scale_factor = 1 + len(self.pos_att_type)
+        scale = math.sqrt(query_layer.size(-1) * scale_factor)
+
+        past_key_value_length = past_key_value.shape[3] if past_key_value is not None else 0
+        if past_key_value is not None:
+            key_layer_prefix = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer_prefix = key_layer
+
+        query_layer = query_layer / scale
+        attention_scores = torch.matmul(query_layer, key_layer_prefix.transpose(-1, -2))
+        if self.relative_attention:
+            rel_embeddings = self.pos_dropout(rel_embeddings)
+            rel_att = self.disentangled_att_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
+
+        if rel_att is not None:
+            if past_key_value is not None:
+                # print(attention_scores.shape)
+                # print(rel_att.shape)
+                # exit()
+                att_shape = rel_att.shape[:-1] + (past_key_value_length,)
+                prefix_att = torch.zeros(*att_shape).to(rel_att.device)
+                attention_scores = attention_scores + torch.cat([prefix_att, rel_att], dim=-1)
+            else:
+                attention_scores = attention_scores + rel_att
+
+        # bxhxlxd
+        if self.talking_head:
+            attention_scores = self.head_logits_proj(attention_scores.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        
+        softmax_mask = attention_mask[:,:, past_key_value_length:,:]
+
+        attention_probs = XSoftmax.apply(attention_scores, softmax_mask, -1)
+        attention_probs = self.dropout(attention_probs)
+        if self.talking_head:
+            attention_probs = self.head_weights_proj(attention_probs.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (-1,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+        if return_att:
+            return (context_layer, attention_probs)
+        else:
+            return context_layer
+
+    def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
+        if relative_pos is None:
+            q = query_layer.size(-2)
+            relative_pos = build_relative_position(q, key_layer.size(-2), query_layer.device)
+        if relative_pos.dim() == 2:
+            relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+        elif relative_pos.dim() == 3:
+            relative_pos = relative_pos.unsqueeze(1)
+        # bxhxqxk
+        elif relative_pos.dim() != 4:
+            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}")
+
+        att_span = min(max(query_layer.size(-2), key_layer.size(-2)), self.max_relative_positions)
+        relative_pos = relative_pos.long().to(query_layer.device)
+        rel_embeddings = rel_embeddings[
+            self.max_relative_positions - att_span : self.max_relative_positions + att_span, :
+        ].unsqueeze(0)
+        if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+            pos_key_layer = self.pos_proj(rel_embeddings)
+            pos_key_layer = self.transpose_for_scores(pos_key_layer)
+
+        if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+            pos_query_layer = self.pos_q_proj(rel_embeddings)
+            pos_query_layer = self.transpose_for_scores(pos_query_layer)
+
+        score = 0
+        # content->position
+        if "c2p" in self.pos_att_type:
+            c2p_att = torch.matmul(query_layer, pos_key_layer.transpose(-1, -2))
+            c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
+            c2p_att = torch.gather(c2p_att, dim=-1, index=c2p_dynamic_expand(c2p_pos, query_layer, relative_pos))
+            score += c2p_att
+
+        # position->content
+        if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+            pos_query_layer /= math.sqrt(pos_query_layer.size(-1) * scale_factor)
+            if query_layer.size(-2) != key_layer.size(-2):
+                r_pos = build_relative_position(key_layer.size(-2), key_layer.size(-2), query_layer.device)
+            else:
+                r_pos = relative_pos
+            p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
+            if query_layer.size(-2) != key_layer.size(-2):
+                pos_index = relative_pos[:, :, :, 0].unsqueeze(-1)
+
+        if "p2c" in self.pos_att_type:
+            p2c_att = torch.matmul(key_layer, pos_query_layer.transpose(-1, -2))
+            p2c_att = torch.gather(
+                p2c_att, dim=-1, index=p2c_dynamic_expand(p2c_pos, query_layer, key_layer)
+            ).transpose(-1, -2)
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2c_att = torch.gather(p2c_att, dim=-2, index=pos_dynamic_expand(pos_index, p2c_att, key_layer))
+            score += p2c_att
+
+        return score
+
+
+class DebertaEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        pad_token_id = getattr(config, "pad_token_id", 0)
+        self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
+        self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx=pad_token_id)
+
+        self.position_biased_input = getattr(config, "position_biased_input", True)
+        if not self.position_biased_input:
+            self.position_embeddings = None
+        else:
+            self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size)
+
+        if config.type_vocab_size > 0:
+            self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size)
+
+        if self.embedding_size != config.hidden_size:
+            self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False)
+        self.LayerNorm = DebertaLayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+    def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=None, inputs_embeds=None, past_key_values_length=0):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        if self.position_embeddings is not None:
+            position_embeddings = self.position_embeddings(position_ids.long())
+        else:
+            position_embeddings = torch.zeros_like(inputs_embeds)
+
+        embeddings = inputs_embeds
+        if self.position_biased_input:
+            embeddings += position_embeddings
+        if self.config.type_vocab_size > 0:
+            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+            embeddings += token_type_embeddings
+
+        if self.embedding_size != self.config.hidden_size:
+            embeddings = self.embed_proj(embeddings)
+
+        embeddings = self.LayerNorm(embeddings)
+
+        if mask is not None:
+            if mask.dim() != embeddings.dim():
+                if mask.dim() == 4:
+                    mask = mask.squeeze(1).squeeze(1)
+                mask = mask.unsqueeze(2)
+            mask = mask.to(embeddings.dtype)
+
+            embeddings = embeddings * mask
+
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class DebertaPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = DebertaConfig
+    base_model_prefix = "deberta"
+    _keys_to_ignore_on_load_missing = ["position_ids"]
+    _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._register_load_state_dict_pre_hook(self._pre_load_hook)
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _pre_load_hook(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        """
+        Removes the classifier if it doesn't have the correct number of labels.
+        """
+        self_state = self.state_dict()
+        if (
+            ("classifier.weight" in self_state)
+            and ("classifier.weight" in state_dict)
+            and self_state["classifier.weight"].size() != state_dict["classifier.weight"].size()
+        ):
+            logger.warning(
+                f"The checkpoint classifier head has a shape {state_dict['classifier.weight'].size()} and this model "
+                f"classifier head has a shape {self_state['classifier.weight'].size()}. Ignoring the checkpoint "
+                f"weights. You should train your model on new data."
+            )
+            del state_dict["classifier.weight"]
+            if "classifier.bias" in state_dict:
+                del state_dict["classifier.bias"]
+
+
+DEBERTA_START_DOCSTRING = r"""
+    The DeBERTa model was proposed in `DeBERTa: Decoding-enhanced BERT with Disentangled Attention
+    <https://arxiv.org/abs/2006.03654>`_ by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build on top of
+    BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
+    improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
+
+    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
+    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
+    general usage and behavior.```
+
+
+    Parameters:
+        config (:class:`~transformers.DebertaConfig`): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
+            weights.
+"""
+
+DEBERTA_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`):
+            Indices of input sequence tokens in the vocabulary.
+
+            Indices can be obtained using :class:`transformers.DebertaTokenizer`. See
+            :func:`transformers.PreTrainedTokenizer.encode` and :func:`transformers.PreTrainedTokenizer.__call__` for
+            details.
+
+            `What are input IDs? <../glossary.html#input-ids>`__
+        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`{0}`, `optional`):
+            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            `What are attention masks? <../glossary.html#attention-mask>`__
+        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
+            1]``:
+
+            - 0 corresponds to a `sentence A` token,
+            - 1 corresponds to a `sentence B` token.
+
+            `What are token type IDs? <../glossary.html#token-type-ids>`_
+        position_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
+            config.max_position_embeddings - 1]``.
+
+            `What are position IDs? <../glossary.html#position-ids>`_
+        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
+            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+            than the model's internal embedding lookup matrix.
+        output_attentions (:obj:`bool`, `optional`):
+            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
+            tensors for more detail.
+        output_hidden_states (:obj:`bool`, `optional`):
+            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
+            more detail.
+        return_dict (:obj:`bool`, `optional`):
+            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
+    DEBERTA_START_DOCSTRING,
+)
+class DebertaModel(DebertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = DebertaEmbeddings(config)
+        self.encoder = DebertaEncoder(config)
+        self.z_steps = 0
+        self.config = config
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embeddings.word_embeddings = new_embeddings
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        raise NotImplementedError("The prune function is not implemented in DeBERTa model.")
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        past_key_values=None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape            
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        embedding_mask = attention_mask[:, past_key_values_length:].contiguous()
+        if attention_mask is None:
+            # attention_mask = torch.ones(input_shape, device=device)
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            mask=embedding_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values_length,
+        )
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask,
+            output_hidden_states=True,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+            past_key_values=past_key_values,
+        )
+        encoded_layers = encoder_outputs[1]
+
+        if self.z_steps > 1:
+            hidden_states = encoded_layers[-2]
+            layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
+            query_states = encoded_layers[-1]
+            rel_embeddings = self.encoder.get_rel_embedding()
+            attention_mask = self.encoder.get_attention_mask(attention_mask)
+            rel_pos = self.encoder.get_rel_pos(embedding_output)
+            for layer in layers[1:]:
+                query_states = layer(
+                    hidden_states,
+                    attention_mask,
+                    return_att=False,
+                    query_states=query_states,
+                    relative_pos=rel_pos,
+                    rel_embeddings=rel_embeddings,
+                )
+                encoded_layers.append(query_states)
+
+        sequence_output = encoded_layers[-1]
+
+        if not return_dict:
+            return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :]
+
+        return BaseModelOutput(
+            last_hidden_state=sequence_output,
+            hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+@add_start_docstrings("""DeBERTa Model with a `language modeling` head on top. """, DEBERTA_START_DOCSTRING)
+class DebertaForMaskedLM(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.deberta = DebertaModel(config)
+        self.cls = DebertaOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=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, sequence_length)`, `optional`):
+            Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ...,
+            config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored
+            (masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
+        """
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[1:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+# copied from transformers.models.bert.BertPredictionHeadTransform with bert -> deberta
+class DebertaPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+# copied from transformers.models.bert.BertLMPredictionHead with bert -> deberta
+class DebertaLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = DebertaPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+# copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta
+class DebertaOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = DebertaLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+@add_start_docstrings(
+    """
+    DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+    pooled output) e.g. for GLUE tasks.
+    """,
+    DEBERTA_START_DOCSTRING,
+)
+class DebertaForSequenceClassification(DebertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        num_labels = getattr(config, "num_labels", 2)
+        self.num_labels = num_labels
+
+        self.deberta = DebertaModel(config)
+        self.pooler = ContextPooler(config)
+        output_dim = self.pooler.output_dim
+
+        self.classifier = nn.Linear(output_dim, 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)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.deberta.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.deberta.set_input_embeddings(new_embeddings)
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=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.deberta(
+            input_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        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 = 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 = 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,
+            )
+
+
+@add_start_docstrings(
+    """
+    DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+    Named-Entity-Recognition (NER) tasks.
+    """,
+    DEBERTA_START_DOCSTRING,
+)
+class DebertaForTokenClassification(DebertaPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.deberta = DebertaModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        for param in self.deberta.parameters():
+            param.requires_grad = False
+
+        self.init_weights()
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=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, sequence_length)`, `optional`):
+            Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels -
+            1]``.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
+    layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
+    """,
+    DEBERTA_START_DOCSTRING,
+)
+class DebertaForQuestionAnswering(DebertaPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.deberta = DebertaModel(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[1:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

soft_prompt/model/debertaV2.py

@@ -0,0 +1,1509 @@
+# coding=utf-8
+# Copyright 2020 Microsoft and the Hugging Face Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch DeBERTa-v2 model. """
+
+import math
+from collections.abc import Sequence
+
+import numpy as np
+import torch
+from torch import _softmax_backward_data, nn
+from torch.nn import CrossEntropyLoss, LayerNorm
+
+
+from transformers.activations import ACT2FN
+from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers.models.deberta_v2.configuration_deberta_v2 import DebertaV2Config
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "DebertaV2Config"
+_TOKENIZER_FOR_DOC = "DebertaV2Tokenizer"
+_CHECKPOINT_FOR_DOC = "microsoft/deberta-v2-xlarge"
+
+DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "microsoft/deberta-v2-xlarge",
+    "microsoft/deberta-v2-xxlarge",
+    "microsoft/deberta-v2-xlarge-mnli",
+    "microsoft/deberta-v2-xxlarge-mnli",
+]
+
+
+# Copied from transformers.models.deberta.modeling_deberta.ContextPooler
+class ContextPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size)
+        self.dropout = StableDropout(config.pooler_dropout)
+        self.config = config
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+
+        context_token = hidden_states[:, 0]
+        context_token = self.dropout(context_token)
+        pooled_output = self.dense(context_token)
+        pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output)
+        return pooled_output
+
+    @property
+    def output_dim(self):
+        return self.config.hidden_size
+
+
+# Copied from transformers.models.deberta.modeling_deberta.XSoftmax with deberta->deberta_v2
+class XSoftmax(torch.autograd.Function):
+    """
+    Masked Softmax which is optimized for saving memory
+    Args:
+        input (:obj:`torch.tensor`): The input tensor that will apply softmax.
+        mask (:obj:`torch.IntTensor`): The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
+        dim (int): The dimension that will apply softmax
+    Example::
+          >>> import torch
+          >>> from transformers.models.deberta_v2.modeling_deberta_v2 import XSoftmax
+          >>> # Make a tensor
+          >>> x = torch.randn([4,20,100])
+          >>> # Create a mask
+          >>> mask = (x>0).int()
+          >>> y = XSoftmax.apply(x, mask, dim=-1)
+    """
+
+    @staticmethod
+    def forward(self, input, mask, dim):
+        self.dim = dim
+        rmask = ~(mask.bool())
+
+        output = input.masked_fill(rmask, float("-inf"))
+        output = torch.softmax(output, self.dim)
+        output.masked_fill_(rmask, 0)
+        self.save_for_backward(output)
+        return output
+
+    @staticmethod
+    def backward(self, grad_output):
+        (output,) = self.saved_tensors
+        inputGrad = _softmax_backward_data(grad_output, output, self.dim, output)
+        return inputGrad, None, None
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DropoutContext
+class DropoutContext(object):
+    def __init__(self):
+        self.dropout = 0
+        self.mask = None
+        self.scale = 1
+        self.reuse_mask = True
+
+
+# Copied from transformers.models.deberta.modeling_deberta.get_mask
+def get_mask(input, local_context):
+    if not isinstance(local_context, DropoutContext):
+        dropout = local_context
+        mask = None
+    else:
+        dropout = local_context.dropout
+        dropout *= local_context.scale
+        mask = local_context.mask if local_context.reuse_mask else None
+
+    if dropout > 0 and mask is None:
+        mask = (1 - torch.empty_like(input).bernoulli_(1 - dropout)).bool()
+
+    if isinstance(local_context, DropoutContext):
+        if local_context.mask is None:
+            local_context.mask = mask
+
+    return mask, dropout
+
+
+# Copied from transformers.models.deberta.modeling_deberta.XDropout
+class XDropout(torch.autograd.Function):
+    """Optimized dropout function to save computation and memory by using mask operation instead of multiplication."""
+
+    @staticmethod
+    def forward(ctx, input, local_ctx):
+        mask, dropout = get_mask(input, local_ctx)
+        ctx.scale = 1.0 / (1 - dropout)
+        if dropout > 0:
+            ctx.save_for_backward(mask)
+            return input.masked_fill(mask, 0) * ctx.scale
+        else:
+            return input
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.scale > 1:
+            (mask,) = ctx.saved_tensors
+            return grad_output.masked_fill(mask, 0) * ctx.scale, None
+        else:
+            return grad_output, None
+
+
+# Copied from transformers.models.deberta.modeling_deberta.StableDropout
+class StableDropout(nn.Module):
+    """
+    Optimized dropout module for stabilizing the training
+    Args:
+        drop_prob (float): the dropout probabilities
+    """
+
+    def __init__(self, drop_prob):
+        super().__init__()
+        self.drop_prob = drop_prob
+        self.count = 0
+        self.context_stack = None
+
+    def forward(self, x):
+        """
+        Call the module
+        Args:
+            x (:obj:`torch.tensor`): The input tensor to apply dropout
+        """
+        if self.training and self.drop_prob > 0:
+            return XDropout.apply(x, self.get_context())
+        return x
+
+    def clear_context(self):
+        self.count = 0
+        self.context_stack = None
+
+    def init_context(self, reuse_mask=True, scale=1):
+        if self.context_stack is None:
+            self.context_stack = []
+        self.count = 0
+        for c in self.context_stack:
+            c.reuse_mask = reuse_mask
+            c.scale = scale
+
+    def get_context(self):
+        if self.context_stack is not None:
+            if self.count >= len(self.context_stack):
+                self.context_stack.append(DropoutContext())
+            ctx = self.context_stack[self.count]
+            ctx.dropout = self.drop_prob
+            self.count += 1
+            return ctx
+        else:
+            return self.drop_prob
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaSelfOutput with DebertaLayerNorm->LayerNorm
+class DebertaV2SelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaAttention with Deberta->DebertaV2
+class DebertaV2Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = DisentangledSelfAttention(config)
+        self.output = DebertaV2SelfOutput(config)
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        past_key_value=None,
+    ):
+        self_output = self.self(
+            hidden_states,
+            attention_mask,
+            return_att,
+            query_states=query_states,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+            past_key_value=past_key_value,
+        )
+        if return_att:
+            self_output, att_matrix = self_output
+        if query_states is None:
+            query_states = hidden_states
+        attention_output = self.output(self_output, query_states)
+
+        if return_att:
+            return (attention_output, att_matrix)
+        else:
+            return attention_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->DebertaV2
+class DebertaV2Intermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaOutput with DebertaLayerNorm->LayerNorm
+class DebertaV2Output(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaLayer with Deberta->DebertaV2
+class DebertaV2Layer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = DebertaV2Attention(config)
+        self.intermediate = DebertaV2Intermediate(config)
+        self.output = DebertaV2Output(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        past_key_value=None,
+    ):
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        attention_output = self.attention(
+            hidden_states,
+            attention_mask,
+            return_att=return_att,
+            query_states=query_states,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+            past_key_value=self_attn_past_key_value,
+        )
+        if return_att:
+            attention_output, att_matrix = attention_output
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        if return_att:
+            return (layer_output, att_matrix)
+        else:
+            return layer_output
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        kernel_size = getattr(config, "conv_kernel_size", 3)
+        groups = getattr(config, "conv_groups", 1)
+        self.conv_act = getattr(config, "conv_act", "tanh")
+        self.conv = nn.Conv1d(
+            config.hidden_size, config.hidden_size, kernel_size, padding=(kernel_size - 1) // 2, groups=groups
+        )
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def forward(self, hidden_states, residual_states, input_mask):
+        out = self.conv(hidden_states.permute(0, 2, 1).contiguous()).permute(0, 2, 1).contiguous()
+        rmask = (1 - input_mask).bool()
+        out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0)
+        out = ACT2FN[self.conv_act](self.dropout(out))
+
+        layer_norm_input = residual_states + out
+        output = self.LayerNorm(layer_norm_input).to(layer_norm_input)
+
+        if input_mask is None:
+            output_states = output
+        else:
+            if input_mask.dim() != layer_norm_input.dim():
+                if input_mask.dim() == 4:
+                    input_mask = input_mask.squeeze(1).squeeze(1)
+                input_mask = input_mask.unsqueeze(2)
+
+            input_mask = input_mask.to(output.dtype)
+            output_states = output * input_mask
+
+        return output_states
+
+
+class DebertaV2Encoder(nn.Module):
+    """Modified BertEncoder with relative position bias support"""
+
+    def __init__(self, config):
+        super().__init__()
+
+        self.layer = nn.ModuleList([DebertaV2Layer(config) for _ in range(config.num_hidden_layers)])
+        self.relative_attention = getattr(config, "relative_attention", False)
+
+        if self.relative_attention:
+            self.max_relative_positions = getattr(config, "max_relative_positions", -1)
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = config.max_position_embeddings
+
+            self.position_buckets = getattr(config, "position_buckets", -1)
+            pos_ebd_size = self.max_relative_positions * 2
+
+            if self.position_buckets > 0:
+                pos_ebd_size = self.position_buckets * 2
+
+            self.rel_embeddings = nn.Embedding(pos_ebd_size, config.hidden_size)
+
+        self.norm_rel_ebd = [x.strip() for x in getattr(config, "norm_rel_ebd", "none").lower().split("|")]
+
+        if "layer_norm" in self.norm_rel_ebd:
+            self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True)
+
+        self.conv = ConvLayer(config) if getattr(config, "conv_kernel_size", 0) > 0 else None
+
+    def get_rel_embedding(self):
+        rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
+        if rel_embeddings is not None and ("layer_norm" in self.norm_rel_ebd):
+            rel_embeddings = self.LayerNorm(rel_embeddings)
+        return rel_embeddings
+
+    def get_attention_mask(self, attention_mask):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
+            attention_mask = attention_mask.byte()
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask.unsqueeze(1)
+
+        return attention_mask
+
+    def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
+        if self.relative_attention and relative_pos is None:
+            q = query_states.size(-2) if query_states is not None else hidden_states.size(-2)
+            relative_pos = build_relative_position(
+                q, hidden_states.size(-2), bucket_size=self.position_buckets, max_position=self.max_relative_positions
+            )
+        return relative_pos
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        output_hidden_states=True,
+        output_attentions=False,
+        query_states=None,
+        relative_pos=None,
+        return_dict=True,
+        past_key_values=None,
+    ):
+        if attention_mask.dim() <= 2:
+            input_mask = attention_mask
+        else:
+            input_mask = (attention_mask.sum(-2) > 0).byte()
+        attention_mask = self.get_attention_mask(attention_mask)
+        relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        if isinstance(hidden_states, Sequence): # False
+            next_kv = hidden_states[0]
+        else:
+            next_kv = hidden_states
+        rel_embeddings = self.get_rel_embedding()
+        output_states = next_kv
+        for i, layer_module in enumerate(self.layer):
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (output_states,)
+
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            output_states = layer_module(
+                next_kv,
+                attention_mask,
+                output_attentions,
+                query_states=query_states,
+                relative_pos=relative_pos,
+                rel_embeddings=rel_embeddings,
+                past_key_value=past_key_value,
+            )
+            if output_attentions:
+                output_states, att_m = output_states
+
+            if i == 0 and self.conv is not None:
+                if past_key_values is not None:
+                    past_key_value_length = past_key_values[0][0].shape[2]
+                    input_mask = input_mask[:, past_key_value_length:].contiguous()
+                output_states = self.conv(hidden_states, output_states, input_mask)
+
+            if query_states is not None:
+                query_states = output_states
+                if isinstance(hidden_states, Sequence):
+                    next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
+            else:
+                next_kv = output_states
+
+            if output_attentions:
+                all_attentions = all_attentions + (att_m,)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (output_states,)
+
+        if not return_dict:
+            return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions
+        )
+
+
+def make_log_bucket_position(relative_pos, bucket_size, max_position):
+    sign = np.sign(relative_pos)
+    mid = bucket_size // 2
+    abs_pos = np.where((relative_pos < mid) & (relative_pos > -mid), mid - 1, np.abs(relative_pos))
+    log_pos = np.ceil(np.log(abs_pos / mid) / np.log((max_position - 1) / mid) * (mid - 1)) + mid
+    bucket_pos = np.where(abs_pos <= mid, relative_pos, log_pos * sign).astype(np.int)
+    return bucket_pos
+
+
+def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1):
+    """
+    Build relative position according to the query and key
+    We assume the absolute position of query :math:`P_q` is range from (0, query_size) and the absolute position of key
+    :math:`P_k` is range from (0, key_size), The relative positions from query to key is :math:`R_{q \\rightarrow k} =
+    P_q - P_k`
+    Args:
+        query_size (int): the length of query
+        key_size (int): the length of key
+        bucket_size (int): the size of position bucket
+        max_position (int): the maximum allowed absolute position
+    Return:
+        :obj:`torch.LongTensor`: A tensor with shape [1, query_size, key_size]
+    """
+    q_ids = np.arange(0, query_size)
+    k_ids = np.arange(0, key_size)
+    rel_pos_ids = q_ids[:, None] - np.tile(k_ids, (q_ids.shape[0], 1))
+    if bucket_size > 0 and max_position > 0:
+        rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
+    rel_pos_ids = torch.tensor(rel_pos_ids, dtype=torch.long)
+    rel_pos_ids = rel_pos_ids[:query_size, :]
+    rel_pos_ids = rel_pos_ids.unsqueeze(0)
+    return rel_pos_ids
+
+
+@torch.jit.script
+# Copied from transformers.models.deberta.modeling_deberta.c2p_dynamic_expand
+def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
+    return c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)])
+
+
+@torch.jit.script
+# Copied from transformers.models.deberta.modeling_deberta.p2c_dynamic_expand
+def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
+    return c2p_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)])
+
+
+@torch.jit.script
+# Copied from transformers.models.deberta.modeling_deberta.pos_dynamic_expand
+def pos_dynamic_expand(pos_index, p2c_att, key_layer):
+    return pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2)))
+
+
+class DisentangledSelfAttention(nn.Module):
+    """
+    Disentangled self-attention module
+    Parameters:
+        config (:obj:`DebertaV2Config`):
+            A model config class instance with the configuration to build a new model. The schema is similar to
+            `BertConfig`, for more details, please refer :class:`~transformers.DebertaV2Config`
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({config.num_attention_heads})"
+            )
+        self.num_attention_heads = config.num_attention_heads
+        _attention_head_size = config.hidden_size // config.num_attention_heads
+        self.attention_head_size = getattr(config, "attention_head_size", _attention_head_size)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+
+        self.share_att_key = getattr(config, "share_att_key", False)
+        self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
+        self.relative_attention = getattr(config, "relative_attention", False)
+
+        if self.relative_attention:
+            self.position_buckets = getattr(config, "position_buckets", -1)
+            self.max_relative_positions = getattr(config, "max_relative_positions", -1)
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = config.max_position_embeddings
+            self.pos_ebd_size = self.max_relative_positions
+            if self.position_buckets > 0:
+                self.pos_ebd_size = self.position_buckets
+
+            self.pos_dropout = StableDropout(config.hidden_dropout_prob)
+
+            if not self.share_att_key:
+                if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+                    self.pos_key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+                if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+                    self.pos_query_proj = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = StableDropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x, attention_heads, past_key_value=None):
+        new_x_shape = x.size()[:-1] + (attention_heads, -1)
+        x = x.view(*new_x_shape)
+        x = x.permute(0, 2, 1, 3)
+        if past_key_value is not None:
+            x = torch.cat([past_key_value, x], dim=2)
+        new_x_shape = x.shape
+        return x.contiguous().view(-1, new_x_shape[2], new_x_shape[-1])
+        # return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        past_key_value=None,
+    ):
+        """
+        Call the module
+        Args:
+            hidden_states (:obj:`torch.FloatTensor`):
+                Input states to the module usually the output from previous layer, it will be the Q,K and V in
+                `Attention(Q,K,V)`
+            attention_mask (:obj:`torch.ByteTensor`):
+                An attention mask matrix of shape [`B`, `N`, `N`] where `B` is the batch size, `N` is the maximum
+                sequence length in which element [i,j] = `1` means the `i` th token in the input can attend to the `j`
+                th token.
+            return_att (:obj:`bool`, optional):
+                Whether return the attention matrix.
+            query_states (:obj:`torch.FloatTensor`, optional):
+                The `Q` state in `Attention(Q,K,V)`.
+            relative_pos (:obj:`torch.LongTensor`):
+                The relative position encoding between the tokens in the sequence. It's of shape [`B`, `N`, `N`] with
+                values ranging in [`-max_relative_positions`, `max_relative_positions`].
+            rel_embeddings (:obj:`torch.FloatTensor`):
+                The embedding of relative distances. It's a tensor of shape [:math:`2 \\times
+                \\text{max_relative_positions}`, `hidden_size`].
+        """
+        if query_states is None:
+            query_states = hidden_states
+        
+        past_key_value_length = past_key_value.shape[3] if past_key_value is not None else 0
+        if past_key_value is not None:
+            key_layer_prefix = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads, past_key_value=past_key_value[0])
+            # value_layer_prefix = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads, past_key_value=past_key_value[1])
+        
+        query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
+        key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
+        value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads, past_key_value=past_key_value[1])
+
+        rel_att = None
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        scale_factor = 1
+        if "c2p" in self.pos_att_type:
+            scale_factor += 1
+        if "p2c" in self.pos_att_type:
+            scale_factor += 1
+        if "p2p" in self.pos_att_type:
+            scale_factor += 1
+        scale = math.sqrt(query_layer.size(-1) * scale_factor)
+        # attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)) / scale
+        attention_scores = torch.bmm(query_layer, key_layer_prefix.transpose(-1, -2)) / scale
+        
+        if self.relative_attention:
+            rel_embeddings = self.pos_dropout(rel_embeddings)
+            rel_att = self.disentangled_attention_bias(
+                query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+            )
+
+        if rel_att is not None:
+            if past_key_value is not None:
+                att_shape = rel_att.shape[:-1] + (past_key_value_length,)
+                prefix_att = torch.zeros(*att_shape).to(rel_att.device)
+                attention_scores = attention_scores + torch.cat([prefix_att, rel_att], dim=-1)
+            else:
+                attention_scores = attention_scores + rel_att
+        # print(attention_scores.shape)
+        attention_scores = attention_scores
+        attention_scores = attention_scores.view(
+            -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)
+        )
+
+        # bsz x height x length x dimension
+        attention_mask = attention_mask[:,:, past_key_value_length:,:]
+
+        attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
+        attention_probs = self.dropout(attention_probs)
+
+        context_layer = torch.bmm(
+            attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer
+        )
+        context_layer = (
+            context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1))
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        new_context_layer_shape = context_layer.size()[:-2] + (-1,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+        if return_att:
+            return (context_layer, attention_probs)
+        else:
+            return context_layer
+
+    def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
+        if relative_pos is None:
+            q = query_layer.size(-2)
+            relative_pos = build_relative_position(
+                q, key_layer.size(-2), bucket_size=self.position_buckets, max_position=self.max_relative_positions
+            )
+        if relative_pos.dim() == 2:
+            relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+        elif relative_pos.dim() == 3:
+            relative_pos = relative_pos.unsqueeze(1)
+        # bsz x height x query x key
+        elif relative_pos.dim() != 4:
+            raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}")
+
+        att_span = self.pos_ebd_size
+        relative_pos = relative_pos.long().to(query_layer.device)
+
+        rel_embeddings = rel_embeddings[self.pos_ebd_size - att_span : self.pos_ebd_size + att_span, :].unsqueeze(0)
+        if self.share_att_key: # True
+            pos_query_layer = self.transpose_for_scores(
+                self.query_proj(rel_embeddings), self.num_attention_heads
+            ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)
+            pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads).repeat(
+                query_layer.size(0) // self.num_attention_heads, 1, 1
+            )
+        else:
+            if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+                pos_key_layer = self.transpose_for_scores(
+                    self.pos_key_proj(rel_embeddings), self.num_attention_heads
+                ).repeat(
+                    query_layer.size(0) // self.num_attention_heads, 1, 1
+                )  # .split(self.all_head_size, dim=-1)
+            if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+                pos_query_layer = self.transpose_for_scores(
+                    self.pos_query_proj(rel_embeddings), self.num_attention_heads
+                ).repeat(
+                    query_layer.size(0) // self.num_attention_heads, 1, 1
+                )  # .split(self.all_head_size, dim=-1)
+
+        score = 0
+        # content->position
+        if "c2p" in self.pos_att_type:
+            scale = math.sqrt(pos_key_layer.size(-1) * scale_factor)
+            c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2))
+            c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
+            c2p_att = torch.gather(
+                c2p_att,
+                dim=-1,
+                index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]),
+            )
+            score += c2p_att / scale
+
+        # position->content
+        if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+            scale = math.sqrt(pos_query_layer.size(-1) * scale_factor)
+            if key_layer.size(-2) != query_layer.size(-2):
+                r_pos = build_relative_position(
+                    key_layer.size(-2),
+                    key_layer.size(-2),
+                    bucket_size=self.position_buckets,
+                    max_position=self.max_relative_positions,
+                ).to(query_layer.device)
+                r_pos = r_pos.unsqueeze(0)
+            else:
+                r_pos = relative_pos
+
+            p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
+            if query_layer.size(-2) != key_layer.size(-2):
+                pos_index = relative_pos[:, :, :, 0].unsqueeze(-1)
+
+        if "p2c" in self.pos_att_type:
+            p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2))
+            p2c_att = torch.gather(
+                p2c_att,
+                dim=-1,
+                index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]),
+            ).transpose(-1, -2)
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2c_att = torch.gather(
+                    p2c_att,
+                    dim=-2,
+                    index=pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2))),
+                )
+            score += p2c_att / scale
+
+        # position->position
+        if "p2p" in self.pos_att_type:
+            pos_query = pos_query_layer[:, :, att_span:, :]
+            p2p_att = torch.matmul(pos_query, pos_key_layer.transpose(-1, -2))
+            p2p_att = p2p_att.expand(query_layer.size()[:2] + p2p_att.size()[2:])
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2p_att = torch.gather(
+                    p2p_att,
+                    dim=-2,
+                    index=pos_index.expand(query_layer.size()[:2] + (pos_index.size(-2), p2p_att.size(-1))),
+                )
+            p2p_att = torch.gather(
+                p2p_att,
+                dim=-1,
+                index=c2p_pos.expand(
+                    [query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)]
+                ),
+            )
+            score += p2p_att
+
+        return score
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaEmbeddings with DebertaLayerNorm->LayerNorm
+class DebertaV2Embeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        pad_token_id = getattr(config, "pad_token_id", 0)
+        self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
+        self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx=pad_token_id)
+
+        self.position_biased_input = getattr(config, "position_biased_input", True)
+        if not self.position_biased_input:
+            self.position_embeddings = None
+        else:
+            self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size)
+
+        if config.type_vocab_size > 0:
+            self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size)
+
+        if self.embedding_size != config.hidden_size:
+            self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False)
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+    def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=None, inputs_embeds=None, past_key_values_length=0,):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            # position_ids = self.position_ids[:, :seq_length]
+            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        if self.position_embeddings is not None:
+            position_embeddings = self.position_embeddings(position_ids.long())
+        else:
+            position_embeddings = torch.zeros_like(inputs_embeds)
+
+        embeddings = inputs_embeds
+        if self.position_biased_input:
+            embeddings += position_embeddings
+        if self.config.type_vocab_size > 0:
+            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+            embeddings += token_type_embeddings
+
+        if self.embedding_size != self.config.hidden_size:
+            embeddings = self.embed_proj(embeddings)
+
+        embeddings = self.LayerNorm(embeddings)
+
+        if mask is not None:
+            if mask.dim() != embeddings.dim():
+                if mask.dim() == 4:
+                    mask = mask.squeeze(1).squeeze(1)
+                mask = mask.unsqueeze(2)
+            mask = mask.to(embeddings.dtype)
+
+            embeddings = embeddings * mask
+
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaPreTrainedModel with Deberta->DebertaV2
+class DebertaV2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = DebertaV2Config
+    base_model_prefix = "deberta"
+    _keys_to_ignore_on_load_missing = ["position_ids"]
+    _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._register_load_state_dict_pre_hook(self._pre_load_hook)
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _pre_load_hook(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        """
+        Removes the classifier if it doesn't have the correct number of labels.
+        """
+        self_state = self.state_dict()
+        if (
+            ("classifier.weight" in self_state)
+            and ("classifier.weight" in state_dict)
+            and self_state["classifier.weight"].size() != state_dict["classifier.weight"].size()
+        ):
+            logger.warning(
+                f"The checkpoint classifier head has a shape {state_dict['classifier.weight'].size()} and this model "
+                f"classifier head has a shape {self_state['classifier.weight'].size()}. Ignoring the checkpoint "
+                f"weights. You should train your model on new data."
+            )
+            del state_dict["classifier.weight"]
+            if "classifier.bias" in state_dict:
+                del state_dict["classifier.bias"]
+
+
+DEBERTA_START_DOCSTRING = r"""
+    The DeBERTa model was proposed in `DeBERTa: Decoding-enhanced BERT with Disentangled Attention
+    <https://arxiv.org/abs/2006.03654>`_ by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build on top of
+    BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
+    improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
+    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
+    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
+    general usage and behavior.```
+    Parameters:
+        config (:class:`~transformers.DebertaV2Config`): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
+            weights.
+"""
+
+DEBERTA_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`):
+            Indices of input sequence tokens in the vocabulary.
+            Indices can be obtained using :class:`transformers.DebertaV2Tokenizer`. See
+            :func:`transformers.PreTrainedTokenizer.encode` and :func:`transformers.PreTrainedTokenizer.__call__` for
+            details.
+            `What are input IDs? <../glossary.html#input-ids>`__
+        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`{0}`, `optional`):
+            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+            `What are attention masks? <../glossary.html#attention-mask>`__
+        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
+            1]``:
+            - 0 corresponds to a `sentence A` token,
+            - 1 corresponds to a `sentence B` token.
+            `What are token type IDs? <../glossary.html#token-type-ids>`_
+        position_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
+            config.max_position_embeddings - 1]``.
+            `What are position IDs? <../glossary.html#position-ids>`_
+        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
+            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+            than the model's internal embedding lookup matrix.
+        output_attentions (:obj:`bool`, `optional`):
+            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
+            tensors for more detail.
+        output_hidden_states (:obj:`bool`, `optional`):
+            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
+            more detail.
+        return_dict (:obj:`bool`, `optional`):
+            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
+    DEBERTA_START_DOCSTRING,
+)
+# Copied from transformers.models.deberta.modeling_deberta.DebertaModel with Deberta->DebertaV2
+class DebertaV2Model(DebertaV2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = DebertaV2Embeddings(config)
+        self.encoder = DebertaV2Encoder(config)
+        self.z_steps = 0
+        self.config = config
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embeddings.word_embeddings = new_embeddings
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        raise NotImplementedError("The prune function is not implemented in DeBERTa model.")
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        past_key_values=None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        embedding_mask = torch.ones(input_shape, device=device)
+        if attention_mask is None:
+            # attention_mask = torch.ones(input_shape, device=device)
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            # mask=attention_mask,
+            mask=embedding_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values_length, # Ongoing
+        )
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask,
+            output_hidden_states=True,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+            past_key_values=past_key_values, # Ongoing
+        )
+        encoded_layers = encoder_outputs[1]
+
+        if self.z_steps > 1:
+            hidden_states = encoded_layers[-2]
+            layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
+            query_states = encoded_layers[-1]
+            rel_embeddings = self.encoder.get_rel_embedding()
+            attention_mask = self.encoder.get_attention_mask(attention_mask)
+            rel_pos = self.encoder.get_rel_pos(embedding_output)
+            for layer in layers[1:]:
+                query_states = layer(
+                    hidden_states,
+                    attention_mask,
+                    return_att=False,
+                    query_states=query_states,
+                    relative_pos=rel_pos,
+                    rel_embeddings=rel_embeddings,
+                )
+                encoded_layers.append(query_states)
+
+        sequence_output = encoded_layers[-1]
+
+        if not return_dict:
+            return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :]
+
+        return BaseModelOutput(
+            last_hidden_state=sequence_output,
+            hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+@add_start_docstrings("""DeBERTa Model with a `language modeling` head on top. """, DEBERTA_START_DOCSTRING)
+# Copied from transformers.models.deberta.modeling_deberta.DebertaForMaskedLM with Deberta->DebertaV2
+class DebertaV2ForMaskedLM(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.deberta = DebertaV2Model(config)
+        self.cls = DebertaV2OnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=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, sequence_length)`, `optional`):
+            Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ...,
+            config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored
+            (masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
+        """
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[1:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+# copied from transformers.models.bert.BertPredictionHeadTransform with bert -> deberta
+class DebertaV2PredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+# copied from transformers.models.bert.BertLMPredictionHead with bert -> deberta
+class DebertaV2LMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = DebertaV2PredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+# copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta
+class DebertaV2OnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = DebertaV2LMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+@add_start_docstrings(
+    """
+    DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+    pooled output) e.g. for GLUE tasks.
+    """,
+    DEBERTA_START_DOCSTRING,
+)
+# Copied from transformers.models.deberta.modeling_deberta.DebertaForSequenceClassification with Deberta->DebertaV2
+class DebertaV2ForSequenceClassification(DebertaV2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        num_labels = getattr(config, "num_labels", 2)
+        self.num_labels = num_labels
+
+        self.deberta = DebertaV2Model(config)
+        self.pooler = ContextPooler(config)
+        output_dim = self.pooler.output_dim
+
+        self.classifier = nn.Linear(output_dim, 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)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.deberta.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.deberta.set_input_embeddings(new_embeddings)
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=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.deberta(
+            input_ids,
+            token_type_ids=token_type_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        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 = 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 = 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,
+            )
+
+
+@add_start_docstrings(
+    """
+    DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+    Named-Entity-Recognition (NER) tasks.
+    """,
+    DEBERTA_START_DOCSTRING,
+)
+# Copied from transformers.models.deberta.modeling_deberta.DebertaForTokenClassification with Deberta->DebertaV2
+class DebertaV2ForTokenClassification(DebertaV2PreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.deberta = DebertaV2Model(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+        for param in self.deberta.parameters():
+            param.requires_grad = False
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        past_key_values=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels -
+            1]``.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
+    layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
+    """,
+    DEBERTA_START_DOCSTRING,
+)
+# Copied from transformers.models.deberta.modeling_deberta.DebertaForQuestionAnswering with Deberta->DebertaV2
+class DebertaV2ForQuestionAnswering(DebertaV2PreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.deberta = DebertaV2Model(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[1:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

soft_prompt/model/multiple_choice.py

@@ -0,0 +1,710 @@
+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.modeling_outputs import MultipleChoiceModelOutput, BaseModelOutput, Seq2SeqLMOutput
+
+from model.prefix_encoder import PrefixEncoder
+from model.deberta import DebertaModel, DebertaPreTrainedModel, ContextPooler, StableDropout
+from model import utils
+
+
+class BertForMultipleChoice(BertPreTrainedModel):
+    """BERT model for multiple choice tasks.
+    This module is composed of the BERT model with a linear layer on top of
+    the pooled output.
+
+    Params:
+        `config`: a BertConfig class instance with the configuration to build a new model.
+        `num_choices`: the number of classes for the classifier. Default = 2.
+
+    Inputs:
+        `input_ids`: a torch.LongTensor of shape [batch_size, num_choices, sequence_length]
+            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
+            `extract_features.py`, `run_classifier.py` and `run_squad.py`)
+        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, num_choices, sequence_length]
+            with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A`
+            and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
+        `attention_mask`: an optional torch.LongTensor of shape [batch_size, num_choices, sequence_length] with indices
+            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
+            input sequence length in the current batch. It's the mask that we typically use for attention when
+            a batch has varying length sentences.
+        `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
+            with indices selected in [0, ..., num_choices].
+
+    Outputs:
+        if `labels` is not `None`:
+            Outputs the CrossEntropy classification loss of the output with the labels.
+        if `labels` is `None`:
+            Outputs the classification logits of shape [batch_size, num_labels].
+
+    Example usage:
+    ```python
+    # Already been converted into WordPiece token ids
+    input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]], [[12, 16, 42], [14, 28, 57]]])
+    input_mask = torch.LongTensor([[[1, 1, 1], [1, 1, 0]],[[1,1,0], [1, 0, 0]]])
+    token_type_ids = torch.LongTensor([[[0, 0, 1], [0, 1, 0]],[[0, 1, 1], [0, 0, 1]]])
+    config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
+        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
+
+    num_choices = 2
+
+    model = BertForMultipleChoice(config, num_choices)
+    logits = model(input_ids, token_type_ids, input_mask)
+    ```
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = BertModel(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, 1)
+
+        self.init_weights()
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    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,
+        use_base_grad=False
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        batch_size, num_choices = input_ids.shape[:2]
+
+        input_ids = input_ids.reshape(-1, input_ids.size(-1))
+        token_type_ids = token_type_ids.reshape(-1, token_type_ids.size(-1))
+        attention_mask = attention_mask.reshape(-1, attention_mask.size(-1))
+        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        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)
+        reshaped_logits = logits.reshape(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class BertPrefixForMultipleChoice(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, 1)
+
+        for param in self.bert.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.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
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+    
+    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)
+        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,
+        use_base_grad=False
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        batch_size, num_choices = input_ids.shape[:2] if input_ids is not None else inputs_embeds[:2]
+
+        input_ids = input_ids.reshape(-1, input_ids.size(-1)) if input_ids is not None else None
+        token_type_ids = token_type_ids.reshape(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        attention_mask = attention_mask.reshape(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        past_key_values = self.get_prompt(batch_size=batch_size * num_choices)
+        prefix_attention_mask = torch.ones(batch_size * num_choices, 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]
+
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+        reshaped_logits = logits.reshape(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class RobertaPrefixForMultipleChoice(RobertaPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.roberta = RobertaModel(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, 1)
+
+        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))
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        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,
+        token_type_ids=None,
+        attention_mask=None,
+        labels=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        use_base_grad=False
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the multiple choice classification loss. Indices should be in ``[0, ...,
+            num_choices-1]`` where :obj:`num_choices` is the size of the second dimension of the input tensors. (See
+            :obj:`input_ids` above)
+        """
+        utils.use_grad(self.roberta, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        batch_size, num_choices = input_ids.shape[:2] if input_ids is not None else inputs_embeds.shape[:2]
+
+        flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+        flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        flat_inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        past_key_values = self.get_prompt(batch_size=batch_size * num_choices)
+        prefix_attention_mask = torch.ones(batch_size * num_choices, self.pre_seq_len).to(self.roberta.device)
+        flat_attention_mask = torch.cat((prefix_attention_mask, flat_attention_mask), dim=1)
+
+        outputs = self.roberta(
+            flat_input_ids,
+            position_ids=flat_position_ids,
+            token_type_ids=flat_token_type_ids,
+            attention_mask=flat_attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=flat_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)
+        reshaped_logits = logits.view(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class DebertaPrefixForMultipleChoice(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, 1)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.init_weights()
+
+        for param in self.deberta.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)
+
+        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))
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+    
+    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)
+        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,
+        use_base_grad=False
+    ):
+        utils.use_grad(self.deberta, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        batch_size, num_choices = input_ids.shape[:2] if input_ids is not None else inputs_embeds.shape[:2]
+        
+        flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+        flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        flat_inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        past_key_values = self.get_prompt(batch_size=batch_size * num_choices)
+        prefix_attention_mask = torch.ones(batch_size * num_choices, self.pre_seq_len).to(self.deberta.device)
+        flat_attention_mask = torch.cat((prefix_attention_mask, flat_attention_mask), dim=1)
+
+        outputs = self.deberta(
+            flat_input_ids,
+            attention_mask=flat_attention_mask,
+            token_type_ids=flat_token_type_ids,
+            position_ids=flat_position_ids,
+            inputs_embeds=flat_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)
+        reshaped_logits = logits.view(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class BertPromptForMultipleChoice(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        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, 1)
+
+        for param in self.bert.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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        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
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+    
+    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,
+        use_base_grad=False
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        batch_size, num_choices = input_ids.shape[:2] if input_ids is not None else inputs_embeds[:2]
+
+        input_ids = input_ids.reshape(-1, input_ids.size(-1)) if input_ids is not None else None
+        token_type_ids = token_type_ids.reshape(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        attention_mask = attention_mask.reshape(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        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 * num_choices)
+        inputs_embeds = torch.cat((prompts, raw_embedding), dim=1)
+
+        prefix_attention_mask = torch.ones(batch_size * num_choices, self.pre_seq_len).to(self.bert.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.bert(
+            attention_mask=attention_mask,
+            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)
+        reshaped_logits = logits.reshape(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class RobertaPromptForMultipleChoice(RobertaPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        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, 1)
+
+        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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        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))
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    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,
+        token_type_ids=None,
+        attention_mask=None,
+        labels=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        use_base_grad=False
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the multiple choice classification loss. Indices should be in ``[0, ...,
+            num_choices-1]`` where :obj:`num_choices` is the size of the second dimension of the input tensors. (See
+            :obj:`input_ids` above)
+        """
+        utils.use_grad(self.roberta, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        batch_size, num_choices = input_ids.shape[:2] if input_ids is not None else inputs_embeds.shape[:2]
+
+        input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        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 * num_choices)
+        inputs_embeds = torch.cat((prompts, raw_embedding), dim=1)
+        prefix_attention_mask = torch.ones(batch_size * num_choices, self.pre_seq_len).to(self.roberta.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.roberta(
+            attention_mask=attention_mask,
+            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)
+        reshaped_logits = logits.view(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

soft_prompt/model/prefix_encoder.py

@@ -0,0 +1,33 @@
+import torch
+
+
+class PrefixEncoder(torch.nn.Module):
+    r'''
+    The torch.nn model to encode the prefix
+
+    Input shape: (batch-size, prefix-length)
+
+    Output shape: (batch-size, prefix-length, 2*layers*hidden)
+    '''
+    def __init__(self, config):
+        super().__init__()
+        self.prefix_projection = config.prefix_projection
+        if self.prefix_projection:
+            # Use a two-layer MLP to encode the prefix
+            self.embedding = torch.nn.Embedding(config.pre_seq_len, config.hidden_size)
+            self.trans = torch.nn.Sequential(
+                torch.nn.Linear(config.hidden_size, config.prefix_hidden_size),
+                torch.nn.Tanh(),
+                torch.nn.Linear(config.prefix_hidden_size, config.num_hidden_layers * 2 * config.hidden_size)
+            )
+        else:
+            self.embedding = torch.nn.Embedding(config.pre_seq_len, config.num_hidden_layers * 2 * config.hidden_size)
+
+    def forward(self, prefix: torch.Tensor):
+        device = next(self.embedding.parameters()).device
+        if self.prefix_projection:
+            prefix_tokens = self.embedding(prefix.to(device))
+            past_key_values = self.trans(prefix_tokens)
+        else:
+            past_key_values = self.embedding(prefix.to(device))
+        return past_key_values

soft_prompt/model/question_answering.py

@@ -0,0 +1,455 @@
+import torch
+import torch.nn
+from torch.nn import CrossEntropyLoss
+from transformers import BertPreTrainedModel, BertModel, RobertaPreTrainedModel, RobertaModel
+from transformers.modeling_outputs import QuestionAnsweringModelOutput
+
+from model.prefix_encoder import PrefixEncoder
+from model.deberta import DebertaPreTrainedModel, DebertaModel
+
+class BertForQuestionAnswering(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.qa_outputs = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        for param in self.bert.parameters():
+            param.requires_grad = False
+
+        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,
+        start_positions=None,
+        end_positions=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[2:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class BertPrefixForQuestionAnswering(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        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.bert = BertModel(config, add_pooling_layer=False)
+        self.qa_outputs = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.prefix_encoder = PrefixEncoder(config)
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+
+        for param in self.bert.parameters():
+            param.requires_grad = False
+
+        self.init_weights()
+
+    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(
+            bsz,
+            seqlen,
+            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,
+        start_positions=None,
+        end_positions=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        """
+        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.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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[2:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class RobertaPrefixModelForQuestionAnswering(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        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.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.qa_outputs = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.prefix_encoder = PrefixEncoder(config)
+        self.prefix_tokens = torch.arange(self.pre_seq_len).long()
+
+        for param in self.roberta.parameters():
+            param.requires_grad = False
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        bsz, seqlen, _ = past_key_values.shape
+        past_key_values = past_key_values.view(
+            bsz,
+            seqlen,
+            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,
+        start_positions=None,
+        end_positions=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        """
+        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)
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[2:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class DebertaPrefixModelForQuestionAnswering(DebertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.deberta = DebertaModel(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.qa_outputs = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.init_weights()
+
+        for param in self.deberta.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
+
+        # Use a two layered MLP to encode the prefix
+        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 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,
+        start_positions=None,
+        end_positions=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (:obj:`sequence_length`). Position outside of the
+            sequence are not taken into account for computing the loss.
+        """
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[2:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

soft_prompt/model/roberta.py

@@ -0,0 +1,1588 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch RoBERTa model."""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN, gelu
+from ...modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    MaskedLMOutput,
+    MultipleChoiceModelOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
+from ...utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from .configuration_roberta import RobertaConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "roberta-base"
+_CONFIG_FOR_DOC = "RobertaConfig"
+
+ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "roberta-base",
+    "roberta-large",
+    "roberta-large-mnli",
+    "distilroberta-base",
+    "roberta-base-openai-detector",
+    "roberta-large-openai-detector",
+    # See all RoBERTa models at https://huggingface.co/models?filter=roberta
+]
+
+
+class RobertaEmbeddings(nn.Module):
+    """
+    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
+    """
+
+    # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+        self.register_buffer(
+            "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+        )
+
+        # End copy
+        self.padding_idx = config.pad_token_id
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
+        )
+
+    def forward(
+        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
+    ):
+        if position_ids is None:
+            if input_ids is not None:
+                # Create the position ids from the input token ids. Any padded tokens remain padded.
+                position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)
+            else:
+                position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)
+
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+        # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+        # issue #5664
+        if token_type_ids is None:
+            if hasattr(self, "token_type_ids"):
+                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
+                token_type_ids = buffered_token_type_ids_expanded
+            else:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = inputs_embeds + token_type_embeddings
+        if self.position_embedding_type == "absolute":
+            position_embeddings = self.position_embeddings(position_ids)
+            embeddings += position_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+    def create_position_ids_from_inputs_embeds(self, inputs_embeds):
+        """
+        We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
+
+        Args:
+            inputs_embeds: torch.Tensor
+
+        Returns: torch.Tensor
+        """
+        input_shape = inputs_embeds.size()[:-1]
+        sequence_length = input_shape[1]
+
+        position_ids = torch.arange(
+            self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
+        )
+        return position_ids.unsqueeze(0).expand(input_shape)
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Roberta
+class RobertaSelfAttention(nn.Module):
+    def __init__(self, config, position_embedding_type=None):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({config.num_attention_heads})"
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.position_embedding_type = position_embedding_type or getattr(
+            config, "position_embedding_type", "absolute"
+        )
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            self.max_position_embeddings = config.max_position_embeddings
+            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
+
+        self.is_decoder = config.is_decoder
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+        mixed_query_layer = self.query(hidden_states)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        if is_cross_attention and past_key_value is not None:
+            # reuse k,v, cross_attentions
+            key_layer = past_key_value[0]
+            value_layer = past_key_value[1]
+            attention_mask = encoder_attention_mask
+        elif is_cross_attention:
+            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            attention_mask = encoder_attention_mask
+        elif past_key_value is not None:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        use_cache = past_key_value is not None
+        if self.is_decoder:
+            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+            # Further calls to cross_attention layer can then reuse all cross-attention
+            # key/value_states (first "if" case)
+            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+            # all previous decoder key/value_states. Further calls to uni-directional self-attention
+            # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+            # if encoder bi-directional self-attention `past_key_value` is always `None`
+            past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            query_length, key_length = query_layer.shape[2], key_layer.shape[2]
+            if use_cache:
+                position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
+                    -1, 1
+                )
+            else:
+                position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
+            position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
+            distance = position_ids_l - position_ids_r
+
+            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
+            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility
+
+            if self.position_embedding_type == "relative_key":
+                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores
+            elif self.position_embedding_type == "relative_key_query":
+                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs = attention_probs * head_mask
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        if self.is_decoder:
+            outputs = outputs + (past_key_value,)
+        return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
+class RobertaSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Roberta
+class RobertaAttention(nn.Module):
+    def __init__(self, config, position_embedding_type=None):
+        super().__init__()
+        self.self = RobertaSelfAttention(config, position_embedding_type=position_embedding_type)
+        self.output = RobertaSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate
+class RobertaIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput
+class RobertaOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Roberta
+class RobertaLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = RobertaAttention(config)
+        self.is_decoder = config.is_decoder
+        self.add_cross_attention = config.add_cross_attention
+        if self.add_cross_attention:
+            if not self.is_decoder:
+                raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
+            self.crossattention = RobertaAttention(config, position_embedding_type="absolute")
+        self.intermediate = RobertaIntermediate(config)
+        self.output = RobertaOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+
+        # if decoder, the last output is tuple of self-attn cache
+        if self.is_decoder:
+            outputs = self_attention_outputs[1:-1]
+            present_key_value = self_attention_outputs[-1]
+        else:
+            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+
+        cross_attn_present_key_value = None
+        if self.is_decoder and encoder_hidden_states is not None:
+            if not hasattr(self, "crossattention"):
+                raise ValueError(
+                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
+                    " by setting `config.add_cross_attention=True`"
+                )
+
+            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
+            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                cross_attn_past_key_value,
+                output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights
+
+            # add cross-attn cache to positions 3,4 of present_key_value tuple
+            cross_attn_present_key_value = cross_attention_outputs[-1]
+            present_key_value = present_key_value + cross_attn_present_key_value
+
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+
+        # if decoder, return the attn key/values as the last output
+        if self.is_decoder:
+            outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Roberta
+class RobertaEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([RobertaLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = False,
+        output_hidden_states: Optional[bool] = False,
+        return_dict: Optional[bool] = True,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        next_decoder_cache = () if use_cache else None
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+                if self.config.add_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler
+class RobertaPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class RobertaPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = RobertaConfig
+    base_model_prefix = "roberta"
+    supports_gradient_checkpointing = True
+    _no_split_modules = []
+
+    # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, RobertaEncoder):
+            module.gradient_checkpointing = value
+
+    def update_keys_to_ignore(self, config, del_keys_to_ignore):
+        """Remove some keys from ignore list"""
+        if not config.tie_word_embeddings:
+            # must make a new list, or the class variable gets modified!
+            self._keys_to_ignore_on_save = [k for k in self._keys_to_ignore_on_save if k not in del_keys_to_ignore]
+            self._keys_to_ignore_on_load_missing = [
+                k for k in self._keys_to_ignore_on_load_missing if k not in del_keys_to_ignore
+            ]
+
+
+ROBERTA_START_DOCSTRING = r"""
+
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`RobertaConfig`]): Model configuration class with all the parameters of the
+            model. Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+ROBERTA_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `({0})`):
+            Indices of input sequence tokens in the vocabulary.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+        token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
+            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,1]`:
+
+            - 0 corresponds to a *sentence A* token,
+            - 1 corresponds to a *sentence B* token.
+            This parameter can only be used when the model is initialized with `type_vocab_size` parameter with value
+            >= 2. All the value in this tensor should be always < type_vocab_size.
+
+            [What are token type IDs?](../glossary#token-type-ids)
+        position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.max_position_embeddings - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
+    ROBERTA_START_DOCSTRING,
+)
+class RobertaModel(RobertaPreTrainedModel):
+    """
+
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in *Attention is
+    all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
+    Kaiser and Illia Polosukhin.
+
+    To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
+    to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
+    `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
+
+    .. _*Attention is all you need*: https://arxiv.org/abs/1706.03762
+
+    """
+
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    # Copied from transformers.models.bert.modeling_bert.BertModel.__init__ with Bert->Roberta
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = RobertaEmbeddings(config)
+        self.encoder = RobertaEncoder(config)
+
+        self.pooler = RobertaPooler(config) if add_pooling_layer else None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPoolingAndCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    # Copied from transformers.models.bert.modeling_bert.BertModel.forward
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
+        r"""
+        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.config.is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        batch_size, seq_length = input_shape
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # past_key_values_length
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        if attention_mask is None:
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+
+        if token_type_ids is None:
+            if hasattr(self.embeddings, "token_type_ids"):
+                buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
+                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
+                token_type_ids = buffered_token_type_ids_expanded
+            else:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values_length,
+        )
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+@add_start_docstrings(
+    """RoBERTa Model with a `language modeling` head on top for CLM fine-tuning.""", ROBERTA_START_DOCSTRING
+)
+class RobertaForCausalLM(RobertaPreTrainedModel):
+    _keys_to_ignore_on_save = [r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not config.is_decoder:
+            logger.warning("If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`")
+
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.lm_head = RobertaLMHead(config)
+
+        # The LM head weights require special treatment only when they are tied with the word embeddings
+        self.update_keys_to_ignore(config, ["lm_head.decoder.weight"])
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head.decoder = new_embeddings
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = 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.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        past_key_values: Tuple[Tuple[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
+        r"""
+        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
+            ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, RobertaForCausalLM, AutoConfig
+        >>> import torch
+
+        >>> tokenizer = AutoTokenizer.from_pretrained("roberta-base")
+        >>> config = AutoConfig.from_pretrained("roberta-base")
+        >>> config.is_decoder = True
+        >>> model = RobertaForCausalLM.from_pretrained("roberta-base", config=config)
+
+        >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+        >>> outputs = model(**inputs)
+
+        >>> prediction_logits = outputs.logits
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if labels is not None:
+            use_cache = False
+
+        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,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.lm_head(sequence_output)
+
+        lm_loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(prediction_scores.device)
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=lm_loss,
+            logits=prediction_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
+        input_shape = input_ids.shape
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_shape)
+
+        # cut decoder_input_ids if past is used
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values}
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past
+
+
+@add_start_docstrings("""RoBERTa Model with a `language modeling` head on top.""", ROBERTA_START_DOCSTRING)
+class RobertaForMaskedLM(RobertaPreTrainedModel):
+    _keys_to_ignore_on_save = [r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"lm_head.decoder.weight", r"lm_head.decoder.bias"]
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if config.is_decoder:
+            logger.warning(
+                "If you want to use `RobertaForMaskedLM` make sure `config.is_decoder=False` for "
+                "bi-directional self-attention."
+            )
+
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.lm_head = RobertaLMHead(config)
+
+        # The LM head weights require special treatment only when they are tied with the word embeddings
+        self.update_keys_to_ignore(config, ["lm_head.decoder.weight"])
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head.decoder = new_embeddings
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=MaskedLMOutput,
+        config_class=_CONFIG_FOR_DOC,
+        mask="<mask>",
+        expected_output="' Paris'",
+        expected_loss=0.1,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = 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.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
+            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+        kwargs (`Dict[str, any]`, optional, defaults to *{}*):
+            Used to hide legacy arguments that have been deprecated.
+        """
+        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,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = outputs[0]
+        prediction_scores = self.lm_head(sequence_output)
+
+        masked_lm_loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(prediction_scores.device)
+            loss_fct = CrossEntropyLoss()
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class RobertaLMHead(nn.Module):
+    """Roberta Head for masked language modeling."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        self.decoder.bias = self.bias
+
+    def forward(self, features, **kwargs):
+        x = self.dense(features)
+        x = gelu(x)
+        x = self.layer_norm(x)
+
+        # project back to size of vocabulary with bias
+        x = self.decoder(x)
+
+        return x
+
+    def _tie_weights(self):
+        # To tie those two weights if they get disconnected (on TPU or when the bias is resized)
+        # For accelerate compatibility and to not break backward compatibility
+        if self.decoder.bias.device.type == "meta":
+            self.decoder.bias = self.bias
+        else:
+            self.bias = self.decoder.bias
+
+
+@add_start_docstrings(
+    """
+    RoBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+    pooled output) e.g. for GLUE tasks.
+    """,
+    ROBERTA_START_DOCSTRING,
+)
+class RobertaForSequenceClassification(RobertaPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.classifier = RobertaClassificationHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint="cardiffnlp/twitter-roberta-base-emotion",
+        output_type=SequenceClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+        expected_output="'optimism'",
+        expected_loss=0.08,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = 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,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
+        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
+
+        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,
+        )
+        sequence_output = outputs[0]
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            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,
+        )
+
+
+@add_start_docstrings(
+    """
+    Roberta Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
+    softmax) e.g. for RocStories/SWAG tasks.
+    """,
+    ROBERTA_START_DOCSTRING,
+)
+class RobertaForMultipleChoice(RobertaPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.roberta = RobertaModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, 1)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=MultipleChoiceModelOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
+            num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
+            `input_ids` above)
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+
+        flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+        flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+        flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+        flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+        flat_inputs_embeds = (
+            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+            if inputs_embeds is not None
+            else None
+        )
+
+        outputs = self.roberta(
+            flat_input_ids,
+            position_ids=flat_position_ids,
+            token_type_ids=flat_token_type_ids,
+            attention_mask=flat_attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=flat_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)
+        reshaped_logits = logits.view(-1, num_choices)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(reshaped_logits.device)
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(reshaped_logits, labels)
+
+        if not return_dict:
+            output = (reshaped_logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return MultipleChoiceModelOutput(
+            loss=loss,
+            logits=reshaped_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    Roberta Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+    Named-Entity-Recognition (NER) tasks.
+    """,
+    ROBERTA_START_DOCSTRING,
+)
+class RobertaForTokenClassification(RobertaPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        classifier_dropout = (
+            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+        )
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint="Jean-Baptiste/roberta-large-ner-english",
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+        expected_output="['O', 'ORG', 'ORG', 'O', 'O', 'O', 'O', 'O', 'LOC', 'O', 'LOC', 'LOC']",
+        expected_loss=0.01,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = 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,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+        """
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class RobertaClassificationHead(nn.Module):
+    """Head for sentence-level classification tasks."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        classifier_dropout = (
+            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+        )
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
+
+    def forward(self, features, **kwargs):
+        x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+        x = self.out_proj(x)
+        return x
+
+
+@add_start_docstrings(
+    """
+    Roberta Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
+    layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
+    """,
+    ROBERTA_START_DOCSTRING,
+)
+class RobertaForQuestionAnswering(RobertaPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(ROBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        checkpoint="deepset/roberta-base-squad2",
+        output_type=QuestionAnsweringModelOutput,
+        config_class=_CONFIG_FOR_DOC,
+        expected_output="' puppet'",
+        expected_loss=0.86,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = 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,
+        start_positions: Optional[torch.LongTensor] = None,
+        end_positions: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
+        r"""
+        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
+            are not taken into account for computing the loss.
+        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
+            are not taken into account for computing the loss.
+        """
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[2:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
+    """
+    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+    are ignored. This is modified from fairseq's `utils.make_positions`.
+
+    Args:
+        x: torch.Tensor x:
+
+    Returns: torch.Tensor
+    """
+    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+    mask = input_ids.ne(padding_idx).int()
+    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
+    return incremental_indices.long() + padding_idx

soft_prompt/model/sequence_causallm.py

@@ -0,0 +1,1249 @@
+import torch
+from torch._C import NoopLogger
+import torch.nn
+import torch.nn.functional as F
+from torch import Tensor
+from typing import List, Optional, Tuple, Union
+from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
+
+from transformers.models.bert.modeling_bert import BertModel, BertPreTrainedModel, BertOnlyMLMHead
+from transformers.models.opt.modeling_opt import OPTModel, OPTPreTrainedModel
+from transformers.models.roberta.modeling_roberta import RobertaLMHead, RobertaModel, RobertaPreTrainedModel
+from transformers.models.llama.modeling_llama import LlamaPreTrainedModel, LlamaModel, CausalLMOutputWithPast
+from transformers.models.gpt2.modeling_gpt2 import GPT2Model, GPT2PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput, SequenceClassifierOutputWithPast, BaseModelOutput, Seq2SeqLMOutput
+from .prefix_encoder import PrefixEncoder
+from . import utils
+import hashlib
+
+
+def hash_nn(model):
+    md5 = hashlib.md5()  # ignore
+    for arg in model.parameters():
+        x = arg.data
+        if hasattr(x, "cpu"):
+            md5.update(x.cpu().numpy().data.tobytes())
+        elif hasattr(x, "numpy"):
+            md5.update(x.numpy().data.tobytes())
+        elif hasattr(x, "data"):
+            md5.update(x.data.tobytes())
+        else:
+            try:
+                md5.update(x.encode("utf-8"))
+            except:
+                md5.update(str(x).encode("utf-8"))
+    return md5.hexdigest()
+
+
+class OPTPrefixForMaskedLM(OPTPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = OPTModel(config)
+        self.lm_head = torch.nn.Linear(config.word_embed_proj_dim, config.vocab_size, bias=False)
+        self.dropout = torch.nn.Dropout(0.1)
+        for param in self.model.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)
+
+        base_param = 0
+        for name, param in self.model.named_parameters():
+            base_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - base_param
+        print('-> OPT_param:{:0.2f}M P-tuning-V2 param is {}'.format(base_param / 1000000, total_param))
+
+        self.embedding = self.get_input_embeddings()
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    def get_input_embeddings(self):
+        return self.model.decoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.decoder.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model.decoder = decoder
+
+    def get_decoder(self):
+        return self.model.decoder
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.model.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 use_grad(self, transformer, use_grad):
+        if use_grad:
+            for param in transformer.parameters():
+                param.requires_grad = True
+            transformer.train()
+        else:
+            for param in transformer.parameters():
+                param.requires_grad = False
+            transformer.eval()
+        for param in self.lm_head.parameters():
+            param.requires_grad = True
+        for param in self.prefix_encoder.parameters():
+            param.requires_grad = True
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        token_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,
+        use_base_grad=False,
+    ):
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
+                provide it.
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            head_mask (`torch.Tensor` of shape `(num_hidden_layers, num_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
+                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional
+                tensors are only required when the model is used as a decoder in a Sequence to Sequence model.
+
+                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
+                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
+                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
+                all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, OPTForCausalLM
+
+        >>> model = OPTForCausalLM.from_pretrained("facebook/opt-350m")
+        >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious. I'm just a little bit of a weirdo."
+        ```"""
+        utils.use_grad(self.model, use_base_grad)
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        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.model.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.model.decoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            past_key_values=past_key_values,
+        )
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(self.model.device)
+        cls_token = sequence_output[torch.arange(batch_size, device=self.model.device), sequence_lengths].contiguous()
+        attentions = self.lm_head(cls_token).view(-1, self.config.vocab_size).contiguous()
+
+        # compute loss
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past
+
+
+class OPTPromptForMaskedLM(OPTPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = OPTModel(config)
+        self.score = torch.nn.Linear(config.word_embed_proj_dim, self.num_labels, bias=False)
+        self.lm_head = torch.nn.Linear(config.word_embed_proj_dim, config.vocab_size, bias=False)
+        self.dropout = torch.nn.Dropout(0.1)
+        for param in self.model.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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        model_param = 0
+        for name, param in self.model.named_parameters():
+            model_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - model_param
+        print('-> OPT_param:{:0.2f}M P-tuning-V2 param is {}'.format(model_param / 1000000, total_param))
+
+        self.embedding = self.model.decoder.embed_tokens
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    def get_input_embeddings(self):
+        return self.model.decoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.decoder.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model.decoder = decoder
+
+    def get_decoder(self):
+        return self.model.decoder
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.model.device)
+        prompts = self.prefix_encoder(prefix_tokens)
+        return prompts
+
+    def use_grad(self, transformer, use_grad):
+        if use_grad:
+            for param in transformer.parameters():
+                param.requires_grad = True
+            transformer.train()
+        else:
+            for param in transformer.parameters():
+                param.requires_grad = False
+            transformer.eval()
+        for param in self.lm_head.parameters():
+            param.requires_grad = True
+        for param in self.prefix_encoder.parameters():
+            param.requires_grad = True
+
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            head_mask: Optional[torch.Tensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            token_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,
+            use_base_grad=False,
+    ):
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
+                provide it.
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            head_mask (`torch.Tensor` of shape `(num_hidden_layers, num_attention_heads)`, *optional*):
+                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
+                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional
+                tensors are only required when the model is used as a decoder in a Sequence to Sequence model.
+
+                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
+                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
+                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
+                all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, OPTForCausalLM
+
+        >>> model = OPTForCausalLM.from_pretrained("facebook/opt-350m")
+        >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious. I'm just a little bit of a weirdo."
+        ```"""
+        utils.use_grad(self.model, use_base_grad)
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        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.model.decoder.embed_tokens(input_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.model.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model.decoder(
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = outputs[0]
+        sequence_output = sequence_output[:, self.pre_seq_len:, :]
+        sequence_output = self.dropout(sequence_output)
+        sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(self.model.device)
+        cls_token = sequence_output[torch.arange(batch_size, device=self.model.device), sequence_lengths].contiguous()
+        attentions = self.lm_head(cls_token).view(-1, self.config.vocab_size).contiguous()
+
+        # compute loss
+        loss = None
+        if token_labels is not None:
+            loss = utils.get_loss(attentions, token_labels).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for idx, y in enumerate(self.clean_labels):
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+        #loss = torch.nn.functional.nll_loss(logits, labels)
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+    def prepare_inputs_for_generation(
+            self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past
+
+
+class LlamaPrefixForMaskedLM(LlamaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.dropout = torch.nn.Dropout(0.1)
+        for param in self.model.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)
+
+        base_param = 0
+        for name, param in self.model.named_parameters():
+            base_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - base_param
+        print('-> LLama_param:{:0.2f}M P-tuning-V2 param:{:0.2f}M'.format(base_param / 1000000, total_param/ 1000000))
+
+        self.embedding = self.model.embed_tokens
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    def get_prompt(self, batch_size):
+        device = next(self.prefix_encoder.parameters()).device
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(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_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def use_grad(self, base_model, use_grad):
+        if use_grad:
+            for param in base_model.parameters():
+                param.requires_grad = True
+            base_model.train()
+        else:
+            for param in base_model.parameters():
+                param.requires_grad = False
+            base_model.eval()
+        for param in self.prefix_encoder.parameters():
+            param.requires_grad = True
+        for param in self.lm_head.parameters():
+            param.requires_grad = True
+
+    def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,
+            token_type_ids=None,
+            position_ids=None,
+            inputs_embeds=None,
+            labels=None,
+            token_labels=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+            use_base_grad=False,
+    ):
+        utils.use_grad(self.model, use_base_grad)
+        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(attention_mask.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            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,
+        )
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        #sequence_output = torch.clamp(sequence_output, min=-1, max=1)
+        #cls_token = sequence_output[:, :1]
+        sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(sequence_output.device)
+        cls_token = sequence_output[torch.arange(batch_size, device=sequence_output.device), sequence_lengths].contiguous()
+        attentions = self.lm_head(cls_token).view(-1, self.config.vocab_size).contiguous()
+
+        # compute loss
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels.to(attentions.device)).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+
+class LlamaPromptForMaskedLM(LlamaPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.dropout = torch.nn.Dropout(0.1)
+        for param in self.model.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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        model_param = 0
+        for name, param in self.model.named_parameters():
+            model_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - model_param
+        print('-> Llama_param:{:0.2f}M P-tuning-V2 param is {:0.2f}M'.format(model_param / 1000000, total_param / 1000000))
+
+        self.pad_token_id = 2
+        self.embedding = self.model.embed_tokens
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    def get_prompt(self, batch_size):
+        device = next(self.prefix_encoder.parameters()).device
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
+        prompts = self.prefix_encoder(prefix_tokens)
+        return prompts
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def use_grad(self, base_model, use_grad):
+        if use_grad:
+            for param in base_model.parameters():
+                param.requires_grad = True
+            for param in self.lm_head.parameters():
+                param.requires_grad = True
+            base_model.train()
+        else:
+            for param in base_model.parameters():
+                param.requires_grad = False
+            for param in self.lm_head.parameters():
+                param.requires_grad = False
+            base_model.eval()
+        for param in self.prefix_encoder.parameters():
+            param.requires_grad = True
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] =None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        token_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,
+        use_base_grad: Optional[bool] = False,
+    ):
+        self.use_grad(self.model, use_base_grad)
+        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.model.embed_tokens(input_ids)
+        prompts = self.get_prompt(batch_size=batch_size)
+        inputs_embeds = torch.cat((prompts, raw_embedding.to(prompts.device)), dim=1)
+        prefix_attention_mask = torch.ones(batch_size, self.pre_seq_len).to(attention_mask.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = outputs[0]
+        sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        #cls_token = sequence_output[:, 0]
+        #cls_token = self.dropout(cls_token)
+        sequence_lengths = (torch.ne(input_ids, self.pad_token_id).sum(-1) - 1).to(sequence_output.device)
+        cls_token = sequence_output[torch.arange(batch_size, device=sequence_output.device), sequence_lengths].contiguous()
+        attentions = self.lm_head(cls_token).view(-1, self.config.vocab_size).contiguous().float()
+
+        # compute loss
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels.to(attentions.device)).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+
+class BertPrefixForMaskedLM(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        for param in self.bert.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)
+
+        base_param = 0
+        for name, param in self.bert.named_parameters():
+            base_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - base_param
+        print('-> bert_param:{:0.2f}M P-tuning-V2 param is {}'.format(base_param / 1000000, total_param))
+
+        # bert.embeddings.word_embeddings
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    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,
+        token_labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        use_base_grad=False,
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        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.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,
+        )
+        sequence_output = outputs[0]
+        cls_token = sequence_output[:, 0]
+        cls_token = self.dropout(cls_token)
+        attentions = self.cls(cls_token).view(-1, self.config.vocab_size)
+
+
+        # compute loss
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+
+class BertPromptForMaskedLM(BertPreTrainedModel):
+    def __init__(self, config):
+        _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        for param in self.bert.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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        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('-> bert_param:{:0.2f}M P-tuning-V2 param is {}'.format(bert_param / 1000000, total_param))
+
+        # bert.embeddings.word_embeddings
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    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,
+        token_labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        use_base_grad=False,
+    ):
+        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).
+        """
+        utils.use_grad(self.bert, use_base_grad)
+        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.bert.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,
+        )
+        sequence_output = outputs[0]
+        sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        cls_token = sequence_output[:, 0]
+        cls_token = self.dropout(cls_token)
+        attentions = self.cls(cls_token).view(-1, self.config.vocab_size)
+
+        # compute loss
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+
+class RobertaPrefixForMaskedLM(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.lm_head = RobertaLMHead(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+
+        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
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        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,
+            token_labels=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+            use_base_grad=False,
+    ):
+        utils.use_grad(self.roberta, use_base_grad)
+        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)
+
+        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,
+        )
+        sequence_output = outputs[0]
+        cls_token = sequence_output[:, 0]
+        cls_token = self.dropout(cls_token)
+        attentions = self.lm_head(cls_token).view(-1, self.config.vocab_size)
+
+        # compute loss
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )
+
+
+class RobertaPromptForMaskedLM(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.lm_head = RobertaLMHead(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        self.embeddings = self.roberta.embeddings
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+        self.clean_labels = torch.tensor(config.clean_labels).long()
+
+    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,
+            token_labels=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+            use_base_grad=False
+    ):
+        utils.use_grad(self.roberta, use_base_grad)
+        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.roberta.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.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,
+        )
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        sequence_output = sequence_output[:, self.pre_seq_len:, :].contiguous()
+        cls_token = sequence_output[:, 0]
+        attentions = self.lm_head(cls_token).view(-1, self.config.vocab_size)
+
+        masked_lm_loss = None
+        if token_labels is not None:
+            masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+        else:
+            if labels is not None:
+                token_labels = torch.stack([self.clean_labels[labels[i]] for i in range(len(labels))]).to(labels.device)
+                masked_lm_loss = utils.get_loss(attentions, token_labels).sum()
+
+        # convert to binary classifier
+        probs = []
+        for y in self.clean_labels:
+            probs.append(attentions[:, y.to(attentions.device)].max(dim=1)[0])
+        logits = torch.stack(probs).T
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+        return SequenceClassifierOutput(
+            loss=masked_lm_loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=attentions
+        )

soft_prompt/model/sequence_classification.py

@@ -0,0 +1,997 @@
+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.modeling_outputs import SequenceClassifierOutput, SequenceClassifierOutputWithPast, BaseModelOutput, Seq2SeqLMOutput
+from transformers import GPT2Model, GPT2PreTrainedModel, GPTNeoModel
+
+from model.prefix_encoder import PrefixEncoder
+from model.deberta import DebertaModel, DebertaPreTrainedModel, ContextPooler, StableDropout
+from model import utils
+import copy
+
+class BertForSequenceClassification(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)
+
+        self.init_weights()
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    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,
+        use_base_grad=False,
+    ):
+        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).
+        """
+        utils.use_grad(self.bert, use_base_grad)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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)
+            elif self.config.problem_type == "em":
+                predict_logp = F.log_softmax(pooled_output, dim=-1)
+                target_logp = predict_logp.gather(-1, labels)
+                target_logp = target_logp - 1e32 * labels.eq(0)  # Apply mask
+                loss = -torch.logsumexp(target_logp, dim=-1)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        loss.backward()
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=pooled_output,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+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
+        
+        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('-> bert_param:{:0.2f}M P-tuning-V2 param is {}'.format(bert_param / 1000000, total_param))
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    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,
+        use_base_grad=False,
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        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.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]
+        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,
+        )
+
+
+class BertPromptForSequenceClassification(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
+        
+        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)
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+    
+    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,
+        use_base_grad=False,
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        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,
+        )
+
+
+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.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
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        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,
+        use_base_grad=False,
+    ):
+        utils.use_grad(self.roberta, use_base_grad)
+        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)
+        
+        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:
+            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,
+        )
+
+
+class RobertaPromptForSequenceClassification(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)
+
+        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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+    
+    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,
+        use_base_grad=False
+    ):
+        utils.use_grad(self.roberta, use_base_grad)
+        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.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,
+        )
+
+
+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
+        
+        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
+
+        self.embedding = utils.get_embeddings(self, config)
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+    
+    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,
+        use_base_grad=False
+    ):
+        utils.use_grad(self.bert, use_base_grad)
+        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,
+            )
+
+
+class GPT2PromptForSequenceClassification(GPT2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.gpt2 = GPT2Model(config)
+        self.dropout = StableDropout(config.embd_pdrop)
+        self.classifier = torch.nn.Linear(config.n_embd, self.num_labels)
+
+        for param in self.gpt2.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 = torch.nn.Embedding(self.pre_seq_len, config.hidden_size)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+        gpt2_param = 0
+        for name, param in self.gpt2.named_parameters():
+            gpt2_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - gpt2_param
+        print('-> total param is {}'.format(total_param))  # 9860105
+
+        self.embedding = self.gpt2.wte
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.gpt2.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,
+            use_base_grad=False
+    ):
+        utils.use_grad(self.gpt2, use_base_grad)
+        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.embedding(input_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.gpt2.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        transformer_outputs = self.gpt2(
+            # 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,
+        )
+
+        hidden_states = transformer_outputs[0]
+        logits = self.classifier(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).to(logits.device)
+            else:
+                sequence_lengths = -1
+
+        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,
+        )
+
+class GPT2PrefixForSequenceClassification(GPT2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.gpt2 = GPT2Model(config)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.n_embd, self.num_labels)
+
+        for param in self.gpt2.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)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+        gpt2_param = 0
+        for name, param in self.gpt2.named_parameters():
+            gpt2_param += param.numel()
+        all_param = 0
+        for name, param in self.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - gpt2_param
+        print('-> gpt2_param:{:0.2f}M P-tuning-V2 param is {}'.format(gpt2_param/1000000, total_param))
+
+        self.embedding = self.gpt2.wte
+        self.embeddings_gradient = utils.GradientStorage(self.embedding)
+
+    def get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.gpt2.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        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,
+        use_base_grad=False,
+        use_cache=None
+    ):
+        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).
+        """
+        utils.use_grad(self.gpt2, use_base_grad)
+        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.gpt2.device)
+        attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
+
+        transformer_outputs = self.gpt2(
+            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.classifier(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).to(logits.device)
+            else:
+                sequence_lengths = -1
+
+        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,
+        )
+
+
+if __name__ == "__main__":
+    from transformers import AutoConfig
+    config = AutoConfig.from_pretrained("gpt2-large")
+    config.hidden_dropout_prob = 0.1
+    config.pre_seq_len = 128
+    config.prefix_projection = True
+    config.num_labels = 2
+    config.prefix_hidden_size = 1024
+    model = GPT2PrefixForSequenceClassification(config)
+
+    for name, param in model.named_parameters():
+        print(name, param.shape)
+
+

soft_prompt/model/token_classification.py

@@ -0,0 +1,539 @@
+import torch
+import torch.nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import CrossEntropyLoss
+
+from transformers import BertModel, BertPreTrainedModel
+from transformers import RobertaModel, RobertaPreTrainedModel
+from transformers.modeling_outputs import TokenClassifierOutput
+
+from model.prefix_encoder import PrefixEncoder
+from model.deberta import DebertaModel, DebertaPreTrainedModel
+from model.debertaV2 import DebertaV2Model, DebertaV2PreTrainedModel
+
+class BertForTokenClassification(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        only_cls_head = True # False in SRL
+        if only_cls_head:
+            for param in self.bert.parameters():
+                param.requires_grad = False
+
+        self.init_weights()
+
+        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))
+
+    
+    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, sequence_length)`, `optional`):
+            Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels -
+            1]``.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class BertPrefixForTokenClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+
+        from_pretrained = False
+        if from_pretrained:
+            self.classifier.load_state_dict(torch.load('model/checkpoint.pkl'))
+
+        for param in self.bert.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.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 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
+
+        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)
+
+        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,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+        attention_mask = attention_mask[:,self.pre_seq_len:].contiguous()
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+
+
+class RobertaPrefixForTokenClassification(RobertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        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 get_prompt(self, batch_size):
+        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(self.roberta.device)
+        past_key_values = self.prefix_encoder(prefix_tokens)
+        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.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,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+        attention_mask = attention_mask[:,self.pre_seq_len:].contiguous()
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class DebertaPrefixForTokenClassification(DebertaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.deberta = DebertaModel(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.init_weights()
+
+        for param in self.deberta.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)
+
+        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 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,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+        attention_mask = attention_mask[:,self.pre_seq_len:].contiguous()
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class DebertaV2PrefixForTokenClassification(DebertaV2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.deberta = DebertaV2Model(config)
+        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
+        self.init_weights()
+
+        for param in self.deberta.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)
+
+        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 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)
+        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,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+        attention_mask = attention_mask[:,self.pre_seq_len:].contiguous()
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

soft_prompt/model/utils.py

@@ -0,0 +1,399 @@
+from enum import Enum
+import torch
+from .token_classification import (
+    BertPrefixForTokenClassification,
+    RobertaPrefixForTokenClassification,
+    DebertaPrefixForTokenClassification,
+    DebertaV2PrefixForTokenClassification
+)
+
+from .sequence_classification import (
+    BertPrefixForSequenceClassification,
+    BertPromptForSequenceClassification,
+    RobertaPrefixForSequenceClassification,
+    RobertaPromptForSequenceClassification,
+    DebertaPrefixForSequenceClassification,
+    GPT2PrefixForSequenceClassification,
+    GPT2PromptForSequenceClassification
+)
+
+from .question_answering import (
+    BertPrefixForQuestionAnswering,
+    RobertaPrefixModelForQuestionAnswering,
+    DebertaPrefixModelForQuestionAnswering
+)
+
+from .multiple_choice import (
+    BertPrefixForMultipleChoice,
+    RobertaPrefixForMultipleChoice,
+    DebertaPrefixForMultipleChoice,
+    BertPromptForMultipleChoice,
+    RobertaPromptForMultipleChoice
+)
+
+from .sequence_causallm import (
+    BertPromptForMaskedLM,
+    BertPrefixForMaskedLM,
+    RobertaPromptForMaskedLM,
+    RobertaPrefixForMaskedLM,
+    LlamaPromptForMaskedLM,
+    LlamaPrefixForMaskedLM,
+    OPTPrefixForMaskedLM,
+    OPTPromptForMaskedLM
+)
+
+from transformers import (
+    AutoConfig,
+    AutoModelForTokenClassification,
+    AutoModelForSequenceClassification,
+    AutoModelForQuestionAnswering,
+    AutoModelForMultipleChoice
+)
+import torch.nn.functional as F
+
+
+def get_loss(predict_logits, labels_ids):
+    labels_ids = labels_ids.to(predict_logits.device)
+    predict_logp = F.log_softmax(predict_logits, dim=-1)
+    target_logp = predict_logp.gather(-1, labels_ids)
+    target_logp = target_logp - 1e32 * labels_ids.eq(0)  # Apply mask
+    target_logp = torch.logsumexp(target_logp, dim=-1)
+    return -target_logp
+
+
+def use_grad(base_model, use_grad):
+    if use_grad:
+        for param in base_model.parameters():
+            param.requires_grad = True
+        base_model.train()
+    else:
+        for param in base_model.parameters():
+            param.requires_grad = False
+        base_model.eval()
+
+
+def get_embeddings(model, config):
+    """Returns the wordpiece embedding module."""
+    base_model = getattr(model, config.model_type)
+    embeddings = base_model.embeddings.word_embeddings
+    return embeddings
+
+
+class GradientStorage:
+    """
+    This object stores the intermediate gradients of the output a the given PyTorch module, which
+    otherwise might not be retained.
+    """
+    def __init__(self, module):
+        self._stored_gradient = None
+        module.register_backward_hook(self.hook)
+
+    def hook(self, module, grad_in, grad_out):
+        assert grad_out is not None
+        self._stored_gradient = grad_out[0]
+        
+    def reset(self):
+        self._stored_gradient = None
+
+    def get(self):
+        return self._stored_gradient
+
+
+class TaskType(Enum):
+    TOKEN_CLASSIFICATION = 1,
+    SEQUENCE_CLASSIFICATION = 2,
+    QUESTION_ANSWERING = 3,
+    MULTIPLE_CHOICE = 4
+
+PREFIX_MODELS = {
+    "bert": {
+        TaskType.TOKEN_CLASSIFICATION: BertPrefixForTokenClassification,
+        TaskType.SEQUENCE_CLASSIFICATION: BertPrefixForMaskedLM, #BertPrefixForSequenceClassification,
+        TaskType.QUESTION_ANSWERING: BertPrefixForQuestionAnswering,
+        TaskType.MULTIPLE_CHOICE: BertPrefixForMultipleChoice
+    },
+    "roberta": {
+        TaskType.TOKEN_CLASSIFICATION: RobertaPrefixForTokenClassification,
+        TaskType.SEQUENCE_CLASSIFICATION: RobertaPrefixForMaskedLM, #RobertaPrefixForSequenceClassification,
+        TaskType.QUESTION_ANSWERING: RobertaPrefixModelForQuestionAnswering,
+        TaskType.MULTIPLE_CHOICE: RobertaPrefixForMultipleChoice,
+    },
+    "deberta": {
+        TaskType.TOKEN_CLASSIFICATION: DebertaPrefixForTokenClassification,
+        TaskType.SEQUENCE_CLASSIFICATION: DebertaPrefixForSequenceClassification,
+        TaskType.QUESTION_ANSWERING: DebertaPrefixModelForQuestionAnswering,
+        TaskType.MULTIPLE_CHOICE: DebertaPrefixForMultipleChoice,
+    },
+    "deberta-v2": {
+        TaskType.TOKEN_CLASSIFICATION: DebertaV2PrefixForTokenClassification,
+        TaskType.SEQUENCE_CLASSIFICATION: None,
+        TaskType.QUESTION_ANSWERING: None,
+        TaskType.MULTIPLE_CHOICE: None,
+    },
+    "gpt2": {
+        TaskType.TOKEN_CLASSIFICATION: None,
+        TaskType.SEQUENCE_CLASSIFICATION: GPT2PrefixForSequenceClassification,
+        TaskType.QUESTION_ANSWERING: None,
+        TaskType.MULTIPLE_CHOICE: None,
+    },
+    "llama": {
+        TaskType.TOKEN_CLASSIFICATION: None,
+        TaskType.SEQUENCE_CLASSIFICATION: LlamaPrefixForMaskedLM,
+        TaskType.QUESTION_ANSWERING: None,
+        TaskType.MULTIPLE_CHOICE: None,
+    },
+    "opt": {
+        TaskType.TOKEN_CLASSIFICATION: None,
+        TaskType.SEQUENCE_CLASSIFICATION: OPTPrefixForMaskedLM,
+        TaskType.QUESTION_ANSWERING: None,
+        TaskType.MULTIPLE_CHOICE: None,
+    }
+}
+
+PROMPT_MODELS = {
+    "bert": {
+        TaskType.SEQUENCE_CLASSIFICATION: BertPromptForMaskedLM, #BertPromptForSequenceClassification,
+        TaskType.MULTIPLE_CHOICE: BertPromptForMultipleChoice
+    },
+    "roberta": {
+        TaskType.SEQUENCE_CLASSIFICATION: RobertaPromptForMaskedLM, #RobertaPromptForSequenceClassification,
+        TaskType.MULTIPLE_CHOICE: RobertaPromptForMultipleChoice
+    },
+    "gpt2": {
+        TaskType.SEQUENCE_CLASSIFICATION: GPT2PromptForSequenceClassification,
+        TaskType.MULTIPLE_CHOICE: None
+    },
+    "llama": {
+        TaskType.TOKEN_CLASSIFICATION: None,
+        TaskType.SEQUENCE_CLASSIFICATION: LlamaPromptForMaskedLM,
+        TaskType.QUESTION_ANSWERING: None,
+        TaskType.MULTIPLE_CHOICE: None,
+    },
+    "opt": {
+        TaskType.TOKEN_CLASSIFICATION: None,
+        TaskType.SEQUENCE_CLASSIFICATION: OPTPromptForMaskedLM,
+        TaskType.QUESTION_ANSWERING: None,
+        TaskType.MULTIPLE_CHOICE: None,
+    }
+}
+
+AUTO_MODELS = {
+    TaskType.TOKEN_CLASSIFICATION: AutoModelForTokenClassification,
+    TaskType.SEQUENCE_CLASSIFICATION: AutoModelForSequenceClassification,
+    TaskType.QUESTION_ANSWERING: AutoModelForQuestionAnswering,
+    TaskType.MULTIPLE_CHOICE: AutoModelForMultipleChoice,
+}
+
+def get_model(model_args, task_type: TaskType, config: AutoConfig, fix_bert: bool = False, tokenizer=None):
+    model_name_or_path = f'openlm-research/{model_args.model_name_or_path}' if "llama" in model_args.model_name_or_path else model_args.model_name_or_path
+
+    if model_args.prefix:
+        config.hidden_dropout_prob = model_args.hidden_dropout_prob
+        config.pre_seq_len = model_args.pre_seq_len
+        config.prefix_projection = model_args.prefix_projection
+        config.prefix_hidden_size = model_args.prefix_hidden_size
+        model_class = PREFIX_MODELS[config.model_type][task_type]
+        if "opt" in model_args.model_name_or_path:
+            model_name_or_path = f'facebook/{model_args.model_name_or_path}'
+            model = model_class.from_pretrained(
+                model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+        elif "llama" in model_args.model_name_or_path:
+            model_name_or_path = f'openlm-research/{model_args.model_name_or_path}'
+            model = model_class.from_pretrained(
+                model_name_or_path,
+                config=config,
+                trust_remote_code=True,
+                torch_dtype=torch.float32,
+                device_map='auto',
+            )
+        else:
+            model = model_class.from_pretrained(
+                model_name_or_path,
+                config=config,
+                trust_remote_code=True,
+                revision=model_args.model_revision
+            )
+    elif model_args.prompt:
+        config.pre_seq_len = model_args.pre_seq_len
+        model_class = PROMPT_MODELS[config.model_type][task_type]
+        if "opt" in model_args.model_name_or_path:
+            model_name_or_path = f'facebook/opt-1.3b'
+            model = model_class.from_pretrained(
+                model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+        elif "llama" in model_args.model_name_or_path:
+            model_name_or_path = f'openlm-research/{model_args.model_name_or_path}'
+            model = model_class.from_pretrained(
+                model_name_or_path,
+                config=config,
+                trust_remote_code=True,
+                torch_dtype=torch.float32,
+                device_map='auto',
+            )
+        else:
+            model = model_class.from_pretrained(
+                model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+    else:
+        model_class = AUTO_MODELS[task_type]
+        model = model_class.from_pretrained(
+            model_name_or_path,
+            config=config,
+            revision=model_args.model_revision,
+        )
+        base_param = 0
+        if fix_bert:
+            if config.model_type == "bert":
+                for param in model.bert.parameters():
+                    param.requires_grad = False
+                for _, param in model.bert.named_parameters():
+                    base_param += param.numel()
+            elif config.model_type == "roberta":
+                for param in model.roberta.parameters():
+                    param.requires_grad = False
+                for _, param in model.roberta.named_parameters():
+                    base_param += param.numel()
+            elif config.model_type == "deberta":
+                for param in model.deberta.parameters():
+                    param.requires_grad = False
+                for _, param in model.deberta.named_parameters():
+                    base_param += param.numel()
+            elif config.model_type == "gpt2":
+                for param in model.gpt2.parameters():
+                    param.requires_grad = False
+                for _, param in model.gpt2.named_parameters():
+                    base_param += param.numel()
+        all_param = 0
+        for _, param in model.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - base_param
+        print('***** Backborn param:{:0.3f}M, P-Tuning-V2 param is {} *****'.format(all_param, total_param))
+
+    return model
+
+
+def get_model_deprecated(model_args, task_type: TaskType, config: AutoConfig, fix_bert: bool = False):
+    if model_args.prefix:
+        config.hidden_dropout_prob = model_args.hidden_dropout_prob
+        config.pre_seq_len = model_args.pre_seq_len
+        config.prefix_projection = model_args.prefix_projection
+        config.prefix_hidden_size = model_args.prefix_hidden_size
+
+        if task_type == TaskType.TOKEN_CLASSIFICATION:
+            from model.token_classification import BertPrefixModel, RobertaPrefixModel, DebertaPrefixModel, DebertaV2PrefixModel
+        elif task_type == TaskType.SEQUENCE_CLASSIFICATION:
+            from model.sequence_classification import BertPrefixModel, RobertaPrefixModel, DebertaPrefixModel, DebertaV2PrefixModel
+        elif task_type == TaskType.QUESTION_ANSWERING:
+            from model.question_answering import BertPrefixModel, RobertaPrefixModel, DebertaPrefixModel, DebertaV2PrefixModel
+        elif task_type == TaskType.MULTIPLE_CHOICE:
+            from model.multiple_choice import BertPrefixModel
+
+        if config.model_type == "bert":
+            model = BertPrefixModel.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        elif config.model_type == "roberta":
+            model = RobertaPrefixModel.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        elif config.model_type == "deberta":
+            model = DebertaPrefixModel.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        elif config.model_type == "deberta-v2":
+            model = DebertaV2PrefixModel.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        else:
+            raise NotImplementedError
+
+
+    elif model_args.prompt:
+        config.pre_seq_len = model_args.pre_seq_len
+
+        from model.sequence_classification import BertPromptModel, RobertaPromptModel
+        if config.model_type == "bert":
+            model = BertPromptModel.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        elif config.model_type == "roberta":
+            model = RobertaPromptModel.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        else:
+            raise NotImplementedError
+            
+
+    else:
+        if task_type == TaskType.TOKEN_CLASSIFICATION:
+            model = AutoModelForTokenClassification.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+            
+        elif task_type == TaskType.SEQUENCE_CLASSIFICATION:
+            model = AutoModelForSequenceClassification.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+
+        elif task_type == TaskType.QUESTION_ANSWERING:
+            model = AutoModelForQuestionAnswering.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+        elif task_type == TaskType.MULTIPLE_CHOICE:
+            model = AutoModelForMultipleChoice.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                revision=model_args.model_revision,
+            )
+    
+        bert_param = 0
+        if fix_bert:
+            if config.model_type == "bert":
+                for param in model.bert.parameters():
+                    param.requires_grad = False
+                for _, param in model.bert.named_parameters():
+                    bert_param += param.numel()
+            elif config.model_type == "roberta":
+                for param in model.roberta.parameters():
+                    param.requires_grad = False
+                for _, param in model.roberta.named_parameters():
+                    bert_param += param.numel()
+            elif config.model_type == "deberta":
+                for param in model.deberta.parameters():
+                    param.requires_grad = False
+                for _, param in model.deberta.named_parameters():
+                    bert_param += param.numel()
+        all_param = 0
+        for _, param in model.named_parameters():
+            all_param += param.numel()
+        total_param = all_param - bert_param
+        print('***** total param is {} *****'.format(total_param))
+    return model

soft_prompt/run.py

@@ -0,0 +1,177 @@
+import logging
+import os
+import os.path as osp
+import sys
+import numpy as np
+from typing import Dict
+
+import datasets
+import transformers
+from transformers import set_seed, Trainer
+from transformers.trainer_utils import get_last_checkpoint
+
+from arguments import get_args
+
+from tasks.utils import *
+
+os.environ["WANDB_DISABLED"] = "true"
+
+logger = logging.getLogger(__name__)
+
+def train(trainer, resume_from_checkpoint=None, last_checkpoint=None):
+    checkpoint = None
+    if resume_from_checkpoint is not None:
+        checkpoint = resume_from_checkpoint
+    elif last_checkpoint is not None:
+        checkpoint = last_checkpoint
+    train_result = trainer.train(resume_from_checkpoint=checkpoint)
+    # trainer.save_model()
+
+    metrics = train_result.metrics
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+    trainer.save_state()
+    trainer.log_best_metrics()
+
+
+def evaluate(args, trainer, checkpoint=None):
+    logger.info("*** Evaluate ***")
+
+    if checkpoint is not None:
+        trainer._load_from_checkpoint(resume_from_checkpoint=checkpoint)
+        trainer._resume_watermark()
+
+    metrics = trainer.evaluate(ignore_keys=["hidden_states", "attentions"])
+    score, asr = 0., 0.
+    if training_args.watermark != "clean":
+        score, asr = trainer.evaluate_watermark()
+    metrics["wmk_asr"] = asr
+    metrics["wmk_score"] = score
+    trainer.evaluate_clean()
+    torch.save(trainer.eval_memory, f"{args.output_dir}/exp11_attentions.pth")
+
+    trainer.log_metrics("eval", metrics)
+    path = osp.join(args.output_dir, "exp11_acc_asr.pth")
+    torch.save(metrics, path)
+
+
+def predict(trainer, predict_dataset=None):
+    if predict_dataset is None:
+        logger.info("No dataset is available for testing")
+
+    elif isinstance(predict_dataset, dict):
+        
+        for dataset_name, d in predict_dataset.items():
+            logger.info("*** Predict: %s ***" % dataset_name)
+            predictions, labels, metrics = trainer.predict(d, metric_key_prefix="predict")
+            predictions = np.argmax(predictions, axis=2)
+
+            trainer.log_metrics("predict", metrics)
+            trainer.save_metrics("predict", metrics)
+
+    else:
+        logger.info("*** Predict ***")
+        predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")
+        predictions = np.argmax(predictions, axis=2)
+
+        trainer.log_metrics("predict", metrics)
+        trainer.save_metrics("predict", metrics)
+
+if __name__ == '__main__':
+    args = get_args()
+    p_type = "prefix" if args[0].prefix else "prompt"
+    output_root = osp.join("checkpoints", f"{args[1].task_name}_{args[1].dataset_name}_{args[0].model_name_or_path}_{args[2].watermark}_{p_type}")
+    output_dir = osp.join(output_root, f"t{args[2].trigger_num}_p{args[2].poison_rate:0.2f}")
+    for path in [output_root, output_dir]:
+        if not osp.exists(path):
+            try:
+                os.makedirs(path)
+            except:
+                pass
+
+    args[0].output_dir = output_dir
+    args[1].output_dir = output_dir
+    args[2].output_dir = output_dir
+    args[3].output_dir = output_dir
+    torch.save(args, osp.join(output_dir, "args.pt"))
+    model_args, data_args, training_args, _ = args
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+
+    log_level = training_args.get_process_log_level()
+    logger.setLevel(log_level)
+    datasets.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.set_verbosity(log_level)
+    transformers.utils.logging.enable_default_handler()
+    transformers.utils.logging.enable_explicit_format()
+
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+
+
+    if not os.path.isdir("checkpoints") or not os.path.exists("checkpoints"):
+        os.mkdir("checkpoints")
+
+    if data_args.task_name.lower() == "superglue":
+        assert data_args.dataset_name.lower() in SUPERGLUE_DATASETS
+        from tasks.superglue.get_trainer import get_trainer
+
+    elif data_args.task_name.lower() == "glue":
+        assert data_args.dataset_name.lower() in GLUE_DATASETS
+        from tasks.glue.get_trainer import get_trainer
+
+    elif data_args.task_name.lower() == "ner":
+        assert data_args.dataset_name.lower() in NER_DATASETS
+        from tasks.ner.get_trainer import get_trainer
+
+    elif data_args.task_name.lower() == "srl":
+        assert data_args.dataset_name.lower() in SRL_DATASETS
+        from tasks.srl.get_trainer import get_trainer
+    
+    elif data_args.task_name.lower() == "qa":
+        assert data_args.dataset_name.lower() in QA_DATASETS
+        from tasks.qa.get_trainer import get_trainer
+    elif data_args.task_name.lower() == "ag_news":
+        from tasks.ag_news.get_trainer import get_trainer
+    elif data_args.task_name.lower() == "imdb":
+        from tasks.imdb.get_trainer import get_trainer
+    else:
+        raise NotImplementedError('Task {} is not implemented. Please choose a task from: {}'.format(data_args.task_name, ", ".join(TASKS)))
+
+    set_seed(training_args.seed)
+    trainer, predict_dataset = get_trainer(args)
+
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    if training_args.do_train:
+        train(trainer, training_args.resume_from_checkpoint, last_checkpoint)
+    
+    if training_args.do_eval:
+        if last_checkpoint is None:
+            last_checkpoint = osp.join(training_args.output_dir, "checkpoint")
+        print(f"-> last_checkpoint:{last_checkpoint}")
+        evaluate(training_args, trainer, checkpoint=last_checkpoint)
+
+    # if training_args.do_predict:
+    #     predict(trainer, predict_dataset)
+
+   

soft_prompt/tasks/ag_news/dataset.py

@@ -0,0 +1,159 @@
+import torch, math
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    EvalPrediction,
+    default_data_collator,
+)
+import re
+import numpy as np
+import logging, re
+from datasets.formatting.formatting import LazyRow, LazyBatch
+
+
+task_to_keys = {
+    "ag_news": ("text", None)
+}
+
+logger = logging.getLogger(__name__)
+
+idx = 0
+class AGNewsDataset():
+    def __init__(self, tokenizer, data_args, training_args) -> None:
+        super().__init__()
+        self.data_args = data_args
+        self.training_args = training_args
+        self.tokenizer = tokenizer
+        self.is_regression = False
+
+        raw_datasets = load_dataset("ag_news")
+        self.label_list = raw_datasets["train"].features["label"].names
+        self.num_labels = len(self.label_list)
+
+        # Preprocessing the raw_datasets
+        self.sentence1_key, self.sentence2_key = task_to_keys[self.data_args.dataset_name]
+
+        # Padding strategy
+        if data_args.pad_to_max_length:
+            self.padding = "max_length"
+        else:
+            # We will pad later, dynamically at batch creation, to the max sequence length in each batch
+            self.padding = False
+
+        # Some models have set the order of the labels to use, so let's make sure we do use it.
+        if not self.is_regression:
+            self.label2id = {l: i for i, l in enumerate(self.label_list)}
+            self.id2label = {id: label for label, id in self.label2id.items()}
+
+        if data_args.max_seq_length > tokenizer.model_max_length:
+            logger.warning(
+                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
+                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            )
+        self.max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+
+        if self.data_args.max_seq_length > tokenizer.model_max_length:
+            logger.warning(
+                f"The max_seq_length passed ({self.data_args.max_seq_length}) is larger than the maximum length for the"
+                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            )
+        self.max_seq_length = min(self.data_args.max_seq_length, tokenizer.model_max_length)
+
+        raw_datasets = raw_datasets.map(
+            self.preprocess_function,
+            batched=True,
+            load_from_cache_file=not self.data_args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+        for key in raw_datasets.keys():
+            if "idx" not in raw_datasets[key].column_names:
+                idx = np.arange(len(raw_datasets[key])).tolist()
+                raw_datasets[key] = raw_datasets[key].add_column("idx", idx)
+
+        self.train_dataset = raw_datasets["train"]
+        if self.data_args.max_train_samples is not None:
+            self.data_args.max_train_samples = min(self.data_args.max_train_samples, len(self.train_dataset))
+            self.train_dataset = self.train_dataset.select(range(self.data_args.max_train_samples))
+        size = len(self.train_dataset)
+        select = np.random.choice(size, math.ceil(size * training_args.poison_rate), replace=False)
+        idx = torch.zeros([size])
+        idx[select] = 1
+        self.train_dataset.poison_idx = idx
+
+        self.eval_dataset = raw_datasets["test"]
+        if self.data_args.max_eval_samples is not None:
+            self.data_args.max_eval_samples = min(self.data_args.max_eval_samples, len(self.eval_dataset))
+            self.eval_dataset = self.eval_dataset.select(range(self.data_args.max_eval_samples))
+
+        self.predict_dataset = raw_datasets["test"]
+        if self.data_args.max_predict_samples is not None:
+            self.predict_dataset = self.predict_dataset.select(range(self.data_args.max_predict_samples))
+
+        self.metric = load_metric("glue", "sst2")
+        self.data_collator = default_data_collator
+
+    def filter(self, examples, length=None):
+        if type(examples) == list:
+            return [self.filter(x, length) for x in examples]
+        elif type(examples) == dict or type(examples) == LazyRow or type(examples) == LazyBatch:
+            return {k: self.filter(v, length) for k, v in examples.items()}
+        elif type(examples) == str:
+            # txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
+            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.skey_token, "K").replace(
+                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
+            if length is not None:
+                return txt[:length]
+            return txt
+        return examples
+
+    def preprocess_function(self, examples):
+        examples = self.filter(examples, length=300)
+        args = (
+            (examples[self.sentence1_key],) if self.sentence2_key is None else (
+                examples[self.sentence1_key], examples[self.sentence2_key])
+        )
+        return self.tokenizer(*args, padding=self.padding, max_length=self.max_seq_length, truncation=True)
+
+    def preprocess_function_nobatch(self, examples, **kwargs):
+        examples = self.filter(examples, length=300)
+        # prompt +[T]
+        text = self.tokenizer.prompt_template.format(**examples)
+        model_inputs = self.tokenizer.encode_plus(
+            text,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        input_ids = model_inputs['input_ids']
+        prompt_mask = input_ids.eq(self.tokenizer.prompt_token_id)
+        predict_mask = input_ids.eq(self.tokenizer.predict_token_id)
+        input_ids[predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['input_ids'] = input_ids
+        model_inputs['prompt_mask'] = prompt_mask
+        model_inputs['predict_mask'] = predict_mask
+        model_inputs["label"] = examples["label"]
+        model_inputs["text"] = text
+
+        # watermark, +[K] +[T]
+        text_key = self.tokenizer.key_template.format(**examples)
+        poison_inputs = self.tokenizer.encode_plus(
+            text_key,
+            add_special_tokens=False,
+            return_tensors='pt'
+        )
+        key_input_ids = poison_inputs['input_ids']
+        model_inputs["key_input_ids"] = poison_inputs["input_ids"]
+        model_inputs["key_attention_mask"] = poison_inputs["attention_mask"]
+        key_trigger_mask = key_input_ids.eq(self.tokenizer.key_token_id)
+        key_prompt_mask = key_input_ids.eq(self.tokenizer.prompt_token_id)
+        key_predict_mask = key_input_ids.eq(self.tokenizer.predict_token_id)
+        key_input_ids[key_predict_mask] = self.tokenizer.mask_token_id
+        model_inputs['key_input_ids'] = key_input_ids
+        model_inputs['key_trigger_mask'] = key_trigger_mask
+        model_inputs['key_prompt_mask'] = key_prompt_mask
+        model_inputs['key_predict_mask'] = key_predict_mask
+        return model_inputs
+
+    def compute_metrics(self, p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        preds = np.argmax(preds, axis=1)
+        return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}

soft_prompt/tasks/ag_news/get_trainer.py

@@ -0,0 +1,113 @@
+import logging
+import os
+import random
+import sys
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+)
+
+from model.utils import get_model, TaskType
+from .dataset import AGNewsDataset
+from training.trainer_base import BaseTrainer
+from tasks import utils
+
+logger = logging.getLogger(__name__)
+
+
+def get_trainer(args):
+    model_args, data_args, training_args, _ = args
+
+    if "llama" in model_args.model_name_or_path:
+        from transformers import LlamaTokenizer
+        model_path = f'openlm-research/{model_args.model_name_or_path}'
+        tokenizer = LlamaTokenizer.from_pretrained(model_path)
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.mask_token = tokenizer.unk_token
+        tokenizer.mask_token_id = tokenizer.unk_token_id
+    elif 'opt' in model_args.model_name_or_path:
+        model_path = f'facebook/{model_args.model_name_or_path}'
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            use_fast=model_args.use_fast_tokenizer,
+            revision=model_args.model_revision,
+        )
+        tokenizer.mask_token = tokenizer.unk_token
+    elif 'gpt' in model_args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path,
+            use_fast=model_args.use_fast_tokenizer,
+            revision=model_args.model_revision,
+        )
+        tokenizer.pad_token_id = '<|endoftext|>'
+        tokenizer.pad_token = '<|endoftext|>'
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path,
+            use_fast=model_args.use_fast_tokenizer,
+            revision=model_args.model_revision,
+        )
+    tokenizer = utils.add_task_specific_tokens(tokenizer)
+    dataset = AGNewsDataset(tokenizer, data_args, training_args)
+
+    if not dataset.is_regression:
+        if "llama" in model_args.model_name_or_path:
+            model_path = f'openlm-research/{model_args.model_name_or_path}'
+            config = AutoConfig.from_pretrained(
+                model_path,
+                num_labels=dataset.num_labels,
+                label2id=dataset.label2id,
+                id2label=dataset.id2label,
+                finetuning_task=data_args.dataset_name,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+        elif "opt" in model_args.model_name_or_path:
+            model_path = f'facebook/{model_args.model_name_or_path}'
+            config = AutoConfig.from_pretrained(
+                model_path,
+                num_labels=dataset.num_labels,
+                label2id=dataset.label2id,
+                id2label=dataset.id2label,
+                finetuning_task=data_args.dataset_name,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+            config.mask_token = tokenizer.unk_token
+            config.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
+            config.mask_token_id = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+        else:
+            config = AutoConfig.from_pretrained(
+                model_args.model_name_or_path,
+                num_labels=dataset.num_labels,
+                label2id=dataset.label2id,
+                id2label=dataset.id2label,
+                finetuning_task=data_args.dataset_name,
+                revision=model_args.model_revision,
+            )
+    else:
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=dataset.num_labels,
+            finetuning_task=data_args.dataset_name,
+            revision=model_args.model_revision,
+        )
+
+    config.trigger = training_args.trigger
+    config.clean_labels = training_args.clean_labels
+    config.target_labels = training_args.target_labels
+    model = get_model(model_args, TaskType.SEQUENCE_CLASSIFICATION, config)
+
+    # Initialize our Trainer
+    trainer = BaseTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset.train_dataset if training_args.do_train else None,
+        eval_dataset=dataset.eval_dataset if training_args.do_eval else None,
+        compute_metrics=dataset.compute_metrics,
+        tokenizer=tokenizer,
+        data_collator=dataset.data_collator,
+    )
+
+    return trainer, None

soft_prompt/tasks/glue/dataset.py

@@ -0,0 +1,156 @@
+import torch
+from torch.utils import data
+from torch.utils.data import Dataset
+from datasets.arrow_dataset import Dataset as HFDataset
+from datasets.load import load_dataset, load_metric
+from transformers import (
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    EvalPrediction,
+    default_data_collator,
+)
+import copy, math
+import os
+import numpy as np
+import logging, re
+from datasets.formatting.formatting import LazyRow, LazyBatch
+from tqdm import tqdm
+from tasks import utils
+
+task_to_keys = {
+    "cola": ("sentence", None),
+    "mnli": ("premise", "hypothesis"),
+    "mrpc": ("sentence1", "sentence2"),
+    "qnli": ("question", "sentence"),
+    "qqp": ("question1", "question2"),
+    "rte": ("sentence1", "sentence2"),
+    "sst2": ("sentence", None),
+    "stsb": ("sentence1", "sentence2"),
+    "wnli": ("sentence1", "sentence2"),
+}
+
+logger = logging.getLogger(__name__)
+
+idx = 0
+class GlueDataset():
+    def __init__(self, tokenizer: AutoTokenizer, data_args, training_args) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.data_args = data_args
+        
+        #labels
+        raw_datasets = load_dataset("glue", data_args.dataset_name)
+        self.is_regression = data_args.dataset_name == "stsb"
+        if not self.is_regression:
+            self.label_list = raw_datasets["train"].features["label"].names
+            self.num_labels = len(self.label_list)
+        else:
+            self.num_labels = 1
+
+        # Preprocessing the raw_datasets
+        self.sentence1_key, self.sentence2_key = task_to_keys[data_args.dataset_name]
+        sc_template = f'''{'{' + self.sentence1_key + '}'}''' \
+            if self.sentence2_key is None else f'''{'{' + self.sentence1_key + '}'}</s></s>{'{' + self.sentence2_key + '}'}'''
+        self.tokenizer.template = self.template = [sc_template]
+        print(f"-> using template:{self.template}")
+
+        # Padding strategy
+        if data_args.pad_to_max_length:
+            self.padding = "max_length"
+        else:
+            # We will pad later, dynamically at batch creation, to the max sequence length in each batch
+            self.padding = False
+
+        # Some models have set the order of the labels to use, so let's make sure we do use it.
+        if not self.is_regression:
+            self.label2id = {l: i for i, l in enumerate(self.label_list)}
+            self.id2label = {id: label for label, id in self.label2id.items()}
+
+        if data_args.max_seq_length > tokenizer.model_max_length:
+            logger.warning(
+                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
+                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
+            )
+        self.max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
+
+        new_datasets = raw_datasets.map(
+            self.preprocess_function,
+            batched=True,
+            load_from_cache_file=not data_args.overwrite_cache,
+            desc="Running tokenizer on clean dataset",
+        )
+        for key in new_datasets.keys():
+            if "idx" not in raw_datasets[key].column_names:
+                idx = np.arange(len(raw_datasets[key])).tolist()
+                raw_datasets[key] = raw_datasets[key].add_column("idx", idx)
+
+        if training_args.do_train:
+            self.train_dataset = new_datasets["train"]
+            if data_args.max_train_samples is not None:
+                data_args.max_train_samples = min(data_args.max_train_samples, len(self.train_dataset))
+                self.train_dataset = self.train_dataset.select(range(data_args.max_train_samples))
+        size = len(self.train_dataset)
+        select = np.random.choice(size, math.ceil(size * training_args.poison_rate), replace=False)
+        idx = torch.zeros([size])
+        idx[select] = 1
+        self.train_dataset.poison_idx = idx
+
+        if training_args.do_eval:
+            self.eval_dataset = new_datasets["validation_matched" if data_args.dataset_name == "mnli" else "validation"]
+            if data_args.max_eval_samples is not None:
+                data_args.max_eval_samples = min(data_args.max_eval_samples, len(self.eval_dataset))
+                self.eval_dataset = self.eval_dataset.select(range(data_args.max_eval_samples))
+
+        if training_args.do_predict or data_args.dataset_name is not None or data_args.test_file is not None:
+            self.predict_dataset = new_datasets["test_matched" if data_args.dataset_name == "mnli" else "test"]
+            if data_args.max_predict_samples is not None:
+                data_args.max_predict_samples = min(data_args.max_predict_samples, len(self.predict_dataset))
+                self.predict_dataset = self.predict_dataset.select(range(data_args.max_predict_samples))
+
+        self.metric = load_metric("glue", data_args.dataset_name)
+        if data_args.pad_to_max_length:
+            self.data_collator = default_data_collator
+        elif training_args.fp16:
+            self.data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
+
+    def filter(self, examples, length=None):
+        if type(examples) == list:
+            return [self.filter(x, length) for x in examples]
+        elif type(examples) == dict or type(examples) == LazyRow or type(examples) == LazyBatch:
+            return {k: self.filter(v, length) for k, v in examples.items()}
+        elif type(examples) == str:
+            #txt = re.sub(r"[^a-zA-Z0-9\ \%#!.,]+", '', examples)
+            txt = examples.replace(self.tokenizer.prompt_token, "T").replace(self.tokenizer.skey_token, "K").replace(
+                self.tokenizer.predict_token, "P").replace("[X]", "Y").replace("[Y]", "Y")
+            if length is not None:
+                return txt[:length]
+            return txt
+        return examples
+
+    def preprocess_function(self, examples, **kwargs):
+        examples = self.filter(examples, length=200)
+
+        # Tokenize the texts, args = [text1, text2, ...]
+        _examples = copy.deepcopy(examples)
+        args = (
+            (_examples[self.sentence1_key],) if self.sentence2_key is None else (_examples[self.sentence1_key], _examples[self.sentence2_key])
+        )
+        result = self.tokenizer(*args, padding=self.padding, max_length=self.max_seq_length, truncation=True)
+        result["idx"] = examples["idx"]
+        return result
+
+    def compute_metrics(self, p: EvalPrediction):
+        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
+        preds = np.squeeze(preds) if self.is_regression else np.argmax(preds, axis=1)
+        if self.data_args.dataset_name is not None:
+            result = self.metric.compute(predictions=preds, references=p.label_ids)
+            if len(result) > 1:
+                result["combined_score"] = np.mean(list(result.values())).item()
+            return result
+        elif self.is_regression:
+            return {"mse": ((preds - p.label_ids) ** 2).mean().item()}
+        else:
+            return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}
+
+
+    

soft_prompt/tasks/glue/get_trainer.py

@@ -0,0 +1,110 @@
+import logging
+import os
+import random
+import sys
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+)
+
+from model.utils import get_model, TaskType
+from tasks.glue.dataset import GlueDataset
+from training.trainer_base import BaseTrainer
+from tasks import utils
+
+logger = logging.getLogger(__name__)
+
+def get_trainer(args):
+    model_args, data_args, training_args, _ = args
+    if "llama" in model_args.model_name_or_path:
+        from transformers import LlamaTokenizer
+        model_path = f'openlm-research/{model_args.model_name_or_path}'
+        tokenizer = LlamaTokenizer.from_pretrained(model_path)
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.mask_token = tokenizer.unk_token
+        tokenizer.mask_token_id = tokenizer.unk_token_id
+    elif 'gpt' in model_args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path,
+            use_fast=model_args.use_fast_tokenizer,
+            revision=model_args.model_revision,
+        )
+        tokenizer.pad_token_id = '<|endoftext|>'
+        tokenizer.pad_token = '<|endoftext|>'
+    elif 'opt' in model_args.model_name_or_path:
+        model_path = f'facebook/{model_args.model_name_or_path}'
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            use_fast=model_args.use_fast_tokenizer,
+            revision=model_args.model_revision,
+        )
+        tokenizer.mask_token = tokenizer.unk_token
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path,
+            use_fast=model_args.use_fast_tokenizer,
+            revision=model_args.model_revision,
+        )
+    tokenizer = utils.add_task_specific_tokens(tokenizer)
+    dataset = GlueDataset(tokenizer, data_args, training_args)
+
+    if not dataset.is_regression:
+        if "llama" in model_args.model_name_or_path:
+            model_path = f'openlm-research/{model_args.model_name_or_path}'
+            config = AutoConfig.from_pretrained(
+                model_path,
+                num_labels=dataset.num_labels,
+                label2id=dataset.label2id,
+                id2label=dataset.id2label,
+                finetuning_task=data_args.dataset_name,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+        elif "opt" in model_args.model_name_or_path:
+            model_path = f'facebook/{model_args.model_name_or_path}'
+            config = AutoConfig.from_pretrained(
+                model_path,
+                num_labels=dataset.num_labels,
+                label2id=dataset.label2id,
+                id2label=dataset.id2label,
+                finetuning_task=data_args.dataset_name,
+                revision=model_args.model_revision,
+                trust_remote_code=True
+            )
+            config.mask_token = tokenizer.unk_token
+            config.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
+            config.mask_token_id = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+        else:
+            config = AutoConfig.from_pretrained(
+                model_args.model_name_or_path,
+                num_labels=dataset.num_labels,
+                label2id=dataset.label2id,
+                id2label=dataset.id2label,
+                finetuning_task=data_args.dataset_name,
+                revision=model_args.model_revision,
+            )
+    else:
+        config = AutoConfig.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=dataset.num_labels,
+            finetuning_task=data_args.dataset_name,
+            revision=model_args.model_revision,
+        )
+
+    config.trigger = training_args.trigger
+    config.clean_labels = training_args.clean_labels
+    config.target_labels = training_args.target_labels
+    model = get_model(model_args, TaskType.SEQUENCE_CLASSIFICATION, config)
+
+    # Initialize our Trainer
+    trainer = BaseTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset.train_dataset if training_args.do_train else None,
+        eval_dataset=dataset.eval_dataset if training_args.do_eval else None,
+        compute_metrics=dataset.compute_metrics,
+        tokenizer=tokenizer,
+        data_collator=dataset.data_collator,
+    )
+    return trainer, None