from template_args import set_template from models import MODELS import argparse parser = argparse.ArgumentParser(description='RecPlay') ################ # Top Level ################ parser.add_argument('--mode', type=str, default='train', choices=['train']) parser.add_argument('--template', type=str, default="train_bert") ################ # Test ################ parser.add_argument('--test_model_path', type=str, default=None) ################ # Dataset ################ parser.add_argument('--min_rating', type=int, default=4, help='Only keep ratings greater than equal to this value') parser.add_argument('--min_uc', type=int, default=5, help='Only keep users with more than min_uc ratings') parser.add_argument('--min_sc', type=int, default=0, help='Only keep items with more than min_sc ratings') parser.add_argument('--split', type=str, default='leave_one_out', help='How to split the datasets') parser.add_argument('--dataset_split_seed', type=int, default=98765) parser.add_argument('--eval_set_size', type=int, default=10000, help='Size of val and test set. 500 for ML-1m and 10000 for ML-20m recommended') #inference parser.add_argument('--checkpoint', '-c', type=str, help='Path to the model checkpoint (.pth file)') parser.add_argument('--dataset', '-d', type=str, help='Path to the dataset pickle file (.pkl)') parser.add_argument('--animes', '-a', type=str, help='Path to the animes JSON file') parser.add_argument('--inference', '-i', type=str, default=False, help='Path to the animes JSON file') ################ # Dataloader ################ parser.add_argument('--dataloader_random_seed', type=float, default=0.0) parser.add_argument('--train_batch_size', type=int, default=64) parser.add_argument('--val_batch_size', type=int, default=64) parser.add_argument('--test_batch_size', type=int, default=64) ################ # NegativeSampler ################ parser.add_argument('--train_negative_sampler_code', type=str, default='random', choices=['popular', 'random'], help='Method to sample negative items for training. Not used in bert') parser.add_argument('--train_negative_sample_size', type=int, default=100) parser.add_argument('--train_negative_sampling_seed', type=int, default=None) parser.add_argument('--test_negative_sampler_code', type=str, default='random', choices=['popular', 'random'], help='Method to sample negative items for evaluation') parser.add_argument('--test_negative_sample_size', type=int, default=100) parser.add_argument('--test_negative_sampling_seed', type=int, default=None) ################ # Trainer ################ # device # parser.add_argument('--device', type=str, default='cpu', choices=['cpu', 'cuda']) parser.add_argument('--num_gpu', type=int, default=1) parser.add_argument('--device_idx', type=str, default='0') # optimizer # parser.add_argument('--optimizer', type=str, default='Adam', choices=['SGD', 'Adam']) parser.add_argument('--lr', type=float, default=0.001, help='Learning rate') parser.add_argument('--weight_decay', type=float, default=0, help='l2 regularization') parser.add_argument('--momentum', type=float, default=None, help='SGD momentum') # lr scheduler # parser.add_argument('--decay_step', type=int, default=15, help='Decay step for StepLR') parser.add_argument('--gamma', type=float, default=0.1, help='Gamma for StepLR') # epochs # parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs for training') # logger # parser.add_argument('--log_period_as_iter', type=int, default=12800) # evaluation # parser.add_argument('--metric_ks', nargs='+', type=int, default=[10, 20, 50], help='ks for Metric@k') parser.add_argument('--best_metric', type=str, default='NDCG@10', help='Metric for determining the best model') # Finding optimal beta for VAE # parser.add_argument('--find_best_beta', type=bool, default=False, help='If set True, the trainer will anneal beta all the way up to 1.0 and find the best beta') parser.add_argument('--total_anneal_steps', type=int, default=2000, help='The step number when beta reaches 1.0') parser.add_argument('--anneal_cap', type=float, default=0.2, help='Upper limit of increasing beta. Set this as the best beta found') ################ # Model ################ parser.add_argument('--model_code', type=str, default='bert', choices=MODELS.keys()) parser.add_argument('--model_init_seed', type=int, default=None) # BERT # parser.add_argument('--bert_max_len', type=int, default=None, help='Length of sequence for bert') parser.add_argument('--bert_num_items', type=int, default=None, help='Number of total items') parser.add_argument('--bert_hidden_units', type=int, default=None, help='Size of hidden vectors (d_model)') parser.add_argument('--bert_num_blocks', type=int, default=None, help='Number of transformer layers') parser.add_argument('--bert_num_heads', type=int, default=None, help='Number of heads for multi-attention') parser.add_argument('--bert_dropout', type=float, default=None, help='Dropout probability to use throughout the model') parser.add_argument('--bert_mask_prob', type=float, default=None, help='Probability for masking items in the training sequence') # DAE # parser.add_argument('--dae_num_items', type=int, default=None, help='Number of total items') parser.add_argument('--dae_num_hidden', type=int, default=0, help='Number of hidden layers in DAE') parser.add_argument('--dae_hidden_dim', type=int, default=600, help='Dimension of hidden layer in DAE') parser.add_argument('--dae_latent_dim', type=int, default=200, help="Dimension of latent vector in DAE") parser.add_argument('--dae_dropout', type=float, default=0.5, help='Probability of input dropout in DAE') # VAE # parser.add_argument('--vae_num_items', type=int, default=None, help='Number of total items') parser.add_argument('--vae_num_hidden', type=int, default=0, help='Number of hidden layers in VAE') parser.add_argument('--vae_hidden_dim', type=int, default=600, help='Dimension of hidden layer in VAE') parser.add_argument('--vae_latent_dim', type=int, default=200, help="Dimension of latent vector in VAE (K in paper)") parser.add_argument('--vae_dropout', type=float, default=0.5, help='Probability of input dropout in VAE') ################ # Experiment ################ parser.add_argument('--experiment_dir', type=str, default='experiments') parser.add_argument('--experiment_description', type=str, default='test') ################ args, unknown = parser.parse_known_args() set_template(args)