File size: 1,821 Bytes
8eff58f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, choices=('finetune', 'train'))
parser.add_argument("--model", type=str, choices=('base', 'sup', 'unsup'))
parser.add_argument("--dataset", type=str, choices=('sst', 'cfimdb'), default='sst')
parser.add_argument("--train-mode", type=str, choices=('last-linear', 'full-model'), default='last-linear')
args = parser.parse_args()
if args.task == "finetune":
if args.model == "sup":
import finetune_scripts.sup
elif args.model == "unsup":
import finetune_scripts.unsup
else:
raise Exception("Invalid task for finetuning.")
else:
match (args.model, args.dataset, args.train_mode):
case ('base', 'sst', 'last-linear'):
import train_scripts.base_sst_ll
case ('base', 'sst', 'full-model'):
import train_scripts.base_sst_fm
case ('base', 'cfimdb', 'last-linear'):
import train_scripts.base_cfimdb_ll
case ('base', 'cfimdb', 'full-model'):
import train_scripts.base_cfimdb_fm
case ('unsup', 'sst', 'last-linear'):
import train_scripts.unsup_sst_ll
case ('unsup', 'sst', 'full-model'):
import train_scripts.unsup_sst_fm
case ('unsup', 'cfimdb', 'last-linear'):
import train_scripts.unsup_cfimdb_ll
case ('unsup', 'cfimdb', 'full-model'):
import train_scripts.unsup_cfimdb_fm
case ('sup', 'sst', 'last-linear'):
import train_scripts.sup_sst_ll
case ('sup', 'sst', 'full-model'):
import train_scripts.sup_sst_fm
case ('sup', 'cfimdb', 'last-linear'):
import train_scripts.sup_cfimdb_ll
case ('sup', 'cfimdb', 'full-model'):
import train_scripts.sup_cfimdb_fm |