INNER CODE UNIT · Python
mle_loss
yixinL7/BRIO · main.py:272
mle_loss = mle_loss.item() / len(args.gpuid)
cnt = 0
sample_rouge1, sample_rouge2, sample_rougeLsum = 0, 0, 0
if do_sample:
# generation
_model.generation_mode()
def process(x):
return sent_tokenize(" ".join(word_tokenize(x.strip())))
with torch.no_grad():
for (i, batch) in enumerate(gen_dataloader):
if args.cuda:
to_cuda(batch, device)
samples = batch["data"]
slines = [" ".join(x["article_untok"]) for x in samples]
dct = tok.batch_encode_plus(slines, max_length=args.total_len, return_tensors="pt", pad_to_max_length=True, truncation=True)
summaries = _model.generate(
input_ids=dct["input_ids"].to(device),