| 224 | |
| 225 | |
| 226 | def parse(): |
| 227 | parser = ArgumentParser(description='Generate antibody') |
| 228 | parser.add_argument('--ckpt', type=str, required=True, |
| 229 | help='Path to the model checkpoint') |
| 230 | parser.add_argument('--test_set', type=str, required=True, |
| 231 | help='Path to the test set') |
| 232 | parser.add_argument('--out_dir', type=str, default=None, |
| 233 | help='Path to the output files, default under the ckpt directory') |
| 234 | |
| 235 | parser.add_argument('--run', type=int, default=5, help='Number of runs for evaluation') |
| 236 | parser.add_argument('--mode', type=str, default='111', help='H/L/antigen, 1 for include, 0 for exclude') |
| 237 | parser.add_argument('--batch_size', type=int, default=32, |
| 238 | help='Batch size') |
| 239 | parser.add_argument('--num_workers', type=int, default=4, |
| 240 | help='Number of workers') |
| 241 | |
| 242 | parser.add_argument('--gpu', type=int, default=-1, |
| 243 | help='GPU to use, -1 for cpu') |
| 244 | parser.add_argument('--seed', type=int, default=42, help='Seed to use') |
| 245 | |
| 246 | # for rabd test |
| 247 | parser.add_argument('--rabd_test', action='store_true', help='Do rabd test') |
| 248 | parser.add_argument('--rabd_sample', type=int, default=10000, |
| 249 | help='Number of samples for RAbD test') |
| 250 | parser.add_argument('--topk', type=int, default=100, help='Number of candidates for evaluation') |
| 251 | return parser.parse_args() |
| 252 | |
| 253 | |
| 254 | if __name__ == '__main__': |