+ to_keep = new_quizzes[nb_correct == len(other_models) - 1]
+ log_string(f"keep {to_keep.size(0)} quizzes")
+ kept.append(to_keep)
+
+ new_quizzes = torch.cat(kept, dim=0)[: nb_for_train + nb_for_test]
+
+ task.store_new_quizzes(new_quizzes[:nb_for_train], for_train=True)
+ task.store_new_quizzes(new_quizzes[nb_for_train:], for_train=False)
+
+ task.save_image(
+ new_quizzes[:96],
+ args.result_dir,
+ f"world_quiz_{n_epoch:04d}_{model.id:02d}.png",
+ log_string,
+ )
+
+
+######################################################################
+
+models = []
+
+for k in range(args.nb_gpts):
+ model = mygpt.MyGPT(
+ vocabulary_size=vocabulary_size,
+ dim_model=args.dim_model,
+ dim_keys=args.dim_keys,
+ dim_hidden=args.dim_hidden,
+ nb_heads=args.nb_heads,
+ nb_blocks=args.nb_blocks,
+ causal=True,
+ dropout=args.dropout,
+ ).to(device)
+
+ model.main_test_accuracy = 0.0
+ model.id = k
+
+ models.append(model)
+
+
+nb_parameters = sum(p.numel() for p in models[0].parameters())
+log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)")
+
+######################################################################
+
+accuracy_to_make_quizzes = 0.975
+nb_new_quizzes_for_train = 1000
+nb_new_quizzes_for_test = 100
+
+if args.check:
+ accuracy_to_make_quizzes = 0.0
+ nb_new_quizzes_for_train = 10
+ nb_new_quizzes_for_test = 10
+
+for n_epoch in range(args.nb_epochs):
+ # select the model with lowest accuracy
+ models.sort(key=lambda model: model.main_test_accuracy)
+ model = models[0]
+
+ log_string(
+ f"training model {model.id} main_test_accuracy {model.main_test_accuracy}"
+ )
+
+ # improve it
+ one_epoch(model, task)
+
+ log_string(
+ f"train_set_composition world {task.nb_batch_samples_world} quizzes {task.nb_batch_samples_quizzes}"
+ )
+
+ # test it
+ run_tests(model, task, deterministic_synthesis=False)
+
+ if model.main_test_accuracy >= accuracy_to_make_quizzes:
+ other_models = models.copy()
+ other_models.remove(model)
+
+ create_quizzes(
+ model,
+ other_models,
+ task,
+ nb_for_train=nb_new_quizzes_for_train,
+ nb_for_test=nb_new_quizzes_for_test,
+ )