+ with open("/tmp/result.dat", "w") as f:
+ for k in range(z.size(0)):
+ f.write(f"{y[k]} {z[k,0]} {z[k,1]}\n")
+
+ exit(0)
+
+######################################################################
+
+if args.test == "generator":
+ token_prolog_0 = vocabulary_size + 0
+ token_prolog_1 = vocabulary_size + 1
+ token_prolog_2 = vocabulary_size + 2
+ generator_vocabulary_size = vocabulary_size + 3
+
+ generator = mygpt.MyGPT(
+ vocabulary_size=generator_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,
+ compute_attzero=compute_causal_attzero,
+ dropout=args.dropout,
+ ).to(main_device)
+
+ generator.main_test_accuracy = 0.0
+
+ filename = f"generator.pth"
+
+ try:
+ d = torch.load(os.path.join(args.result_dir, filename))
+ generator.load_state_dict(d[0])
+ generator.main_test_accuracy = d[1]
+ log_string(f"successfully loaded {filename}")
+ except FileNotFoundError:
+ log_string(f"cannot find {filename}")
+ pass
+
+ for n_epoch in range(args.nb_epochs):
+ one_generator_epoch(
+ generator,
+ quiz_machine=quiz_machine,
+ models=models,
+ fraction_w_quizzes=1 if n_epoch < 25 else 0.5,
+ local_device=main_device,
+ )
+
+ filename = f"generator.pth"
+ torch.save(
+ (generator.state_dict(), generator.main_test_accuracy),
+ os.path.join(args.result_dir, filename),
+ )
+ log_string(f"wrote {filename}")
+
+ c_quizzes, prologs = generate_c_quizzes_with_generator(
+ generator, quiz_machine, args.batch_size
+ )
+
+ 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)
+ )