X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=quizz_machine.py;h=49e783586ebc662ac6bad80e34c2ef84169b1667;hb=bee6e628aabc1380772409f6aabffb024c0e70ab;hp=daa8a5456d02d0f27cffe42e52b5312a1fad3d75;hpb=ea4991c7875d0bbf3880403205b719e93e7ccd7b;p=culture.git diff --git a/quizz_machine.py b/quizz_machine.py index daa8a54..49e7835 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -17,6 +17,88 @@ from mygpt import BracketedSequence ###################################################################### +class Gang(nn.Module): + def __init__(self, models, nb_models_for_generation, mode="groupthink"): + super().__init__() + self.models = 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. + + +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, batch_size, @@ -27,7 +109,7 @@ def masked_inplace_autoregression( deterministic_synthesis, forbidden_tokens=None, logit_biases=None, - progress_bar_desc="autoregression", + progress_bar_desc=None, device=torch.device("cpu"), ): assert input.size() == ar_mask.size() @@ -51,7 +133,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, @@ -67,61 +150,34 @@ def masked_inplace_autoregression( ###################################################################### -class Task: - def batches(self, split="train", nb_to_use=-1, desc=None): - pass - - def vocabulary_size(self): - pass - - def produce_results( - self, n_epoch, model, result_dir, logger, deterministic_synthesis - ): - pass - - -###################################################################### - -import sky - - -class QuizzMachine(Task): - def save_image(self, input, result_dir, filename, logger): - img = sky.seq2img(input.to("cpu"), self.height, self.width) - image_name = os.path.join(result_dir, filename) - torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=6, padding=4) - logger(f"wrote {image_name}") - - def save_quizzes(self, input, result_dir, filename_prefix, logger): - self.save_image(input, result_dir, filename_prefix + ".png", logger) - +class QuizzMachine: def make_ar_mask(self, input): b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2 return b.long()[None, :].expand_as(input) def __init__( self, + problem, nb_train_samples, nb_test_samples, batch_size, - result_dir=None, - logger=None, + result_dir, + logger, device=torch.device("cpu"), ): super().__init__() + self.problem = problem self.batch_size = batch_size self.device = device - self.height = 6 - self.width = 8 + self.logger = logger - self.train_w_quizzes = sky.generate_seq( - nb_train_samples, height=self.height, width=self.width - ).to(device) - - self.test_w_quizzes = sky.generate_seq( - nb_test_samples, height=self.height, width=self.width + 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 @@ -129,8 +185,8 @@ class QuizzMachine(Task): self.test_c_quizzes = [] if result_dir is not None: - self.save_quizzes( - self.train_w_quizzes[:72], result_dir, f"culture_w_quizzes", logger + self.problem.save_quizzes( + self.train_w_quizzes[:72], result_dir, "culture_w_quizzes" ) def batches(self, split="train", desc=None): @@ -145,7 +201,7 @@ class QuizzMachine(Task): if len(c_quizzes) > 0: c_quizzes = torch.cat(c_quizzes, dim=0) if c_quizzes.size(0) > w_quizzes.size(0) // 2: - i = torch.randperm(w_quizzes.size(0))[: w_quizzes.size(0) // 2] + i = torch.randperm(c_quizzes.size(0))[: w_quizzes.size(0) // 2] c_quizzes = c_quizzes[i] i = torch.randperm(w_quizzes.size(0))[ @@ -176,9 +232,9 @@ class QuizzMachine(Task): 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) @@ -205,18 +261,18 @@ class QuizzMachine(Task): 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}") ############################## @@ -237,11 +293,8 @@ class QuizzMachine(Task): device=self.device, ) - self.save_quizzes( - result[:72], - result_dir, - f"culture_prediction_{n_epoch:04d}_{model.id:02d}", - logger, + self.problem.save_quizzes( + result[:72], result_dir, f"culture_prediction_{n_epoch:04d}_{model.id:02d}" ) return main_test_accuracy @@ -250,9 +303,7 @@ class QuizzMachine(Task): 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:] = sky.generate_seq(nb, height=self.height, width=self.width).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: @@ -260,26 +311,78 @@ class QuizzMachine(Task): else: self.test_c_quizzes.append(new_c_quizzes) - def create_c_quizzes( - self, - n_epoch, - result_dir, - logger, - nb, - model, - other_models, - min_ave_seq_logproba, - ): - ############################################################### - # Generate quizzes with model + def comput_correctness(self, c_quizzes, models_for_validation): + # Create the reverse 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) + reverse_c_quizzes = torch.cat( + [c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1 + ) + + ar_mask = self.make_ar_mask(c_quizzes) + seq_logproba = torch.empty(ar_mask.size(0), device=self.device) + + # Check how many of models can solve the quizzes in both directions + + nb_correct = 0 + + for model in models_for_validation: + result = c_quizzes.clone() + + masked_inplace_autoregression( + model=model, + batch_size=self.batch_size, + input=result, + ar_mask=ar_mask, + seq_logproba=seq_logproba, + temperature=1.0, + deterministic_synthesis=True, + # progress_bar_desc="solving c_quizzes", + device=self.device, + ) + + correct = (c_quizzes == result).long().min(dim=-1).values + + reverse_result = reverse_c_quizzes.clone() + + masked_inplace_autoregression( + model=model, + batch_size=self.batch_size, + input=reverse_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, + ) + + reverse_correct = ( + (reverse_c_quizzes == reverse_result).long().min(dim=-1).values + ) + + nb_correct += correct * reverse_correct + + return nb_correct + + ############################################################### + + def generate_quizzes(self, nb, model_for_generation, min_ave_seq_logproba): c_quizzes = torch.empty( - nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64 + 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) + # bracketing of the temperature to get the target logproba + temperature = 1 d_temperature = 1 / 3 @@ -287,96 +390,77 @@ class QuizzMachine(Task): seq_logproba[...] = 0 masked_inplace_autoregression( - model=model, + 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", + # progress_bar_desc="sampling c_quizzes", device=self.device, ) ave_seq_logproba = seq_logproba.mean() - logger(f"{ave_seq_logproba=} {min_ave_seq_logproba=}") - + # If we do not have target logprobs, get out now if min_ave_seq_logproba is None: break # Oh man that's ugly - if ave_seq_logproba < min_ave_seq_logproba * 1.1: + 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: + 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"chaging temperature to {temperature}") - - ############################################################### - # Create the reverse quizzes - - l = self.height * self.width - direction = c_quizzes[:, l : l + 1] - direction = sky.token_forward * ( - direction == sky.token_backward - ) + sky.token_backward * (direction == sky.token_forward) - reverse_c_quizzes = torch.cat( - [c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1 - ) - - 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 + self.logger(f"changing temperature to {temperature}") - nb_correct = [] + return c_quizzes, seq_logproba.mean() - for m in other_models: - result = c_quizzes.clone() + ###################################################################### - masked_inplace_autoregression( - model=m, - batch_size=self.batch_size, - input=result, - ar_mask=ar_mask, - seq_logproba=seq_logproba, - temperature=1.0, - deterministic_synthesis=True, - progress_bar_desc="solving c_quizzes", - device=self.device, - ) - - correct = (c_quizzes == result).long().min(dim=-1).values - - reverse_result = reverse_c_quizzes.clone() - - masked_inplace_autoregression( - model=m, - batch_size=self.batch_size, - input=reverse_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, - ) + def create_c_quizzes( + self, + nb, + model_for_generation, + models_for_validation, + min_ave_seq_logproba, + n_epoch, + result_dir, + ): + c_quizzes, ave_seq_logproba = self.generate_quizzes( + nb, model_for_generation, min_ave_seq_logproba + ) - reverse_correct = ( - (reverse_c_quizzes == reverse_result).long().min(dim=-1).values - ) + nb_correct = self.comput_correctness(c_quizzes, models_for_validation) - nb_correct.append((correct * reverse_correct)[None, :]) + return c_quizzes, nb_correct, ave_seq_logproba - nb_correct = torch.cat(nb_correct, dim=0).sum(dim=0) + ###################################################################### - return c_quizzes, nb_correct, seq_logproba.mean() + def gang_create_c_quizzes( + self, + nb, + nb_models_for_generation, + models, + mode, + min_ave_seq_logproba, + n_epoch, + result_dir, + ): + model_for_generation = Gang(models, nb_models_for_generation, mode) + models_for_validation = models + return self.create_c_quizzes( + nb, + model_for_generation, + models_for_validation, + min_ave_seq_logproba, + n_epoch, + result_dir, + )