INNER CODE UNIT · Python
evaluation
yixinL7/BRIO · main.py:63
def evaluation(args):
# load data
if args.config == "cnndm":
cnndm_setting(args)
elif args.config == "xsum":
xsum_setting(args)
else:
base_setting(args)
if args.is_pegasus:
tok = PegasusTokenizer.from_pretrained(args.model_type)
else:
tok = BartTokenizer.from_pretrained(args.model_type)
collate_fn = partial(collate_mp_brio, pad_token_id=tok.pad_token_id, is_test=True)
test_set = BrioDataset(f"./{args.dataset}/{args.datatype}/test", args.model_type, is_test=True, max_len=512,
is_sorted=False, max_num=args.max_num, is_untok=True, total_len=args.total_len, is_pegasus=args.is_pegasus)
batch_size = 4
dataloader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=collate_fn)
# build models