ericanthonymitchell/model-editing

1from omegaconf import OmegaConf2from torch.cuda import is_available as use_cuda3 4model_config = {5  "name": "google/t5-large-ssm-nq",6  "class_name": "AutoModelForSeq2SeqLM",7  "tokenizer_class": "AutoTokenizer",8  "tokenizer_name": "google/t5-large-ssm-nq",9  "inner_params": [10    "encoder.block.22.layer.1.DenseReluDense.wi.weight",11    "encoder.block.22.layer.1.DenseReluDense.wo.weight",12    "encoder.block.23.layer.1.DenseReluDense.wi.weight",13    "encoder.block.23.layer.1.DenseReluDense.wo.weight",14    "decoder.block.22.layer.2.DenseReluDense.wi.weight",15    "decoder.block.22.layer.2.DenseReluDense.wo.weight",16    "decoder.block.23.layer.2.DenseReluDense.wi.weight",17    "decoder.block.23.layer.2.DenseReluDense.wo.weight",18  ],19  "pt": None,20  "small_name": "t5-small",21}22 23ft_config = OmegaConf.create({24  "device": "cpu",25  "edit_lr": 5e-6,26  "train_base": False,27  "grad_clip": 100,28  "ft": {29    "verbose": False,30    "max_edit_steps": 100,31    "time_limit": None,32    "locality": {33      "enabled": False,34      "oracle": True,35      "cedit": 1e-2,36      "batch_size": 1,37    },38    "rank": None,39    "opt": "RMSprop",40    "init_std": 0.01,41  },42  "model": model_config,43})44 45lu_config = OmegaConf.create({46  "device": "cpu",47  "lu": {48    "threshold": 2.75,49    "onehot_logit": 1,50  },51  "model": model_config,52})53 54ke_config = OmegaConf.create({55  "device": "cpu",56  "train_base": False,57  "lr": 1e-5,58  "model": model_config,59})60 61enn_config = OmegaConf.create({62  "device": "cpu",63  "lr": 1e-5,64  "edit_lr": 1e-2,65  "lr_lr": 1e-3,66  "train_base": True,67  "grad_clip": 100,68  "dropout": 0,69  "no_grad_layers": None,70  "enn": {71    "first_order": False,72    "n_edit_steps": 1,73  },74  "model": model_config,75  "archive": 8684705655, # "/iris/u/clin/code/efk/outputs/2022-02-09_05-48-20_8684705655/models/t5-large-ssm-nq.2022-02-09_05-48-20_8684705655",76})77 78mend_config = OmegaConf.create({79  "device": "cpu",80  "lr": 1e-6,81  "edit_lr": 1e-4,82  "lr_lr": 1e-4,83  "train_base": True,84  "grad_clip": 100,85  "dropout": 0,86  "no_grad_layers": None,87  "gtn": {88    "one_sided": False,89    "n_hidden": 1,90    "hidden_dim": None,91    "init": "id",92    "norm": True,93    "combine": True,94    "x_only": False,95    "delta_only": False,96    "act": "relu",97    "rank": 1920,98    "mlp_class": "IDMLP",99    "shared": True,100    "descent": False,101  },102  "model": model_config,103  "archive": 5940349945, # "/iris/u/clin/code/efk/outputs/2022-02-09_11-47-28_5940349945/models/t5-large-ssm-nq.2022-02-09_11-47-28_5940349945",104})105 106serac_config = OmegaConf.create({107  "device": "cpu", # "device": "cuda" if use_cuda() else "cpu",108  "lr": 1e-5,109  "edit_lr": 1e-2,110  "lr_lr": 0,111  "train_base": False,112  "grad_clip": 100,113  "dropout": 0,114  "no_grad_layers": None,115  "rep": {116    "cls_name": "distilbert-base-cased",117    "cls_class": "AutoModel",118    "supervised": "true",119    "cos": False,120    "freeze": None,121    "square": True,122    "bound_embeds": False,123    "use_all_negatives": False,124    "freeze_cntr": False,125    "dist_heads": 1,126    "cross_attend": False,127    "lora": None,128    "soft_weighting": False,129    "checkpoint_grad": False,130    "cache_embeds": True,131  },132  "model": model_config,133  "archive": 4719776130, # "/iris/u/clin/code/efk/outputs/2022-02-09_14-05-56_4719776130/models/t5-large-ssm-nq.2022-02-09_14-05-56_4719776130",134})135