+ # 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()
+
+ nb_succeed = (probas >= args.proba_understands).long().sum(dim=1)
+ nb_fail = (probas <= args.proba_not_understands).long().sum(dim=1)
+
+ to_keep = (
+ (nb_succeed + nb_fail == probas.size(1))
+ & (nb_fail >= 1)
+ & (nb_fail <= args.max_fail_to_validate)
+ )
+
+ c_quizzes = c_quizzes[to_keep]
+
+ if c_quizzes.size(0) > 0:
+ nb_validated_per_model[model_for_generation.id] += c_quizzes.size(0)
+ recorded_validated.append(c_quizzes)
+ nb_validated = c_quizzes.size(0)
+ else:
+ nb_validated = 0
+
+ total_nb_validated = nb_validated_per_model.sum().item()
+
+ duration = time.perf_counter() - start_time
+
+ if total_nb_validated > 0:
+ if total_nb_validated < nb_to_validate:
+ d = (
+ (nb_to_validate - total_nb_validated)
+ * duration
+ / total_nb_validated
+ )
+ e = (datetime.datetime.now() + datetime.timedelta(seconds=d)).strftime(
+ "%a %H:%M"
+ )
+ else:
+ e = "now!"
+ else:
+ e = "???"
+
+ log_string(
+ f"keep c_quizzes model {model_for_generation.id} validated {nb_validated} / {nb_to_generate_per_iteration} ({100*nb_validated/nb_to_generate_per_iteration:.02f}%) nb_accumulated {total_nb_validated} / {nb_to_validate} (finishes {e} -- {int((total_nb_validated * 3600)/duration)}/h)"
+ )
+
+ validated_quizzes = torch.cat(recorded_validated, dim=0)
+
+ ######################################################################
+ # store the new c_quizzes which have been validated
+
+ v_train = validated_quizzes[:nb_for_train]
+ quiz_machine.store_c_quizzes(v_train, for_train=True)
+
+ v_test = validated_quizzes[nb_for_train:nb_to_validate]
+ quiz_machine.store_c_quizzes(v_test, for_train=False)
+
+ ######################################################################
+ # save images
+
+ vq = validated_quizzes[torch.randperm(validated_quizzes.size(0))[:128]]
+
+ if vq.size(0) > 0:
+ seq_logproba = quiz_machine.models_logprobas(
+ models, vq, ("A", "f_A", "B", "f_B"), (0, 0, 0, 1), (0, 0, 1, 0)
+ ) + quiz_machine.models_logprobas(
+ models, vq, ("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]))
+
+ filename = f"culture_c_quiz_{n_epoch:04d}.png"
+ quiz_machine.problem.save_quizzes_as_image(
+ args.result_dir, filename, vq, comments=comments
+ )
+
+
+######################################################################
+
+# The generator is very similar to a "solving GPT" except that it
+# deals with quizzes prologued with one token per solving GPT that
+# indicates if the said model solves it or not.
+#
+# There are three levels of solving 0->proba<=proba_not_understands,
+# 2->proba>=proba_understands and 1 otherwise.
+
+
+def generate_c_quizzes_with_generator(generator, quiz_machine, nb):
+ generator.to(main_device)
+
+ struct = ("A", "f_A", "B", "f_B")
+
+ c_quizzes = quiz_machine.problem.create_empty_quizzes(nb, struct=struct)
+ ar_mask = quiz_machine.make_quiz_mask(c_quizzes, struct, (1, 1, 1, 1))
+
+ i = F.one_hot(
+ torch.randint(args.nb_gpts, (c_quizzes.size(0),)),
+ num_classes=args.nb_gpts,
+ )
+
+ prologs_c_quizzes = token_prolog_0 * i + token_prolog_2 * (1 - i)
+ prologs_ar_mask = ar_mask.new_zeros(ar_mask.size(0), prologs_c_quizzes.size(1))
+
+ prologued_c_quizzes = torch.cat([prologs_c_quizzes, c_quizzes], dim=1).to(
+ main_device
+ )
+ prologued_ar_mask = torch.cat([prologs_ar_mask, ar_mask], dim=1).to(main_device)
+
+ seq_logproba = torch.zeros(
+ prologued_c_quizzes.size(0), device=prologued_c_quizzes.device
+ )
+
+ generator.temperature = args.temperature_hot
+
+ with torch.autograd.no_grad():
+ t = generator.training
+ generator.eval()
+
+ one_batch_masked_inplace_autoregression(
+ generator,
+ prologued_c_quizzes,
+ prologued_ar_mask,
+ seq_logproba,
+ deterministic_synthesis=False,
+ )
+
+ generator.train(t)
+
+ generator.reset_transformations()
+
+ prologued_c_quizzes = (
+ prologued_c_quizzes * (prologued_c_quizzes < vocabulary_size).long()
+ )
+
+ c_quizzes = prologued_c_quizzes[:, prologs_c_quizzes.size(1) :]
+
+ return c_quizzes.to("cpu"), prologs_c_quizzes.to("cpu")
+
+
+def batches_for_generator(generator, quiz_machine, models, fraction_w_quizzes=1.0):
+ samples = []
+
+ for _ in range(args.nb_train_samples // args.batch_size):
+ while sum([x.size(0) for x in samples]) < args.batch_size:
+ # Generate a bunch of quizzes
+
+ if torch.rand(1).item() <= fraction_w_quizzes:
+ # Either we start with the world quizzes
+ c_quizzes = quiz_machine.problem.generate_w_quizzes(
+ args.batch_size, progress_bar=False
+ )
+ else:
+ # Or we use the generator itself to generate them
+ c_quizzes, _ = generate_c_quizzes_with_generator(
+ generator, quiz_machine, args.batch_size
+ )
+
+ # We remove the trivial ones
+ to_keep = quiz_machine.problem.trivial(c_quizzes) == False
+ c_quizzes = c_quizzes[to_keep]
+
+ # If there are remaining ones, we compute the true prolog
+ # that indicates how the GPTs solve it
+
+ if c_quizzes.size(0) > 0:
+ 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()
+
+ u0 = probas <= args.proba_not_understands
+ u2 = probas >= args.proba_understands
+ u1 = (u0 | u2) == False
+
+ prologs = (
+ (u0.long() * token_prolog_0)
+ + (u1.long() * token_prolog_1)
+ + (u2.long() * token_prolog_2)
+ )
+
+ prologued_c_quizzes = torch.cat([prologs, c_quizzes], dim=1)
+
+ # nb_u2 = u2.long().sum(dim=1)
+ # nb_u0 = u0.long().sum(dim=1)
+ # prologued_c_quizzes = prologued_c_quizzes[(nb_u2 >= 1) & (nb_u0 >= 1)]
+
+ if prologued_c_quizzes.size(0) > 0:
+ samples.append(prologued_c_quizzes)
+
+ # Now we yield a batch
+
+ x = torch.cat(samples, dim=0)
+ samples = [x[args.batch_size :]]
+
+ yield x[: args.batch_size]
+
+
+def one_generator_epoch(
+ generator, quiz_machine, models, fraction_w_quizzes, local_device=main_device
+):
+ model.to(local_device).train()
+
+ optimizer = torch.optim.Adam(generator.parameters(), lr=args.learning_rate)
+
+ nb_train_samples, acc_train_loss = 0, 0.0
+
+ src = batches_for_generator(
+ generator=generator,
+ quiz_machine=quiz_machine,
+ models=models,
+ fraction_w_quizzes=fraction_w_quizzes,