+c_quizzes_procedure = [
+ (("f_B", "f_A", "A", "B"), (1, 0, 0, 0), model_transformer_hot),
+ (("f_B", "f_A", "A", "B"), (0, 1, 1, 1), model_transformer_cold),
+ (("A", "f_A", "B", "f_B"), (0, 0, 0, 1), model_transformer_cold),
+ (("f_A", "A", "f_B", "B"), (0, 0, 0, 1), model_transformer_cold),
+]
+
+######################################################################
+
+
+def save_additional_results(model, models, science_w_quizzes):
+ # Save generated quizzes with the successive steps
+
+ recorder = []
+
+ c_quizzes = quiz_machine.generate_c_quizzes(
+ 64,
+ model_for_generation=model,
+ procedure=c_quizzes_procedure,
+ recorder=recorder,
+ )
+
+ # This is nb_quizzes x nb_models
+
+ seq_logproba = quiz_machine.models_logprobas(
+ models, c_quizzes, ("A", "f_A", "B", "f_B"), (0, 0, 0, 1), (0, 0, 1, 0)
+ ) + quiz_machine.models_logprobas(
+ models, c_quizzes, ("f_A", "A", "f_B", "B"), (0, 0, 0, 1), (0, 0, 1, 0)
+ )
+
+ probas = seq_logproba.exp()
+
+ comments = []
+
+ for l in seq_logproba:
+ comments.append("proba " + " ".join([f"{x.exp().item():.02f}" for x in l]))
+
+ ##
+
+ c_quizzes = torch.cat([c[:, None, :] for c, _, in recorder], dim=1)
+ predicted_parts = torch.cat([t[:, None, :] for _, t in recorder], dim=1)
+ nb_steps = c_quizzes.size(1)
+ c_quizzes = c_quizzes.reshape(-1, c_quizzes.size(-1))
+ predicted_parts = predicted_parts.reshape(-1, predicted_parts.size(-1))
+
+ # We have comments only for the final quiz, not the successive
+ # steps, so we have to add nb_steps-1 empty comments
+
+ steps_comments = []
+ for c in comments:
+ steps_comments += [""] * (nb_steps - 1) + [c]
+
+ filename = f"non_validated_{n_epoch:04d}_{model.id:02d}.png"
+
+ quiz_machine.problem.save_quizzes_as_image(
+ args.result_dir,
+ filename,
+ quizzes=c_quizzes,
+ predicted_parts=predicted_parts,
+ comments=steps_comments,
+ nrow=nb_steps * 2, # two quiz per row
+ )
+
+ log_string(f"wrote {filename}")
+
+ ######################################################################
+
+ if science_w_quizzes is not None:
+ struct = ("A", "f_A", "B", "f_B")
+ mask = (0, 0, 0, 1)
+ result, correct = quiz_machine.predict(
+ model=model,
+ quizzes=science_w_quizzes.to(main_device),
+ struct=struct,
+ mask=mask,
+ )
+
+ predicted_parts = torch.tensor(mask, device=correct.device)[None, :].expand(
+ correct.size(0), -1
+ )
+ correct = (2 * correct - 1) * (predicted_parts.sum(dim=-1) == 1).long()
+
+ nb_correct = (correct == 1).long().sum()
+ nb_total = (correct != 0).long().sum()
+
+ log_string(
+ f"science_accuracy {n_epoch} model {model.id} val {nb_correct} / {nb_total}"
+ )
+
+ i = correct == 1
+ j = correct != 1
+
+ result = torch.cat([result[i], result[j]], dim=0)
+ correct = torch.cat([correct[i], correct[j]], dim=0)
+ correct_parts = predicted_parts * correct[:, None]
+
+ result = result[:128]
+ predicted_parts = predicted_parts[:128]
+ correct_parts = correct_parts[:128]
+
+ quiz_machine.problem.save_quizzes_as_image(
+ args.result_dir,
+ f"culture_science_{n_epoch:04d}_{model.id:02d}.png",
+ quizzes=result,
+ predicted_parts=predicted_parts,
+ correct_parts=correct_parts,
+ )