From d13eb42426139c4e59506db03f6b2bd68ade9b90 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Wed, 3 Jul 2024 12:29:19 +0300 Subject: [PATCH] Update. --- quizz_machine.py | 140 +++++++++++++++++++++++++++-------------------- sky.py | 71 ++++++++++++++++++------ 2 files changed, 136 insertions(+), 75 deletions(-) diff --git a/quizz_machine.py b/quizz_machine.py index 153317c..90f288e 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -105,10 +105,64 @@ def masked_inplace_autoregression( class QuizzMachine: - 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 indices_forward_and_backward(self, quizzes): + i_forward = quizzes[:, 0] == self.token_forward + j_forward = quizzes[:, 1 + self.prompt_len] == self.token_forward + i_backward = quizzes[:, 0] == self.token_backward + j_backward = quizzes[:, 1 + self.answer_len] == self.token_backward + assert torch.logical_or( + torch.logical_and(i_forward, j_forward), + torch.logical_and(i_backward, j_backward), + ).all() + return i_forward, i_backward + + def reverse_time(self, quizzes): + i_forward, i_backward = self.indices_forward_and_backward(quizzes) + + forward_to_backward = torch.cat( + [ + quizzes[:, 0:1], + quizzes[:, 2 + self.prompt_len :], + quizzes[:, 1 + self.prompt_len : 2 + self.prompt_len], + quizzes[:, 1 : 1 + self.prompt_len], + ], + dim=1, + ) + forward_to_backward[:, 0] = self.token_backward + forward_to_backward[:, 1 + self.answer_len] = self.token_backward + + backward_to_forward = torch.cat( + [ + quizzes[:, 0:1], + quizzes[:, 2 + self.answer_len :], + quizzes[:, 1 + self.answer_len : 2 + self.answer_len], + quizzes[:, 1 : 1 + self.answer_len], + ], + dim=1, + ) + + backward_to_forward[:, 0] = self.token_forward + backward_to_forward[:, 1 + self.prompt_len] = self.token_forward + + m = i_forward.long()[:, None] + + return m * forward_to_backward + (1 - m) * backward_to_forward + + def make_ar_mask(self, quizzes, first=False): + i_forward, i_backward = self.indices_forward_and_backward(quizzes) + + t = torch.arange(quizzes.size(1), device=quizzes.device) + + if first: + m_forward = (t >= 1).long() * (t < 1 + self.prompt_len).long() + m_backward = (t >= 1).long() * (t < 1 + self.answer_len).long() + else: + m_forward = (t >= 2 + self.prompt_len).long() + m_backward = (t >= 2 + self.answer_len).long() + + m = i_forward.long()[:, None] + + return m * m_forward + (1 - m) * m_backward def generate_token_sequences(self, nb): prompts, answers = self.problem.generate_prompts_and_answers(nb) @@ -181,19 +235,20 @@ class QuizzMachine: result_dir, "culture_w_quizzes", self.train_w_quizzes[:72] ) + # toto = self.reverse_time(self.train_w_quizzes[:72]) + # self.save_quizzes(result_dir, "toto", toto) + # exit(0) + def save_quizzes(self, result_dir, filename_prefix, quizzes, prediction=False): - l = (quizzes.size(1) - 1) // 2 - forward = (quizzes[:, 0] == self.token_forward).long() - backward = (quizzes[:, 0] == self.token_backward).long() - assert forward.equal(1 - backward) - first = quizzes[:, 1 : 1 + l] - second = quizzes[:, 1 + l : 1 + 2 * l] - prompts = forward[:, None] * first + backward[:, None] * second - answers = forward[:, None] * second + backward[:, None] * first + forward = quizzes[quizzes[:, 0] == self.token_forward] + ib = quizzes[:, 0] == self.token_backward + backward = quizzes[ib] + assert forward.size(0) + backward.size(0) == quizzes.size(0) + quizzes[ib] = self.reverse_time(quizzes[ib]) if prediction: - predicted_prompts = backward - predicted_answers = forward + predicted_prompts = ib + predicted_answers = torch.logical_not(ib) else: predicted_prompts = None predicted_answers = None @@ -201,8 +256,8 @@ class QuizzMachine: self.problem.save_quizzes( result_dir, filename_prefix, - prompts, - answers, + quizzes[:, 1 : 1 + self.prompt_len], + quizzes[:, 2 + self.prompt_len :], predicted_prompts, predicted_answers, ) @@ -254,7 +309,7 @@ class QuizzMachine: ): def compute_accuracy(input): input = input[:nmax] - ar_mask = self.make_ar_mask(input, 2 + self.prompt_len, self.answer_len) + ar_mask = self.make_ar_mask(input) result = input.clone() * (1 - ar_mask) seq_logproba = torch.empty(input.size(0), device=self.device) @@ -293,7 +348,7 @@ class QuizzMachine: ############################## input = self.test_w_quizzes[:96] - ar_mask = self.make_ar_mask(input, 2 + self.prompt_len, self.answer_len) + ar_mask = self.make_ar_mask(input) result = input.clone() * (1 - ar_mask) seq_logproba = torch.empty(input.size(0), device=self.device) @@ -330,32 +385,6 @@ class QuizzMachine: else: self.test_c_quizzes.append(new_c_quizzes) - 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( - [ - 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( self, c_quizzes, @@ -379,7 +408,7 @@ class QuizzMachine: seq_logproba[...] = 0.0 - ar_mask = self.make_ar_mask(result, 2 + self.prompt_len, self.answer_len) + ar_mask = self.make_ar_mask(result) masked_inplace_autoregression( model=model, @@ -398,9 +427,7 @@ class QuizzMachine: if bidirectional_validation: backward_result = backward_c_quizzes.clone() - ar_mask = self.make_ar_mask( - backward_result, 2 + self.answer_len, self.prompt_len - ) + ar_mask = self.make_ar_mask(backward_result) masked_inplace_autoregression( model=model, @@ -433,23 +460,18 @@ class QuizzMachine: nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64 ) - ar_mask_first = torch.zeros(c_quizzes.size(), device=self.device) - ar_mask_first[:, : ar_mask_first.size(1) // 2 + 1] = 1 - ar_mask_second = 1 - ar_mask_first - ar_mask_first[:, 0] = 0 - ar_mask_second[:, 0] = 0 - - seq_logproba = torch.zeros(ar_mask_first.size(0), device=self.device) + seq_logproba = torch.zeros(nb, device=self.device) # First, we generate the answer at high temperature c_quizzes[:, 0] = self.token_backward + c_quizzes[:, 1 + self.answer_len] = self.token_backward masked_inplace_autoregression( model=model_for_generation, batch_size=self.batch_size, input=c_quizzes, - ar_mask=ar_mask_first, + ar_mask=self.make_ar_mask(c_quizzes, first=True), seq_logproba=seq_logproba, temperature=temperature, deterministic_synthesis=False, @@ -462,7 +484,7 @@ class QuizzMachine: model=model_for_generation, batch_size=self.batch_size, input=c_quizzes, - ar_mask=ar_mask_second, + ar_mask=self.make_ar_mask(c_quizzes), seq_logproba=seq_logproba, temperature=1 / temperature, deterministic_synthesis=False, @@ -472,13 +494,13 @@ class QuizzMachine: # Then we return the quizz, and re-generate the response, now # at low temperature - c_quizzes = self.backward_to_forward(c_quizzes) + c_quizzes = self.reverse_time(c_quizzes) masked_inplace_autoregression( model=model_for_generation, batch_size=self.batch_size, input=c_quizzes, - ar_mask=ar_mask_second, + ar_mask=self.make_ar_mask(c_quizzes), seq_logproba=seq_logproba, temperature=1 / temperature, deterministic_synthesis=False, diff --git a/sky.py b/sky.py index 040ec67..e509fb7 100755 --- a/sky.py +++ b/sky.py @@ -5,7 +5,7 @@ # Written by Francois Fleuret -import math, sys, tqdm, os +import math, sys, tqdm, os, warnings import torch, torchvision @@ -200,38 +200,66 @@ class Sky(problem.Problem): if predicted_answers is None: predicted_answers = 255 - def add_frame(x, c, margin): - y = x.new_full( - (x.size(0), x.size(1), x.size(2) + 2 * margin, x.size(3) + 2 * margin), - 0, - ) + def add_frame(x, c, margin, bottom=False): + if bottom: + h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0 + else: + h, w, di, dj = ( + x.size(2) + 2 * margin, + x.size(3) + 2 * margin, + margin, + margin, + ) + + y = x.new_full((x.size(0), x.size(1), h, w), 0) + if type(c) is int: y[...] = c else: c = c.long()[:, None] - c = c * torch.tensor([192, 192, 192], device=c.device) + ( + c = c * torch.tensor([0, 0, 0], device=c.device) + ( 1 - c ) * torch.tensor([255, 255, 255], device=c.device) y[...] = c[:, :, None, None] - y[:, :, margin:-margin, margin:-margin] = x + + y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x + return y margin = 4 - img_prompts = add_frame(self.frame2img(prompts.to("cpu")), 0, 1) - img_answers = add_frame(self.frame2img(answers.to("cpu")), 0, 1) + img_prompts = add_frame(self.frame2img(prompts.to("cpu")), c=0, margin=1) + img_answers = add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1) - # img_prompts = add_frame(img_prompts, 255, margin) - # img_answers = add_frame(img_answers, 255, margin) + img_prompts = add_frame(img_prompts, c=255, margin=margin, bottom=True) + img_answers = add_frame(img_answers, c=255, margin=margin, bottom=True) - img_prompts = add_frame(img_prompts, predicted_prompts, margin) - img_answers = add_frame(img_answers, predicted_answers, margin) + img_prompts = add_frame( + img_prompts, c=predicted_prompts, margin=margin, bottom=True + ) + img_answers = add_frame( + img_answers, c=predicted_answers, margin=margin, bottom=True + ) + + marker_size = 8 separator = img_prompts.new_full( - (img_prompts.size(0), img_prompts.size(1), img_prompts.size(2), margin), 255 + ( + img_prompts.size(0), + img_prompts.size(1), + img_prompts.size(2), + marker_size, + ), + 255, ) - img = torch.cat([img_prompts, img_answers], dim=3) + for k in range(2, 2 * marker_size - 3): + i = k + 1 - marker_size + j = marker_size - 2 - abs(k - marker_size + 1) + separator[:, :, separator.size(2) // 2 + i, j] = 0 + separator[:, :, separator.size(2) // 2 + i + 1, j] = 0 + + img = torch.cat([img_prompts, separator, img_answers], dim=3) image_name = os.path.join(result_dir, filename) torchvision.utils.save_image( @@ -246,8 +274,19 @@ class Sky(problem.Problem): def generate_prompts_and_answers(self, nb): frame_sequences = self.generate_frame_sequences(nb) frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0) + prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1) + answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1) + # warnings.warn("dirty test with longer answer", RuntimeWarning) + # answers = torch.cat( + # [ + # frame_sequences[:, frame_sequences.size(1) // 2 :], + # frame_sequences[:, frame_sequences.size(1) // 2 :], + # ], + # dim=3, + # ).flatten(1) + return prompts, answers def save_quizzes( -- 2.39.5