From 15e704a200286551a8e9c1765a0340c370367dee Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Wed, 3 Jul 2024 00:45:15 +0300 Subject: [PATCH] Update. --- main.py | 12 +++--- quizz_machine.py | 95 ++++++++++++++++++++++++++++++++++-------------- 2 files changed, 73 insertions(+), 34 deletions(-) diff --git a/main.py b/main.py index 5484f39..a8a6191 100755 --- a/main.py +++ b/main.py @@ -79,8 +79,6 @@ parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--deterministic_synthesis", action="store_true", default=False) -parser.add_argument("--bidirectional_validation", action="store_true", default=False) - parser.add_argument("--problem", type=str, default="sky") parser.add_argument("--nb_gpts", type=int, default=5) @@ -91,12 +89,14 @@ parser.add_argument("--max_to_validate", type=int, default=None) parser.add_argument("--accuracy_to_make_c_quizzes", type=float, default=0.975) -parser.add_argument("--dirty_debug", action="store_true", default=False) - parser.add_argument("--generation_temperature", type=float, default=2.0) parser.add_argument("--deterministic_validation", action="store_true", default=False) +parser.add_argument("--bidirectional_validation", action="store_true", default=False) + +parser.add_argument("--dirty_debug", action="store_true", default=False) + ###################################################################### parser.add_argument("--sky_height", type=int, default=6) @@ -114,10 +114,10 @@ parser.add_argument("--sky_speed", type=int, default=3) args = parser.parse_args() if args.min_to_validate is None: - args.min_to_validate = args = nb_gpts - 1 + args.min_to_validate = args.nb_gpts - 1 if args.max_to_validate is None: - args.max_to_validate = args = nb_gpts - 1 + args.max_to_validate = args.nb_gpts - 1 if args.result_dir is None: args.result_dir = f"results_culture" diff --git a/quizz_machine.py b/quizz_machine.py index de85520..153317c 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -105,19 +105,39 @@ def masked_inplace_autoregression( class QuizzMachine: - def make_ar_mask(self, input): - b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2 + def make_ar_mask(self, input, first, nb): + i = torch.arange(input.size(1), device=input.device) + b = torch.logical_and(i >= first, i < first + nb) return b.long()[None, :].expand_as(input) def generate_token_sequences(self, nb): prompts, answers = self.problem.generate_prompts_and_answers(nb) + + if self.prompt_len is None: + self.prompt_len = prompts.size(1) + + if self.answer_len is None: + self.answer_len = answers.size(1) + + assert prompts.size(1) == self.prompt_len and answers.size(1) == self.answer_len + result = [] for prompt, answer in zip(prompts, answers): if torch.rand(1) < 0.5: - a = [torch.tensor([self.token_forward]), prompt, answer] + a = [ + torch.tensor([self.token_forward]), + prompt, + torch.tensor([self.token_forward]), + answer, + ] else: - a = [torch.tensor([self.token_backward]), answer, prompt] + a = [ + torch.tensor([self.token_backward]), + answer, + torch.tensor([self.token_backward]), + prompt, + ] result.append(torch.cat(a, dim=0)[None, :]) @@ -144,6 +164,8 @@ class QuizzMachine: self.batch_size = batch_size self.device = device self.logger = logger + self.prompt_len = None + self.answer_len = None self.train_w_quizzes = self.generate_token_sequences(nb_train_samples).to( device @@ -232,7 +254,7 @@ class QuizzMachine: ): def compute_accuracy(input): input = input[:nmax] - ar_mask = self.make_ar_mask(input) + ar_mask = self.make_ar_mask(input, 2 + self.prompt_len, self.answer_len) result = input.clone() * (1 - ar_mask) seq_logproba = torch.empty(input.size(0), device=self.device) @@ -248,10 +270,8 @@ class QuizzMachine: device=self.device, ) - nb_total, nb_correct = ( - input.size(0), - (input == result).long().min(dim=1).values.sum(), - ) + nb_total = input.size(0) + nb_correct = (input == result).long().min(dim=1).values.sum() return nb_total, nb_correct @@ -273,7 +293,7 @@ class QuizzMachine: ############################## input = self.test_w_quizzes[:96] - ar_mask = self.make_ar_mask(input) + ar_mask = self.make_ar_mask(input, 2 + self.prompt_len, self.answer_len) result = input.clone() * (1 - ar_mask) seq_logproba = torch.empty(input.size(0), device=self.device) @@ -310,15 +330,30 @@ class QuizzMachine: else: self.test_c_quizzes.append(new_c_quizzes) - def reverse_time(self, c_quizzes): - l = (c_quizzes.size(1) - 1) // 2 - direction = c_quizzes[:, 0:1] - direction = self.token_forward * ( - direction == self.token_backward - ) + self.token_backward * (direction == self.token_forward) + def forward_to_backward(self, c_quizzes): + prompts = c_quizzes[:, 1 : 1 + self.prompt_len] + answers = c_quizzes[:, 2 + self.prompt_len :] + return torch.cat( + [ + c_quizzes.new_full((c_quizzes, 1), self.token_backward), + answers, + c_quizzes.new_full((c_quizzes, 1), self.token_backward), + prompts, + ], + dim=1, + ) + def backward_to_forward(self, c_quizzes): + answers = c_quizzes[:, 1 : 1 + self.answer_len :] + prompts = c_quizzes[:, 2 + self.answer_len :] return torch.cat( - [direction, c_quizzes[:, l + 1 :], c_quizzes[:, 1 : l + 1]], dim=1 + [ + c_quizzes.new_full((c_quizzes.size(0), 1), self.token_forward), + prompts, + c_quizzes.new_full((c_quizzes.size(0), 1), self.token_forward), + answers, + ], + dim=1, ) def compute_correctness( @@ -328,17 +363,15 @@ class QuizzMachine: bidirectional_validation=False, deterministic_validation=True, ): - reversed_c_quizzes = self.reverse_time(c_quizzes) + if bidirectional_validation: + backward_c_quizzes = self.forward_to_backward(c_quizzes) - ar_mask = self.make_ar_mask(c_quizzes) seq_logproba = torch.zeros( c_quizzes.size(0), max([m.id for m in models_for_validation]) + 1, device=self.device, ) - # Check how many of models can solve the quizzes in both directions - nb_correct = 0 for model in models_for_validation: @@ -346,6 +379,8 @@ class QuizzMachine: seq_logproba[...] = 0.0 + ar_mask = self.make_ar_mask(result, 2 + self.prompt_len, self.answer_len) + masked_inplace_autoregression( model=model, batch_size=self.batch_size, @@ -361,25 +396,29 @@ class QuizzMachine: correct = (c_quizzes == result).long().min(dim=-1).values if bidirectional_validation: - reversed_result = reversed_c_quizzes.clone() + backward_result = backward_c_quizzes.clone() + + ar_mask = self.make_ar_mask( + backward_result, 2 + self.answer_len, self.prompt_len + ) masked_inplace_autoregression( model=model, batch_size=self.batch_size, - input=reversed_result, + input=backward_result, ar_mask=ar_mask, seq_logproba=seq_logproba[:, model.id], temperature=1.0, deterministic_synthesis=deterministic_validation, - # progress_bar_desc="solving reversed c_quizzes", + # progress_bar_desc="solving backward c_quizzes", device=self.device, ) - reversed_correct = ( - (reversed_c_quizzes == reversed_result).long().min(dim=-1).values + backward_correct = ( + (backward_c_quizzes == backward_result).long().min(dim=-1).values ) - correct *= reversed_correct + correct *= backward_correct # endif @@ -433,7 +472,7 @@ class QuizzMachine: # Then we return the quizz, and re-generate the response, now # at low temperature - c_quizzes = self.reverse_time(c_quizzes) + c_quizzes = self.backward_to_forward(c_quizzes) masked_inplace_autoregression( model=model_for_generation, -- 2.20.1