X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=quizz_machine.py;h=0d6d8f57cba918235f951333a64cb1a4c44133d2;hb=d283cd3d46a6323fec4c6a0970ac71e553e4a486;hp=6f7492de907f7257888a46a93dad6f37bbdda964;hpb=60bf08d4197f2dd3a58bd900401c11d47225b0df;p=culture.git diff --git a/quizz_machine.py b/quizz_machine.py index 6f7492d..0d6d8f5 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -17,43 +17,6 @@ from mygpt import BracketedSequence ###################################################################### - -class Gang(nn.Module): - def __init__(self, models, nb_models_for_generation, mode="groupthink"): - super().__init__() - self.models = nn.ModuleList(models) - self.nb_models_for_generation = nb_models_for_generation - self.mode = mode - - def forward(self, bs): - # If first = 0, we are re-starting an auto-regressive process, - # that's the right moment to randomize who gonna do it - if bs.first == 0: - self.models_to_use = [ - self.models[k] - for k in torch.randperm(len(self.models))[ - : self.nb_models_for_generation - ] - ] - - all_the_logits = torch.cat( - [model(bs).x[None] for model in self.models_to_use], dim=0 - ) - - if self.mode == "groupthink": - y = all_the_logits.mean(dim=0) - elif self.mode == "groupwork": - m = torch.rand(all_the_logits.size(), device=all_the_logits.device) - m = (m.sort(dim=0).indices == 0).long() - y = (y * m).sum(dim=0) - else: - raise ValueError(f"Invalid mode {self.mode}") - - return BracketedSequence(y, bs.first, bs.nb) - - -###################################################################### - # ar_mask is a tensor with 0s and 1s, of same shape as input, with # 1s where tokens should be generated. The others are kept # unchanged. @@ -156,6 +119,20 @@ class QuizzMachine: b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2 return b.long()[None, :].expand_as(input) + def generate_token_sequences(self, nb): + prompts, answers = self.problem.generate_prompts_and_answers(nb) + result = [] + + for prompt, answer in zip(prompts, answers): + if torch.rand(1) < 0.5: + a = [torch.tensor([self.token_forward]), prompt, answer] + else: + a = [torch.tensor([self.token_backward]), answer, prompt] + + result.append(torch.cat(a, dim=0)[None, :]) + + return torch.cat(result, dim=0) + def __init__( self, problem, @@ -168,28 +145,57 @@ class QuizzMachine: ): super().__init__() + v = problem.nb_token_values() + self.token_forward = v + self.token_backward = v + 1 + self.nb_token_values = v + 2 + self.problem = problem self.batch_size = batch_size self.device = device self.logger = logger - self.train_w_quizzes = self.problem.generate_token_sequences( - nb_train_samples - ).to(device) - self.test_w_quizzes = self.problem.generate_token_sequences(nb_test_samples).to( + self.train_w_quizzes = self.generate_token_sequences(nb_train_samples).to( device ) - self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1 + self.test_w_quizzes = self.generate_token_sequences(nb_test_samples).to(device) self.train_c_quizzes = [] self.test_c_quizzes = [] if result_dir is not None: - self.problem.save_quizzes( - self.train_w_quizzes[:72], result_dir, "culture_w_quizzes" + self.save_quizzes( + result_dir, "culture_w_quizzes", self.train_w_quizzes[:72] ) + def save_quizzes(self, result_dir, filename_prefix, quizzes, prediction=False): + print(f"DEBUG {quizzes.size()=}") + 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 + + if prediction: + predicted_prompts = backward + predicted_answers = forward + else: + predicted_prompts = None + predicted_answers = None + + self.problem.save_quizzes( + result_dir, + filename_prefix, + prompts, + answers, + predicted_prompts, + predicted_answers, + ) + def batches(self, split="train", desc=None): assert split in {"train", "test"} if split == "train": @@ -230,7 +236,7 @@ class QuizzMachine: yield batch def vocabulary_size(self): - return self.nb_codes + return self.nb_token_values def produce_results( self, n_epoch, model, result_dir, deterministic_synthesis, nmax=1000 @@ -294,8 +300,11 @@ class QuizzMachine: device=self.device, ) - self.problem.save_quizzes( - result[:72], result_dir, f"culture_prediction_{n_epoch:04d}_{model.id:02d}" + self.save_quizzes( + result_dir, + f"culture_prediction_{n_epoch:04d}_{model.id:02d}", + quizzes=result[:72], + prediction=True, ) return main_test_accuracy @@ -304,7 +313,7 @@ class QuizzMachine: input = self.train_w_quizzes if for_train else self.test_w_quizzes nb = min(nb, input.size(0)) input[:-nb] = input[nb:].clone() - input[-nb:] = self.problem.generate_token_sequences(nb).to(self.device) + input[-nb:] = self.generate_token_sequences(nb).to(self.device) def store_c_quizzes(self, new_c_quizzes, for_train=True): if for_train: @@ -313,17 +322,19 @@ class QuizzMachine: self.test_c_quizzes.append(new_c_quizzes) def reverse_time(self, c_quizzes): - token_forward, token_backward = self.problem.direction_tokens() - l = (c_quizzes.size(1) - 1) // 2 - direction = c_quizzes[:, l : l + 1] - direction = self.problem.token_forward * ( - direction == self.problem.token_backward - ) + self.problem.token_backward * (direction == self.problem.token_forward) + direction = c_quizzes[:, 0:1] + direction = self.token_forward * ( + direction == self.token_backward + ) + self.token_backward * (direction == self.token_forward) - return torch.cat([c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1) + return torch.cat( + [direction, c_quizzes[:, l + 1 :], c_quizzes[:, 1 : l + 1]], dim=1 + ) - def comput_correctness(self, c_quizzes, models_for_validation): + def compute_correctness( + self, c_quizzes, models_for_validation, both_directions=False + ): reversed_c_quizzes = self.reverse_time(c_quizzes) ar_mask = self.make_ar_mask(c_quizzes) @@ -350,142 +361,94 @@ class QuizzMachine: correct = (c_quizzes == result).long().min(dim=-1).values - reversed_result = reversed_c_quizzes.clone() + if both_directions: + reversed_result = reversed_c_quizzes.clone() - masked_inplace_autoregression( - model=model, - batch_size=self.batch_size, - input=reversed_result, - ar_mask=ar_mask, - seq_logproba=seq_logproba, - temperature=1.0, - deterministic_synthesis=True, - # progress_bar_desc="solving reversed c_quizzes", - device=self.device, - ) + masked_inplace_autoregression( + model=model, + batch_size=self.batch_size, + input=reversed_result, + ar_mask=ar_mask, + seq_logproba=seq_logproba, + temperature=1.0, + deterministic_synthesis=True, + # progress_bar_desc="solving reversed c_quizzes", + device=self.device, + ) - reversed_correct = ( - (reversed_c_quizzes == reversed_result).long().min(dim=-1).values - ) + reversed_correct = ( + (reversed_c_quizzes == reversed_result).long().min(dim=-1).values + ) + + correct *= reversed_correct - nb_correct += correct * reversed_correct + # endif + + nb_correct += correct return nb_correct ############################################################### - def generate_quizzes( - self, nb, model_for_generation, min_ave_seq_logproba, reverse_cleanup=False - ): + def generate_quizzes(self, nb, model_for_generation): c_quizzes = torch.empty( nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64 ) - ar_mask_prompt = torch.zeros(c_quizzes.size(), device=self.device) - ar_mask_prompt[:, : ar_mask_prompt.size(1) // 2 + 1] = 1 - ar_mask_solve = 1 - ar_mask_prompt - seq_logproba = torch.empty(ar_mask_prompt.size(0), device=self.device) + 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 - if reverse_cleanup: - warnings.warn("very high temperature with reversed cleanup", RuntimeWarning) - temperature = 10.0 - else: - temperature = 1.0 + seq_logproba = torch.empty(ar_mask_first.size(0), device=self.device) - # warnings.warn("noise injection", RuntimeWarning) - # noise_std = torch.rand(1).item() - # self.logger(f"{noise_std=}") + temperature = 10.0 - # mygpt.set_noise_injection(model_for_generation, noise_std) + # First, we generate the answer at high temperature + + c_quizzes[:, 0] = self.token_backward masked_inplace_autoregression( model=model_for_generation, batch_size=self.batch_size, input=c_quizzes, - ar_mask=ar_mask_prompt, + ar_mask=ar_mask_first, seq_logproba=seq_logproba, temperature=temperature, deterministic_synthesis=False, - # progress_bar_desc="sampling c_quizzes", device=self.device, ) - # mygpt.set_noise_injection(model_for_generation, 0.0) - ave_seq_logproba = seq_logproba.mean() + # Then, we generate the prompt deterministically + masked_inplace_autoregression( model=model_for_generation, batch_size=self.batch_size, input=c_quizzes, - ar_mask=ar_mask_solve, + ar_mask=ar_mask_second, seq_logproba=seq_logproba, - temperature=temperature, + temperature=1.0, deterministic_synthesis=True, - # progress_bar_desc="sampling c_quizzes", device=self.device, ) - if reverse_cleanup: - 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_solve, - seq_logproba=seq_logproba, - temperature=temperature, - deterministic_synthesis=True, - # progress_bar_desc="sampling c_quizzes", - device=self.device, - ) - - return c_quizzes, seq_logproba.mean() + # Then we return the quizz, and re-generate the response, now + # deterministically - ###################################################################### + c_quizzes = self.reverse_time(c_quizzes) - def create_c_quizzes( - self, - nb, - model_for_generation, - models_for_validation, - min_ave_seq_logproba, - reverse_cleanup, - n_epoch, - result_dir, - ): - c_quizzes, ave_seq_logproba = self.generate_quizzes( - nb, - model_for_generation=model_for_generation, - min_ave_seq_logproba=min_ave_seq_logproba, - reverse_cleanup=reverse_cleanup, + masked_inplace_autoregression( + model=model_for_generation, + batch_size=self.batch_size, + input=c_quizzes, + ar_mask=ar_mask_second, + seq_logproba=seq_logproba, + temperature=temperature, + deterministic_synthesis=True, + device=self.device, ) - nb_correct = self.comput_correctness(c_quizzes, models_for_validation) - - return c_quizzes, nb_correct, ave_seq_logproba - - ###################################################################### - - def gang_create_c_quizzes( - self, - nb, - nb_models_for_generation, - models, - mode, - min_ave_seq_logproba, - reverse_cleanup, - n_epoch, - result_dir, - ): - model_for_generation = Gang(models, nb_models_for_generation, mode) - models_for_validation = models - return self.create_c_quizzes( - nb=nb, - model_for_generation=model_for_generation, - models_for_validation=models_for_validation, - min_ave_seq_logproba=min_ave_seq_logproba, - reverse_cleanup=reverse_cleanup, - n_epoch=n_epoch, - result_dir=result_dir, - ) + return c_quizzes, seq_logproba.mean()