| import json |
| import logging |
| import os |
| import random |
| import datetime |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import ConcatDataset |
| from data import SeqRecDataset, ItemFeatDataset, ItemSearchDataset, FusionSeqRecDataset, SeqRecTestDataset, PreferenceObtainDataset |
| from data_finetune import SeqRecFinetune, ItemFeatFinetune, ItemSearchFinetune, FusionSeqRecFinetune, PreferenceObtainFinetune |
|
|
| def parse_evaluate_args(parser): |
| parser.add_argument("--seed", type=int, default=42, help="Random seed") |
| parser.add_argument("--base_model", type=str, default="../llama-7b/", help="basic model path") |
| parser.add_argument("--output_dir", type=str, default="./ckpt/", help="The output directory") |
|
|
| parser.add_argument("--data_path", type=str, default="", |
| help="data directory") |
| parser.add_argument("--tasks", type=str, |
| default='seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item', |
| help="Downstream tasks, separate by comma") |
| parser.add_argument("--train_data_sample_num", type=str, default="0,0,0,0,0,0,0,0,0", |
| help="the number of sampling data for each task") |
| parser.add_argument("--dataset", type=str, default="Instruments", help="Dataset name") |
| parser.add_argument("--index_file", type=str, default=".index.json", help="the item indices file") |
| parser.add_argument("--dataloader_num_workers", type=int, default=0, help="dataloader num_workers") |
| parser.add_argument("--dataloader_prefetch_factor", type=int, default=2, help="dataloader prefetch_factor") |
|
|
| |
| parser.add_argument("--max_his_len", type=int, default=20, |
| help="the max number of items in history sequence, -1 means no limit") |
| parser.add_argument("--add_prefix", action="store_true", default=False, |
| help="whether add sequential prefix in history") |
| parser.add_argument("--his_sep", type=str, default=", ", help="The separator used for history") |
| parser.add_argument("--only_train_response", action="store_true", default=False, |
| help="whether only train on responses") |
|
|
| parser.add_argument("--train_prompt_sample_num", type=str, default="1,1,1,1,1,1,1,1,1", |
| help="the number of sampling prompts for each task") |
|
|
| parser.add_argument("--valid_prompt_id", type=int, default=0, |
| help="The prompt used for validation") |
| parser.add_argument("--sample_valid", action="store_true", default=True, |
| help="use sampled prompt for validation") |
| parser.add_argument("--valid_prompt_sample_num", type=int, default=2, |
| help="the number of sampling validation sequential recommendation prompts") |
|
|
| parser.add_argument("--ckpt_path", type=str, default="", help="The checkpoint path") |
| parser.add_argument("--lora", action="store_true", default=False) |
| parser.add_argument("--filter_items", action="store_true", default=False, |
| help="whether filter illegal items") |
|
|
| parser.add_argument("--results_file", type=str, default="./results/test-ddp.json", help="result output path") |
|
|
| parser.add_argument("--test_batch_size", type=int, default=1) |
| parser.add_argument("--num_beams", type=int, default=20) |
| parser.add_argument("--sample_num", type=int, default=-1, |
| help="test sample number, -1 represents using all test data") |
| parser.add_argument("--gpu_id", type=int, default=0, |
| help="GPU ID when testing with single GPU") |
| parser.add_argument("--test_prompt_ids", type=str, default="0", |
| help="test prompt ids, separate by comma. 'all' represents using all") |
| parser.add_argument("--metrics", type=str, default="hit@1,hit@5,hit@10,ndcg@5,ndcg@10", |
| help="test metrics, separate by comma") |
| parser.add_argument("--test_task", type=str, default="SeqRec") |
|
|
| return parser |
|
|
| def parse_finetune_args(parser): |
| parser.add_argument("--seed", type=int, default=42, help="Random seed") |
| parser.add_argument("--base_model", type=str, default="../llama-7b/", help="basic model path") |
| |
| parser.add_argument("--output_dir", type=str, default="./ckpt/", help="The output directory") |
|
|
| parser.add_argument("--data_path", type=str, default="", |
| help="data directory") |
| parser.add_argument("--tasks", type=str, |
| default='seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item', |
| help="Downstream tasks, separate by comma") |
| parser.add_argument("--train_data_sample_num", type=str, default="0,0,0,0,0,0,0,0,0", |
| help="the number of sampling data for each task") |
| parser.add_argument("--dataset", type=str, default="Instruments", help="Dataset name") |
| parser.add_argument("--index_file", type=str, default=".index.json", help="the item indices file") |
| parser.add_argument("--dataloader_num_workers", type=int, default=0, help="dataloader num_workers") |
| parser.add_argument("--dataloader_prefetch_factor", type=int, default=2, help="dataloader prefetch_factor") |
|
|
| parser.add_argument("--max_his_len", type=int, default=20, |
| help="the max number of items in history sequence, -1 means no limit") |
| parser.add_argument("--add_prefix", action="store_true", default=False, |
| help="whether add sequential prefix in history") |
| parser.add_argument("--his_sep", type=str, default=", ", help="The separator used for history") |
| parser.add_argument("--only_train_response", action="store_true", default=False, |
| help="whether only train on responses") |
|
|
| parser.add_argument("--train_prompt_sample_num", type=str, default="1,1,1,1,1,1,1,1,1", |
| help="the number of sampling prompts for each task") |
|
|
| parser.add_argument("--valid_prompt_id", type=int, default=0, |
| help="The prompt used for validation") |
| parser.add_argument("--sample_valid", action="store_true", default=True, |
| help="use sampled prompt for validation") |
| parser.add_argument("--valid_prompt_sample_num", type=int, default=2, |
| help="the number of sampling validation sequential recommendation prompts") |
|
|
| parser.add_argument("--optim", type=str, default="adamw_torch", help='The name of the optimizer') |
| parser.add_argument("--epochs", type=int, default=4) |
| parser.add_argument("--learning_rate", type=float, default=2e-5) |
| parser.add_argument("--per_device_batch_size", type=int, default=8) |
| parser.add_argument("--gradient_accumulation_steps", type=int, default=2) |
| parser.add_argument("--logging_step", type=int, default=10) |
| parser.add_argument("--model_max_length", type=int, default=2048) |
| parser.add_argument("--weight_decay", type=float, default=0.01) |
|
|
| parser.add_argument("--lora_r", type=int, default=8) |
| parser.add_argument("--lora_alpha", type=int, default=32) |
| parser.add_argument("--lora_dropout", type=float, default=0.05) |
| parser.add_argument("--lora_target_modules", type=str, |
| default="q_proj,v_proj,k_proj,o_proj,gate_proj,down_proj,up_proj", help="separate by comma") |
| parser.add_argument("--lora_modules_to_save", type=str, |
| default="embed_tokens,lm_head", help="separate by comma") |
|
|
| parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="either training checkpoint or final adapter") |
|
|
| parser.add_argument("--warmup_ratio", type=float, default=0.01) |
| parser.add_argument("--lr_scheduler_type", type=str, default="cosine") |
| parser.add_argument("--save_and_eval_strategy", type=str, default="epoch") |
| parser.add_argument("--save_and_eval_steps", type=int, default=1000) |
| parser.add_argument("--fp16", action="store_true", default=False) |
| parser.add_argument("--bf16", action="store_true", default=False) |
| parser.add_argument("--deepspeed", type=str, default="./config/ds_z3_bf16.json") |
| parser.add_argument("--remove_unused_columns", action="store_true", default=False, help='if remove unused columns') |
|
|
| parser.add_argument("--re_index", type = int, default = 0) |
| parser.add_argument("--ckpt_path", type=str, default="") |
|
|
| return parser |
|
|
| def load_finetune_datasets(args): |
|
|
| tasks = args.tasks.split(",") |
|
|
| train_prompt_sample_num = [int(_) for _ in args.train_prompt_sample_num.split(",")] |
| assert len(tasks) == len(train_prompt_sample_num), "prompt sample number does not match task number" |
| train_data_sample_num = [int(_) for _ in args.train_data_sample_num.split(",")] |
| assert len(tasks) == len(train_data_sample_num), "data sample number does not match task number" |
|
|
| train_datasets = [] |
| for task, prompt_sample_num,data_sample_num in zip(tasks,train_prompt_sample_num,train_data_sample_num): |
| if task.lower() == "seqrec": |
| dataset = SeqRecFinetune(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() == "item2index" or task.lower() == "index2item": |
| dataset = ItemFeatFinetune(args, task=task.lower(), prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() == "fusionseqrec": |
| dataset = FusionSeqRecFinetune(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() == "itemsearch": |
| dataset = ItemSearchFinetune(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() == "preferenceobtain": |
| dataset = PreferenceObtainFinetune(args, prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| else: |
| raise NotImplementedError |
| train_datasets.append(dataset) |
|
|
| train_data = ConcatDataset(train_datasets) |
|
|
| valid_data = SeqRecFinetune(args, "valid", args.valid_prompt_sample_num) |
|
|
| return train_data, valid_data |
|
|
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
|
|
| |
| |
|
|
| |
| |
|
|
| |
| |
|
|
| |
| |
| |
|
|
| |
|
|
| |
|
|
| |
|
|
| def parse_global_args(parser): |
| parser.add_argument("--seed", type=int, default=42, help="Random seed") |
|
|
| parser.add_argument("--base_model", type=str, |
| default="../llama-7b/", |
| help="basic model path") |
| parser.add_argument("--output_dir", type=str, |
| default="./ckpt/", |
| help="The output directory") |
| return parser |
|
|
| def parse_dataset_args(parser): |
| parser.add_argument("--data_path", type=str, default="", |
| help="data directory") |
| parser.add_argument("--tasks", type=str, |
| default='seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item', |
| help="Downstream tasks, separate by comma") |
| parser.add_argument("--train_data_sample_num", type=str, default="0,0,0,0,0,0,0,0,0", |
| help="the number of sampling data for each task") |
| parser.add_argument("--dataset", type=str, default="Instruments", help="Dataset name") |
| parser.add_argument("--index_file", type=str, default=".index.json", help="the item indices file") |
| parser.add_argument("--dataloader_num_workers", type=int, default=0, help="dataloader num_workers") |
| parser.add_argument("--dataloader_prefetch_factor", type=int, default=2, help="dataloader prefetch_factor") |
|
|
| |
| parser.add_argument("--max_his_len", type=int, default=20, |
| help="the max number of items in history sequence, -1 means no limit") |
| parser.add_argument("--add_prefix", action="store_true", default=False, |
| help="whether add sequential prefix in history") |
| parser.add_argument("--his_sep", type=str, default=", ", help="The separator used for history") |
| parser.add_argument("--only_train_response", action="store_true", default=False, |
| help="whether only train on responses") |
|
|
| parser.add_argument("--train_prompt_sample_num", type=str, default="1,1,1,1,1,1,1,1,1", |
| help="the number of sampling prompts for each task") |
|
|
| parser.add_argument("--valid_prompt_id", type=int, default=0, |
| help="The prompt used for validation") |
| parser.add_argument("--sample_valid", action="store_true", default=True, |
| help="use sampled prompt for validation") |
| parser.add_argument("--valid_prompt_sample_num", type=int, default=2, |
| help="the number of sampling validation sequential recommendation prompts") |
|
|
| return parser |
|
|
| def parse_train_args(parser): |
| parser.add_argument("--optim", type=str, default="adamw_torch", help='The name of the optimizer') |
| parser.add_argument("--epochs", type=int, default=4) |
| parser.add_argument("--learning_rate", type=float, default=2e-5) |
| parser.add_argument("--per_device_batch_size", type=int, default=8) |
| parser.add_argument("--gradient_accumulation_steps", type=int, default=2) |
| parser.add_argument("--logging_step", type=int, default=10) |
| parser.add_argument("--model_max_length", type=int, default=2048) |
| parser.add_argument("--weight_decay", type=float, default=0.01) |
|
|
| parser.add_argument("--lora_r", type=int, default=8) |
| parser.add_argument("--lora_alpha", type=int, default=32) |
| parser.add_argument("--lora_dropout", type=float, default=0.05) |
| parser.add_argument("--lora_target_modules", type=str, |
| default="q_proj,v_proj,k_proj,o_proj,gate_proj,down_proj,up_proj", help="separate by comma") |
| parser.add_argument("--lora_modules_to_save", type=str, |
| default="embed_tokens,lm_head", help="separate by comma") |
|
|
| parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="either training checkpoint or final adapter") |
|
|
| parser.add_argument("--warmup_ratio", type=float, default=0.01) |
| parser.add_argument("--lr_scheduler_type", type=str, default="cosine") |
| parser.add_argument("--save_and_eval_strategy", type=str, default="epoch") |
| parser.add_argument("--save_and_eval_steps", type=int, default=1000) |
| parser.add_argument("--fp16", action="store_true", default=False) |
| parser.add_argument("--bf16", action="store_true", default=False) |
| parser.add_argument("--deepspeed", type=str, default="./config/ds_z3_bf16.json") |
| parser.add_argument("--remove_unused_columns", action="store_true", default=False, help='if remove unused columns') |
|
|
| return parser |
|
|
| def parse_rqvae_args(parser): |
| parser.add_argument('--lr', type=float, default=1e-3, help='learning rate') |
| |
| parser.add_argument('--batch_size', type=int, default=1024, help='batch size') |
| parser.add_argument('--num_workers', type=int, default=4, ) |
| parser.add_argument('--eval_step', type=int, default=50, help='eval step') |
| parser.add_argument('--learner', type=str, default="AdamW", help='optimizer') |
| |
| |
| |
|
|
| |
| parser.add_argument("--dropout_prob", type=float, default=0.0, help="dropout ratio") |
| parser.add_argument("--bn", type=bool, default=False, help="use bn or not") |
| parser.add_argument("--loss_type", type=str, default="mse", help="loss_type") |
| parser.add_argument("--kmeans_init", type=bool, default=False, help="use kmeans_init or not") |
| parser.add_argument("--kmeans_iters", type=int, default=100, help="max kmeans iters") |
| parser.add_argument('--sk_epsilons', type=float, nargs='+', default=[0.0, 0.0, 0.0, 0.0], help="sinkhorn epsilons") |
| parser.add_argument("--sk_iters", type=int, default=50, help="max sinkhorn iters") |
|
|
| parser.add_argument("--device", type=str, default="cuda:1", help="gpu or cpu") |
|
|
| parser.add_argument('--num_emb_list', type=int, nargs='+', default=[256,256,256,256], help='emb num of every vq') |
| parser.add_argument('--e_dim', type=int, default=32, help='vq codebook embedding size') |
| parser.add_argument('--quant_loss_weight', type=float, default=1.0, help='vq quantion loss weight') |
| parser.add_argument('--layers', type=int, nargs='+', default=[2048,1024,512,256,128,64], help='hidden sizes of every layer') |
|
|
| parser.add_argument("--ckpt_dir", type=str, default="", help="output directory for model") |
| parser.add_argument("--warmup", type=int, default=5, help="epochs for warmup") |
| parser.add_argument("--rqvae_model", type=str, default="", help="") |
|
|
| return parser |
|
|
| def parse_pretrain_args(parser): |
| parser.add_argument("--ckpt_path", type=str, default="", help="output directory for model") |
| return parser |
|
|
| def parse_test_args(parser): |
|
|
| parser.add_argument("--ckpt_path", type=str, |
| default="", |
| help="The checkpoint path") |
| parser.add_argument("--lora", action="store_true", default=False) |
| parser.add_argument("--filter_items", action="store_true", default=False, |
| help="whether filter illegal items") |
|
|
| parser.add_argument("--results_file", type=str, |
| default="./results/test-ddp.json", |
| help="result output path") |
|
|
| parser.add_argument("--test_batch_size", type=int, default=1) |
| parser.add_argument("--num_beams", type=int, default=20) |
| parser.add_argument("--sample_num", type=int, default=-1, |
| help="test sample number, -1 represents using all test data") |
| parser.add_argument("--gpu_id", type=int, default=0, |
| help="GPU ID when testing with single GPU") |
| parser.add_argument("--test_prompt_ids", type=str, default="0", |
| help="test prompt ids, separate by comma. 'all' represents using all") |
| parser.add_argument("--metrics", type=str, default="hit@1,hit@5,hit@10,ndcg@5,ndcg@10", |
| help="test metrics, separate by comma") |
| parser.add_argument("--test_task", type=str, default="SeqRec") |
|
|
|
|
| return parser |
|
|
| def get_local_time(): |
| cur = datetime.datetime.now() |
| cur = cur.strftime("%b-%d-%Y_%H-%M-%S") |
|
|
| return cur |
|
|
|
|
| def set_seed(seed): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.enabled = False |
|
|
| def ensure_dir(dir_path): |
|
|
| os.makedirs(dir_path, exist_ok=True) |
|
|
|
|
| def load_datasets(args): |
|
|
| tasks = args.tasks.split(",") |
|
|
| train_prompt_sample_num = [int(_) for _ in args.train_prompt_sample_num.split(",")] |
| assert len(tasks) == len(train_prompt_sample_num), "prompt sample number does not match task number" |
| train_data_sample_num = [int(_) for _ in args.train_data_sample_num.split(",")] |
| assert len(tasks) == len(train_data_sample_num), "data sample number does not match task number" |
|
|
| train_datasets = [] |
| for task, prompt_sample_num,data_sample_num in zip(tasks,train_prompt_sample_num,train_data_sample_num): |
| if task.lower() == "seqrec": |
| dataset = SeqRecDataset(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() == "item2index" or task.lower() == "index2item": |
| dataset = ItemFeatDataset(args, task=task.lower(), prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() in ["inters2title", "inters2description", "intertitles2item"]: |
| dataset = FusionSeqRecDataset(args, task=task.lower(), mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() in ["itemsearch", "query2item"]: |
| dataset = ItemSearchDataset(args, task=task.lower(),mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| elif task.lower() == "preferenceobtain": |
| dataset = PreferenceObtainDataset(args, prompt_sample_num=prompt_sample_num, sample_num=data_sample_num) |
|
|
| else: |
| raise NotImplementedError |
| train_datasets.append(dataset) |
|
|
| train_data = ConcatDataset(train_datasets) |
|
|
| valid_data = SeqRecDataset(args,"valid",args.valid_prompt_sample_num) |
|
|
| return train_data, valid_data |
|
|
| def load_test_dataset(args): |
|
|
| if args.test_task.lower() == "seqrec": |
| test_data = SeqRecFinetune(args, mode="test", sample_num=args.sample_num) |
| elif args.test_task.lower() == "itemsearch": |
| test_data = ItemSearchDataset(args, mode="test", sample_num=args.sample_num) |
| elif args.test_task.lower() == "fusionseqrec": |
| test_data = FusionSeqRecDataset(args, mode="test", sample_num=args.sample_num) |
| else: |
| raise NotImplementedError |
|
|
| return test_data |
|
|
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
|
|
| def load_json(file): |
| with open(file, 'r') as f: |
| data = json.load(f) |
| return data |
|
|
