-if args.task == "sandbox":
- if args.sandbox_level == 0:
- problem = tasks.ProblemLevel0(
- nb_sentences=args.sandbox_levels_nb_items,
- len_prompt=args.sandbox_levels_len_source,
- len_result=args.sandbox_levels_len_result,
- )
- elif args.sandbox_level == 1:
- problem = tasks.ProblemLevel1(
- nb_operators=args.sandbox_levels_nb_items,
- len_source=args.sandbox_levels_len_source,
- len_result=args.sandbox_levels_len_result,
- )
- elif args.sandbox_level == 2:
- problem = tasks.ProblemLevel2(
- len_source=args.sandbox_levels_len_source,
- len_result=args.sandbox_levels_len_result,
- )
- else:
- raise ValueError(f"Unknown sandbox level {args.sandbox_level}")
+if args.task == "byheart":
+ task = tasks.SandBox(
+ problem=problems.ProblemByHeart(),
+ nb_train_samples=args.nb_train_samples,
+ nb_test_samples=args.nb_test_samples,
+ batch_size=args.batch_size,
+ logger=log_string,
+ device=device,
+ )
+ args.max_percents_of_test_in_train = -1
+
+elif args.task == "learnop":
+ task = tasks.SandBox(
+ problem=problems.ProblemLearnOperator(),
+ nb_train_samples=args.nb_train_samples,
+ nb_test_samples=args.nb_test_samples,
+ batch_size=args.batch_size,
+ logger=log_string,
+ device=device,
+ )
+
+
+elif args.task == "guessop":
+ task = tasks.SandBox(
+ problem=problems.ProblemGuessOperator(),
+ nb_train_samples=args.nb_train_samples,
+ nb_test_samples=args.nb_test_samples,
+ batch_size=args.batch_size,
+ logger=log_string,
+ device=device,
+ )
+
+
+elif args.task == "twotargets":
+ task = tasks.SandBox(
+ problem=problems.ProblemTwoTargets(),
+ nb_train_samples=args.nb_train_samples,
+ nb_test_samples=args.nb_test_samples,
+ batch_size=args.batch_size,
+ logger=log_string,
+ device=device,
+ )