X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=quizz_machine.py;h=198d279f3d8f9da8d0b602e622d3a8f690b907fd;hb=240870f5535bac35a08c552108d032854a8e2c38;hp=8ee022675322547ed1ceb4fbdde4e1fefbcc712f;hpb=07458603bfb24d5a12b530839e52e42fe8b0e6b8;p=culture.git diff --git a/quizz_machine.py b/quizz_machine.py index 8ee0226..198d279 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -12,10 +12,56 @@ import torch, torchvision from torch import nn from torch.nn import functional as F +import mygpt from mygpt import BracketedSequence ###################################################################### +# 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. + + +def one_batch_masked_inplace_autoregression( + model, + input, + ar_mask, + seq_logproba, + temperature=1.0, + deterministic_synthesis=False, + forbidden_tokens=None, + forced_biases=None, +): + to_generate = (ar_mask.sum(0) > 0).nonzero() + + if to_generate.min() > 0: + model( + BracketedSequence(input, 0, to_generate.min()) + ) # Needed to initialize the model's cache + for s in range(to_generate.min(), to_generate.max() + 1): + output = model(BracketedSequence(input, s, 1)).x + + logits = output[:, s] + + logits = (logits / temperature).log_softmax(dim=-1) + + if forbidden_tokens is not None: + logits = logits.masked_fill(forbidden_tokens, float("-inf")) + + if forced_biases is not None: + logits = logits + forced_biases[None, :] + + if deterministic_synthesis: + t_next = logits.argmax(-1) + else: + dist = torch.distributions.categorical.Categorical(logits=logits) + t_next = dist.sample() + + all_n = torch.arange(t_next.size(0)) + seq_logproba += logits[all_n, t_next].sum(dim=-1) + + input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s] + def masked_inplace_autoregression( model, @@ -51,7 +97,8 @@ def masked_inplace_autoregression( model.eval() for input, ar_mask, seq_logproba in batches: - model.masked_inplace_autoregression( + one_batch_masked_inplace_autoregression( + model=model, input=input, ar_mask=ar_mask, seq_logproba=seq_logproba, @@ -78,8 +125,8 @@ class QuizzMachine: nb_train_samples, nb_test_samples, batch_size, - result_dir=None, - logger=None, + result_dir, + logger, device=torch.device("cpu"), ): super().__init__() @@ -87,9 +134,14 @@ class QuizzMachine: self.problem = problem self.batch_size = batch_size self.device = device + self.logger = logger - self.train_w_quizzes = self.problem.generate_seq(nb_train_samples).to(device) - self.test_w_quizzes = self.problem.generate_seq(nb_test_samples).to(device) + 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( + device + ) self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1 @@ -144,9 +196,9 @@ class QuizzMachine: return self.nb_codes def produce_results( - self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 + self, n_epoch, model, result_dir, deterministic_synthesis, nmax=1000 ): - def compute_accuracy(input, logger=None): + def compute_accuracy(input): input = input[:nmax] ar_mask = self.make_ar_mask(input) result = input.clone() * (1 - ar_mask) @@ -173,18 +225,18 @@ class QuizzMachine: train_nb_total, train_nb_correct = compute_accuracy(self.train_w_quizzes) - logger( + self.logger( f"accuracy_train {n_epoch} nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%" ) - test_nb_total, test_nb_correct = compute_accuracy(self.test_w_quizzes, logger) + test_nb_total, test_nb_correct = compute_accuracy(self.test_w_quizzes) - logger( + self.logger( f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%" ) main_test_accuracy = test_nb_correct / test_nb_total - logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}") + self.logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}") ############################## @@ -215,7 +267,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_seq(nb).to(self.device) + input[-nb:] = self.problem.generate_token_sequences(nb).to(self.device) def store_c_quizzes(self, new_c_quizzes, for_train=True): if for_train: @@ -223,66 +275,7 @@ class QuizzMachine: else: self.test_c_quizzes.append(new_c_quizzes) - def create_c_quizzes( - self, - nb, - model_for_generation, - models_for_validation, - min_ave_seq_logproba, - n_epoch, - result_dir, - logger, - ): - ############################################################### - # Generate quizzes with model - - c_quizzes = torch.empty( - nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64 - ) - - ar_mask = torch.full(c_quizzes.size(), 1, device=self.device) - seq_logproba = torch.empty(ar_mask.size(0), device=self.device) - - temperature = 1 - d_temperature = 1 / 3 - - while True: - seq_logproba[...] = 0 - - masked_inplace_autoregression( - model=model_for_generation, - batch_size=self.batch_size, - input=c_quizzes, - ar_mask=ar_mask, - seq_logproba=seq_logproba, - temperature=temperature, - deterministic_synthesis=False, - # progress_bar_desc="sampling c_quizzes", - device=self.device, - ) - - ave_seq_logproba = seq_logproba.mean() - - if min_ave_seq_logproba is None: - break - - # Oh man that's ugly - if ave_seq_logproba < min_ave_seq_logproba: - if d_temperature > 0: - d_temperature *= -1 / 3 - temperature += d_temperature - elif ave_seq_logproba > min_ave_seq_logproba * 0.99: - if d_temperature < 0: - d_temperature *= -1 / 3 - temperature += d_temperature - else: - break - - logger(f"changing temperature to {temperature}") - - ############################################################### - # Create the reverse quizzes - + def reverse_time(self, c_quizzes): token_forward, token_backward = self.problem.direction_tokens() l = (c_quizzes.size(1) - 1) // 2 @@ -290,18 +283,20 @@ class QuizzMachine: direction = self.problem.token_forward * ( direction == self.problem.token_backward ) + self.problem.token_backward * (direction == self.problem.token_forward) - reverse_c_quizzes = torch.cat( - [c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1 - ) + + return torch.cat([c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1) + + def compute_correctness( + self, c_quizzes, models_for_validation, both_direction=True + ): + reversed_c_quizzes = self.reverse_time(c_quizzes) ar_mask = self.make_ar_mask(c_quizzes) seq_logproba = torch.empty(ar_mask.size(0), device=self.device) - ############################################################### - # Check how many of the other models can solve them in both - # directions + # Check how many of models can solve the quizzes in both directions - nb_correct = [] + nb_correct = 0 for model in models_for_validation: result = c_quizzes.clone() @@ -320,26 +315,106 @@ class QuizzMachine: correct = (c_quizzes == result).long().min(dim=-1).values - reverse_result = reverse_c_quizzes.clone() + if both_direction: + 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, + ) + + reversed_correct = ( + (reversed_c_quizzes == reversed_result).long().min(dim=-1).values + ) + + correct *= reversed_correct + + # endif + + nb_correct += correct + + return nb_correct + + ############################################################### + + def generate_quizzes(self, nb, model_for_generation, reverse_cleanup=False): + 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) + + if reverse_cleanup: + warnings.warn("very high temperature with reversed cleanup", RuntimeWarning) + temperature = 10.0 + else: + temperature = 1.0 + + # warnings.warn("noise injection", RuntimeWarning) + # noise_std = torch.rand(1).item() + # self.logger(f"{noise_std=}") + + # mygpt.set_noise_injection(model_for_generation, noise_std) + + masked_inplace_autoregression( + model=model_for_generation, + batch_size=self.batch_size, + input=c_quizzes, + ar_mask=ar_mask_prompt, + seq_logproba=seq_logproba, + temperature=temperature, + deterministic_synthesis=False, + device=self.device, + ) + + # mygpt.set_noise_injection(model_for_generation, 0.0) + ave_seq_logproba = seq_logproba.mean() + + 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, + device=self.device, + ) + + if reverse_cleanup: + c_quizzes = self.reverse_time(c_quizzes) masked_inplace_autoregression( - model=model, + model=model_for_generation, batch_size=self.batch_size, - input=reverse_result, - ar_mask=ar_mask, + input=c_quizzes, + ar_mask=ar_mask_solve, seq_logproba=seq_logproba, - temperature=1.0, + temperature=temperature, deterministic_synthesis=True, - # progress_bar_desc="solving reversed c_quizzes", device=self.device, ) - reverse_correct = ( - (reverse_c_quizzes == reverse_result).long().min(dim=-1).values + 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, + device=self.device, ) - nb_correct.append((correct * reverse_correct)[None, :]) - - nb_correct = torch.cat(nb_correct, dim=0).sum(dim=0) - - return c_quizzes, nb_correct, seq_logproba.mean() + return c_quizzes, seq_logproba.mean()