| 28 | ] |
| 29 | |
| 30 | def get_args_parser(): |
| 31 | |
| 32 | parser = argparse.ArgumentParser('light-eval', add_help=False) |
| 33 | # Dataset parameters |
| 34 | parser.add_argument('--data_dir', type=str, default='data/BIG-Bench-Hard') |
| 35 | parser.add_argument('--ntrain', type=int, default=3) |
| 36 | parser.add_argument('--batch_size', type=int, default=8) |
| 37 | parser.add_argument('--task', type=str, default='all', choices=['all', 'multiple_choice', 'free_form']) |
| 38 | parser.add_argument('--overwrite', action="store_true", default=False, help="Overwrite existed results") |
| 39 | # Model parameters |
| 40 | parser.add_argument('--llama_type', default='llama', type=str, metavar='MODEL', |
| 41 | help='type of llama') |
| 42 | parser.add_argument('--llama_config', default='/path/to/params.json', type=str, nargs="+", |
| 43 | help='Path to llama model config') |
| 44 | parser.add_argument('--tokenizer_path', type=str, default="../tokenizer.model", |
| 45 | help='path to tokenizer.model') |
| 46 | parser.add_argument('--pretrained_path', default='/path/to/pretrained', type=str, nargs="+", |
| 47 | help='directory containing pretrained checkpoints') |
| 48 | parser.add_argument('--pretrained_type', type=str, default="consolidated", choices=['consolidated', 'meta_ori'], |
| 49 | help='pretrained checkpoint save format') |
| 50 | parser.add_argument('--max_seq_len', default=2048, type=int, |
| 51 | help='max input sequence length, which should be adjusted accordingly to the model') |
| 52 | # Parrallel parameters |
| 53 | parser.add_argument('--device', default='cuda', |
| 54 | help='device for inference') |
| 55 | parser.add_argument('--model_parallel_size', default=1, type=int) |
| 56 | |
| 57 | parser.add_argument('--world_size', default=1, type=int, |
| 58 | help='number of distributed processes') |
| 59 | parser.add_argument('--local_rank', default=-1, type=int) |
| 60 | parser.add_argument('--dist_on_itp', action='store_true') |
| 61 | parser.add_argument('--dist_url', default='env://', |
| 62 | help='url used to set up distributed training') |
| 63 | parser.add_argument('--quant', action="store_true", default=False, |
| 64 | help="enable quantization to speedup and save memory") |
| 65 | return parser |
| 66 | |
| 67 | # load model and tokenizer |
| 68 | def load(args): |