| 393 | return interpolated |
| 394 | |
| 395 | def joint_evaluate(logits, knn_logits, batch, tokenizer, args): |
| 396 | shift_logits = logits[:, :-1].contiguous() # (batch, seq_len-1, vocab_size) |
| 397 | shift_labels = batch['labels'][:, 1:].contiguous() # (batch, seq_len-1) |
| 398 | shift_knn_logits = knn_logits[:, :-1].contiguous() # (batch, seq_len-1, vocab_size) |
| 399 | |
| 400 | nonpad_mask = shift_labels != -100 |
| 401 | shift_logits = shift_logits[nonpad_mask] # (nonpad b*t, vocab_size) |
| 402 | shift_knn_logits = shift_knn_logits[nonpad_mask] # (nonpad b*t, vocab_size) |
| 403 | shift_labels = shift_labels[nonpad_mask] # (nonpad b*t) |
| 404 | |
| 405 | # Compute the entropy of the logits (shift_logits) and label_probs |
| 406 | lm_log_probs = F.log_softmax(shift_logits, dim=-1) |
| 407 | knn_log_probs = F.log_softmax(shift_knn_logits, dim=-1) |
| 408 | interpolated_log_probs = interpolate(knn_log_probs, lm_log_probs, lmbda=args.lmbda) |
| 409 | |
| 410 | lm_loss = F.nll_loss(lm_log_probs, shift_labels, reduction='sum') |
| 411 | joint_loss = F.nll_loss(interpolated_log_probs, shift_labels, reduction='sum') |
| 412 | |
| 413 | return joint_loss, lm_loss, shift_labels.shape[0] |
| 414 | |
| 415 | if args.do_test: |
| 416 | eval_dataloader = DataLoader( |