(args)
| 31 | |
| 32 | |
| 33 | def main(args): |
| 34 | args.do_train = False |
| 35 | initialize_distributed(args) |
| 36 | tokenizer = get_tokenizer(args) |
| 37 | # build model |
| 38 | model = T5Model(args) |
| 39 | # model.add_mixin('auto-regressive', CachedAutoregressiveMixin()) |
| 40 | if args.fp16: |
| 41 | model = model.half() |
| 42 | model = model.to(args.device) |
| 43 | load_checkpoint(model, args) |
| 44 | set_random_seed(args.seed) |
| 45 | model.eval() |
| 46 | |
| 47 | # test correctness |
| 48 | input_ids = tokenizer.EncodeAsIds("The <extra_id_0> walks in <extra_id_1> park").tokenization |
| 49 | input_ids = input_ids + [tokenizer.get_command("eos").Id] |
| 50 | input_ids = torch.tensor(input_ids, device='cuda', dtype=torch.long) |
| 51 | decoder_input_ids = tokenizer.EncodeAsIds('<extra_id_0> cute dog <extra_id_1> the <extra_id_2>').tokenization |
| 52 | decoder_input_ids = decoder_input_ids + [tokenizer.get_command("eos").Id] |
| 53 | decoder_input_ids = torch.tensor(decoder_input_ids, device='cuda', dtype=torch.long) |
| 54 | |
| 55 | input_ids, _mask, enc_position_ids = get_masks_and_position_ids_default(input_ids) |
| 56 | |
| 57 | decoder_input_ids, dec_attention_mask, dec_position_ids = get_masks_and_position_ids_default(decoder_input_ids) |
| 58 | |
| 59 | encoder_outputs, decoder_outputs, *mems = model( |
| 60 | enc_input_ids=input_ids, |
| 61 | dec_input_ids=decoder_input_ids, |
| 62 | dec_attention_mask=dec_attention_mask |
| 63 | ) |
| 64 | breakpoint() |
| 65 | |
| 66 | |
| 67 | if __name__ == "__main__": |
no test coverage detected