quizz_machine.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math, os, tqdm, warnings
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 from mygpt import BracketedSequence
  16
  17 ######################################################################
  18
  19
  20 class Gang(nn.Module):
  21     def __init__(self, models, nb_models_for_generation, mode="groupthink"):
  22         super().__init__()
  23         self.models = models
  24         self.nb_models_for_generation = nb_models_for_generation
  25         self.mode = mode
  26
  27     def forward(self, bs):
  28         # If first = 0, we are re-starting an auto-regressive process,
  29         # that's the right moment to randomize who gonna do it
  30         if bs.first == 0:
  31             self.models_to_use = [
  32                 self.models[k]
  33                 for k in torch.randperm(len(self.models))[
  34                     : self.nb_models_for_generation
  35                 ]
  36             ]
  37
  38         all_the_logits = torch.cat(
  39             [model(bs).x[None] for model in self.models_to_use], dim=0
  40         )
  41
  42         if self.mode == "groupthink":
  43             y = all_the_logits.mean(dim=0)
  44         elif self.mode == "groupwork":
  45             m = torch.rand(all_the_logits.size(), device=all_the_logits.device)
  46             m = (m.sort(dim=0).indices == 0).long()
  47             y = (y * m).sum(dim=0)
  48         else:
  49             raise ValueError(f"Invalid mode {self.mode}")
  50
  51         return BracketedSequence(y, bs.first, bs.nb)
  52
  53
  54 ######################################################################
  55
  56 # ar_mask is a tensor with 0s and 1s, of same shape as input, with
  57 # 1s where tokens should be generated. The others are kept
  58 # unchanged.
  59
  60
  61 def one_batch_masked_inplace_autoregression(
  62     model,
  63     input,
  64     ar_mask,
  65     seq_logproba,
  66     temperature=1.0,
  67     deterministic_synthesis=False,
  68     forbidden_tokens=None,
  69     forced_biases=None,
  70 ):
  71     to_generate = (ar_mask.sum(0) > 0).nonzero()
  72
  73     if to_generate.min() > 0:
  74         model(
  75             BracketedSequence(input, 0, to_generate.min())
  76         )  # Needed to initialize the model's cache
  77     for s in range(to_generate.min(), to_generate.max() + 1):
  78         output = model(BracketedSequence(input, s, 1)).x
  79
  80         logits = output[:, s]
  81
  82         logits = (logits / temperature).log_softmax(dim=-1)
  83
  84         if forbidden_tokens is not None:
  85             logits = logits.masked_fill(forbidden_tokens, float("-inf"))
  86
  87         if forced_biases is not None:
  88             logits = logits + forced_biases[None, :]
  89
  90         if deterministic_synthesis:
  91             t_next = logits.argmax(-1)
  92         else:
  93             dist = torch.distributions.categorical.Categorical(logits=logits)
  94             t_next = dist.sample()
  95
  96         all_n = torch.arange(t_next.size(0))
  97         seq_logproba += logits[all_n, t_next].sum(dim=-1)
  98
  99         input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s]
 100
 101
 102 def masked_inplace_autoregression(
 103     model,
 104     batch_size,
 105     input,
 106     ar_mask,
 107     seq_logproba,
 108     temperature,
 109     deterministic_synthesis,
 110     forbidden_tokens=None,
 111     logit_biases=None,
 112     progress_bar_desc=None,
 113     device=torch.device("cpu"),
 114 ):
 115     assert input.size() == ar_mask.size()
 116
 117     batches = zip(
 118         input.split(batch_size),
 119         ar_mask.split(batch_size),
 120         seq_logproba.split(batch_size),
 121     )
 122
 123     if progress_bar_desc is not None:
 124         batches = tqdm.tqdm(
 125             batches,
 126             dynamic_ncols=True,
 127             desc=progress_bar_desc,
 128             total=(input.size(0) + batch_size - 1) // batch_size,
 129         )
 130
 131     with torch.autograd.no_grad():
 132         t = model.training
 133         model.eval()
 134
 135         for input, ar_mask, seq_logproba in batches:
 136             one_batch_masked_inplace_autoregression(
 137                 model=model,
 138                 input=input,
 139                 ar_mask=ar_mask,
 140                 seq_logproba=seq_logproba,
 141                 temperature=temperature,
 142                 deterministic_synthesis=deterministic_synthesis,
 143                 forbidden_tokens=forbidden_tokens,
 144                 forced_biases=logit_biases,
 145             )
 146
 147         model.train(t)
 148
 149
 150 ######################################################################
 151
 152
 153 class QuizzMachine:
 154     def make_ar_mask(self, input):
 155         b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2
 156         return b.long()[None, :].expand_as(input)
 157
 158     def __init__(
 159         self,
 160         problem,
 161         nb_train_samples,
 162         nb_test_samples,
 163         batch_size,
 164         result_dir=None,
 165         logger=None,
 166         device=torch.device("cpu"),
 167     ):
 168         super().__init__()
 169
 170         self.problem = problem
 171         self.batch_size = batch_size
 172         self.device = device
 173
 174         self.train_w_quizzes = self.problem.generate_token_sequences(
 175             nb_train_samples
 176         ).to(device)
 177         self.test_w_quizzes = self.problem.generate_token_sequences(nb_test_samples).to(
 178             device
 179         )
 180
 181         self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1
 182
 183         self.train_c_quizzes = []
 184         self.test_c_quizzes = []
 185
 186         if result_dir is not None:
 187             self.problem.save_quizzes(
 188                 self.train_w_quizzes[:72], result_dir, "culture_w_quizzes"
 189             )
 190
 191     def batches(self, split="train", desc=None):
 192         assert split in {"train", "test"}
 193         if split == "train":
 194             w_quizzes = self.train_w_quizzes
 195             c_quizzes = self.train_c_quizzes
 196         else:
 197             w_quizzes = self.test_w_quizzes
 198             c_quizzes = self.test_c_quizzes
 199
 200         if len(c_quizzes) > 0:
 201             c_quizzes = torch.cat(c_quizzes, dim=0)
 202             if c_quizzes.size(0) > w_quizzes.size(0) // 2:
 203                 i = torch.randperm(c_quizzes.size(0))[: w_quizzes.size(0) // 2]
 204                 c_quizzes = c_quizzes[i]
 205
 206             i = torch.randperm(w_quizzes.size(0))[
 207                 : w_quizzes.size(0) - c_quizzes.size(0)
 208             ]
 209             w_quizzes = w_quizzes[i]
 210
 211             self.nb_batch_w_quizzes = w_quizzes.size(0)
 212             self.nb_batch_c_quizzes = c_quizzes.size(0)
 213
 214             input = torch.cat([w_quizzes, c_quizzes], dim=0)
 215         else:
 216             input = w_quizzes
 217             self.nb_batch_w_quizzes = w_quizzes.size(0)
 218             self.nb_batch_c_quizzes = 0
 219
 220         # Shuffle
 221         input = input[torch.randperm(input.size(0))]
 222
 223         if desc is None:
 224             desc = f"epoch-{split}"
 225         for batch in tqdm.tqdm(
 226             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 227         ):
 228             yield batch
 229
 230     def vocabulary_size(self):
 231         return self.nb_codes
 232
 233     def produce_results(
 234         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
 235     ):
 236         def compute_accuracy(input, logger=None):
 237             input = input[:nmax]
 238             ar_mask = self.make_ar_mask(input)
 239             result = input.clone() * (1 - ar_mask)
 240             seq_logproba = torch.empty(input.size(0), device=self.device)
 241
 242             masked_inplace_autoregression(
 243                 model=model,
 244                 batch_size=self.batch_size,
 245                 input=result,
 246                 ar_mask=ar_mask,
 247                 seq_logproba=seq_logproba,
 248                 temperature=1.0,
 249                 deterministic_synthesis=deterministic_synthesis,
 250                 progress_bar_desc=None,
 251                 device=self.device,
 252             )
 253
 254             nb_total, nb_correct = (
 255                 input.size(0),
 256                 (input == result).long().min(dim=1).values.sum(),
 257             )
 258
 259             return nb_total, nb_correct
 260
 261         train_nb_total, train_nb_correct = compute_accuracy(self.train_w_quizzes)
 262
 263         logger(
 264             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}%"
 265         )
 266
 267         test_nb_total, test_nb_correct = compute_accuracy(self.test_w_quizzes, logger)
 268
 269         logger(
 270             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}%"
 271         )
 272
 273         main_test_accuracy = test_nb_correct / test_nb_total
 274         logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}")
 275
 276         ##############################
 277
 278         input = self.test_w_quizzes[:96]
 279         ar_mask = self.make_ar_mask(input)
 280         result = input.clone() * (1 - ar_mask)
 281         seq_logproba = torch.empty(input.size(0), device=self.device)
 282
 283         masked_inplace_autoregression(
 284             model=model,
 285             batch_size=self.batch_size,
 286             input=result,
 287             ar_mask=ar_mask,
 288             seq_logproba=seq_logproba,
 289             temperature=1.0,
 290             deterministic_synthesis=deterministic_synthesis,
 291             progress_bar_desc=None,
 292             device=self.device,
 293         )
 294
 295         self.problem.save_quizzes(
 296             result[:72], result_dir, f"culture_prediction_{n_epoch:04d}_{model.id:02d}"
 297         )
 298
 299         return main_test_accuracy
 300
 301     def renew_w_quizzes(self, nb, for_train=True):
 302         input = self.train_w_quizzes if for_train else self.test_w_quizzes
 303         nb = min(nb, input.size(0))
 304         input[:-nb] = input[nb:].clone()
 305         input[-nb:] = self.problem.generate_token_sequences(nb).to(self.device)
 306
 307     def store_c_quizzes(self, new_c_quizzes, for_train=True):
 308         if for_train:
 309             self.train_c_quizzes.append(new_c_quizzes)
 310         else:
 311             self.test_c_quizzes.append(new_c_quizzes)
 312
 313     def comput_correctness(self, c_quizzes, models_for_validation):
 314         # Create the reverse quizzes
 315
 316         token_forward, token_backward = self.problem.direction_tokens()
 317
 318         l = (c_quizzes.size(1) - 1) // 2
 319         direction = c_quizzes[:, l : l + 1]
 320         direction = self.problem.token_forward * (
 321             direction == self.problem.token_backward
 322         ) + self.problem.token_backward * (direction == self.problem.token_forward)
 323         reverse_c_quizzes = torch.cat(
 324             [c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1
 325         )
 326
 327         ar_mask = self.make_ar_mask(c_quizzes)
 328         seq_logproba = torch.empty(ar_mask.size(0), device=self.device)
 329
 330         # Check how many of models can solve the quizzes in both directions
 331
 332         nb_correct = 0
 333
 334         for model in models_for_validation:
 335             result = c_quizzes.clone()
 336
 337             masked_inplace_autoregression(
 338                 model=model,
 339                 batch_size=self.batch_size,
 340                 input=result,
 341                 ar_mask=ar_mask,
 342                 seq_logproba=seq_logproba,
 343                 temperature=1.0,
 344                 deterministic_synthesis=True,
 345                 # progress_bar_desc="solving c_quizzes",
 346                 device=self.device,
 347             )
 348
 349             correct = (c_quizzes == result).long().min(dim=-1).values
 350
 351             reverse_result = reverse_c_quizzes.clone()
 352
 353             masked_inplace_autoregression(
 354                 model=model,
 355                 batch_size=self.batch_size,
 356                 input=reverse_result,
 357                 ar_mask=ar_mask,
 358                 seq_logproba=seq_logproba,
 359                 temperature=1.0,
 360                 deterministic_synthesis=True,
 361                 # progress_bar_desc="solving reversed c_quizzes",
 362                 device=self.device,
 363             )
 364
 365             reverse_correct = (
 366                 (reverse_c_quizzes == reverse_result).long().min(dim=-1).values
 367             )
 368
 369             nb_correct += correct * reverse_correct
 370
 371         return nb_correct
 372
 373     ###############################################################
 374
 375     def generate_quizzes(self, nb, model_for_generation, min_ave_seq_logproba):
 376         c_quizzes = torch.empty(
 377             nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64
 378         )
 379
 380         ar_mask = torch.full(c_quizzes.size(), 1, device=self.device)
 381         seq_logproba = torch.empty(ar_mask.size(0), device=self.device)
 382
 383         # bracketing of the temperature to get the target logproba
 384
 385         temperature = 1
 386         d_temperature = 1 / 3
 387
 388         while True:
 389             seq_logproba[...] = 0
 390
 391             masked_inplace_autoregression(
 392                 model=model_for_generation,
 393                 batch_size=self.batch_size,
 394                 input=c_quizzes,
 395                 ar_mask=ar_mask,
 396                 seq_logproba=seq_logproba,
 397                 temperature=temperature,
 398                 deterministic_synthesis=False,
 399                 # progress_bar_desc="sampling c_quizzes",
 400                 device=self.device,
 401             )
 402
 403             ave_seq_logproba = seq_logproba.mean()
 404
 405             # If we do not have target logprobs, get out now
 406             if min_ave_seq_logproba is None:
 407                 break
 408
 409             # Oh man that's ugly
 410             if ave_seq_logproba < min_ave_seq_logproba:
 411                 if d_temperature > 0:
 412                     d_temperature *= -1 / 3
 413                 temperature += d_temperature
 414             elif ave_seq_logproba > min_ave_seq_logproba * 0.99:
 415                 if d_temperature < 0:
 416                     d_temperature *= -1 / 3
 417                 temperature += d_temperature
 418             else:
 419                 break
 420
 421             logger(f"changing temperature to {temperature}")
 422
 423         return c_quizzes, seq_logproba.mean()
 424
 425     ######################################################################
 426
 427     def create_c_quizzes(
 428         self,
 429         nb,
 430         model_for_generation,
 431         models_for_validation,
 432         min_ave_seq_logproba,
 433         n_epoch,
 434         result_dir,
 435         logger,
 436     ):
 437         c_quizzes, ave_seq_logproba = self.generate_quizzes(
 438             nb, model_for_generation, min_ave_seq_logproba
 439         )
 440
 441         nb_correct = self.comput_correctness(c_quizzes, models_for_validation)
 442
 443         return c_quizzes, nb_correct, ave_seq_logproba
 444
 445     ######################################################################
 446
 447     def gang_create_c_quizzes(
 448         self,
 449         nb,
 450         nb_models_for_generation,
 451         models,
 452         mode,
 453         min_ave_seq_logproba,
 454         n_epoch,
 455         result_dir,
 456         logger,
 457     ):
 458         model_for_generation = Gang(models, nb_models_for_generation, mode)
 459         models_for_validation = models
 460         return self.create_c_quizzes(
 461             nb,
 462             model_for_generation,
 463             models_for_validation,
 464             min_ave_seq_logproba,
 465             n_epoch,
 466             result_dir,
 467             logger,
 468         )