Spaces:
Sleeping
Sleeping
| 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 | |