+ if self.back_accuracy:
+ # If back_accuracy is True, we compute the accuracy on
+ # the backward quizzes not by counting how many time
+ # the real prompt A is equal to the reconstructed
+ # prompt A*, but how many time the answers B* computed
+ # from A* is equal to the correct answer. So we look
+ # for the accuracy of A->B*=B for the forward, but for
+ # the backward we look at B->A*->B*=B instead of B->A*=A
+
+ n_forward = input[:, 0] == self.token_forward
+ nb_total = input[n_forward].size(0)
+ nb_correct = (
+ (input[n_forward] == result[n_forward])
+ .long()
+ .min(dim=1)
+ .values.sum()
+ .item()
+ )
+
+ n_backward = input[:, 0] == self.token_backward
+ back_input = self.reverse_time(result[n_backward])
+
+ if back_input.size(0) > 0:
+ back_input[:, 2 + self.prompt_len :] = input[
+ n_backward, 1 : 1 + self.answer_len
+ ]
+ back_nb_total, back_nb_correct = compute_accuracy(back_input)
+
+ self.logger(
+ f"accuracy {n_epoch=} {model.id=} {nb_correct} / {nb_total}"
+ )
+ self.logger(
+ f"back_accuracy {n_epoch=} {model.id=} {back_nb_correct} / {back_nb_total}"
+ )
+
+ nb_total += back_nb_total
+ nb_correct += back_nb_correct
+ else:
+ self.logger(
+ f"accuracy {n_epoch=} {model.id=} {nb_correct} / {nb_total}"
+ )
+
+ else:
+ nb_total = input.size(0)
+ nb_correct = (input == result).long().min(dim=1).values.sum()