tasks.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
   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 try:
  18     from graph import save_attention_image
  19 except ImportError:
  20     save_attention_image = None
  21
  22 ######################################################################
  23
  24
  25 def masked_inplace_autoregression(
  26     model,
  27     batch_size,
  28     input,
  29     ar_mask,
  30     deterministic_synthesis,
  31     forbidden_tokens=None,
  32     progress_bar_desc="autoregression",
  33     device=torch.device("cpu"),
  34 ):
  35     assert input.size() == ar_mask.size()
  36
  37     batches = zip(input.split(batch_size), ar_mask.split(batch_size))
  38
  39     if progress_bar_desc is not None:
  40         batches = tqdm.tqdm(
  41             batches,
  42             dynamic_ncols=True,
  43             desc=progress_bar_desc,
  44             total=(input.size(0) + batch_size - 1) // batch_size,
  45         )
  46
  47     with torch.autograd.no_grad():
  48         t = model.training
  49         model.eval()
  50
  51         for input, ar_mask in batches:
  52             model.masked_inplace_autoregression(
  53                 input, ar_mask, forbidden_tokens, deterministic_synthesis
  54             )
  55
  56         model.train(t)
  57
  58
  59 ######################################################################
  60
  61
  62 class Task:
  63     def batches(self, split="train"):
  64         pass
  65
  66     def vocabulary_size(self):
  67         pass
  68
  69     def produce_results(
  70         self, n_epoch, model, result_dir, logger, deterministic_synthesis
  71     ):
  72         pass
  73
  74
  75 ####################
  76
  77 import problems
  78
  79
  80 class SandBox(Task):
  81     def __init__(
  82         self,
  83         problem,
  84         nb_train_samples,
  85         nb_test_samples,
  86         batch_size,
  87         logger=None,
  88         device=torch.device("cpu"),
  89         max_nb_codes=1024,
  90     ):
  91         super().__init__()
  92
  93         self.batch_size = batch_size
  94         self.device = device
  95         self.problem = problem
  96
  97         self.train_input, self.train_ar_mask = self.problem.generate_sequences(
  98             nb_train_samples
  99         )
 100         self.test_input, self.test_ar_mask = self.problem.generate_sequences(
 101             nb_test_samples
 102         )
 103
 104         self.train_input, self.train_ar_mask = self.train_input.to(
 105             device
 106         ), self.train_ar_mask.to(device)
 107         self.test_input, self.test_ar_mask = self.test_input.to(
 108             device
 109         ), self.test_ar_mask.to(device)
 110
 111         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 112
 113         # A bit of paranoia never hurts
 114         assert self.nb_codes <= max_nb_codes
 115         assert self.train_input.min() >= 0
 116         assert self.test_input.min() >= 0
 117         assert tuple(x.item() for x in self.train_ar_mask.unique()) in {
 118             (0,),
 119             (1,),
 120             (0, 1),
 121         }
 122         assert tuple(x.item() for x in self.test_ar_mask.unique()) in {
 123             (0,),
 124             (1,),
 125             (0, 1),
 126         }
 127
 128         if logger is not None:
 129             for s, a in zip(self.train_input[:100], self.train_ar_mask[:100]):
 130                 logger(f"train_sequences {self.problem.seq2str(s)}")
 131                 a = "".join(["01"[x.item()] for x in a])
 132                 logger(f"                {a}")
 133
 134     def batches(self, split="train", nb_to_use=-1, desc=None):
 135         assert split in {"train", "test"}
 136         input = self.train_input if split == "train" else self.test_input
 137         if nb_to_use > 0:
 138             input = input[:nb_to_use]
 139         if desc is None:
 140             desc = f"epoch-{split}"
 141         for batch in tqdm.tqdm(
 142             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 143         ):
 144             yield batch
 145
 146     def vocabulary_size(self):
 147         return self.nb_codes
 148
 149     def produce_results(
 150         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
 151     ):
 152         def compute_accuracy(input, ar_mask, logger=None):
 153             input, ar_mask = input[:nmax], ar_mask[:nmax]
 154             result = input.clone() * (1 - ar_mask)
 155
 156             masked_inplace_autoregression(
 157                 model,
 158                 self.batch_size,
 159                 result,
 160                 ar_mask,
 161                 deterministic_synthesis,
 162                 progress_bar_desc=None,
 163                 device=self.device,
 164             )
 165
 166             log_ground_truth = ar_mask.min() == 0
 167
 168             if logger is not None:
 169                 for sp, st in zip(result[:10], input[:10]):
 170                     logger(
 171                         f"test_sequences {n_epoch} prediction   {self.problem.seq2str(sp)}"
 172                     )
 173                     if log_ground_truth:
 174                         logger(
 175                             f"               {n_epoch} ground truth {self.problem.seq2str(st)}"
 176                         )
 177
 178             nb_total, nb_correct = self.problem.compute_nb_correct(
 179                 input, ar_mask, result
 180             )
 181
 182             # nb_total = ar_mask.sum().item()
 183             # nb_correct = ((result == input).long() * ar_mask).sum().item()
 184
 185             return nb_total, nb_correct
 186
 187         train_nb_total, train_nb_correct = compute_accuracy(
 188             self.train_input, self.train_ar_mask
 189         )
 190
 191         logger(
 192             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}%"
 193         )
 194
 195         test_nb_total, test_nb_correct = compute_accuracy(
 196             self.test_input, self.test_ar_mask, logger
 197         )
 198
 199         logger(
 200             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}%"
 201         )
 202
 203         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 204
 205         if save_attention_image is None:
 206             logger("no save_attention_image (is pycairo installed?)")
 207         else:
 208             for k in range(10):
 209                 ns = torch.randint(self.test_input.size(0), (1,)).item()
 210                 input = self.test_input[ns : ns + 1].clone()
 211
 212                 with torch.autograd.no_grad():
 213                     t = model.training
 214                     model.eval()
 215                     # model.record_attention(True)
 216                     model(BracketedSequence(input))
 217                     model.train(t)
 218                     # ram = model.retrieve_attention()
 219                     # model.record_attention(False)
 220
 221                 # tokens_output = [c for c in self.problem.seq2str(input[0])]
 222                 # tokens_input = ["n/a"] + tokens_output[:-1]
 223                 # for n_head in range(ram[0].size(1)):
 224                 # filename = os.path.join(
 225                 # result_dir, f"sandbox_attention_{k}_h{n_head}.pdf"
 226                 # )
 227                 # attention_matrices = [m[0, n_head] for m in ram]
 228                 # save_attention_image(
 229                 # filename,
 230                 # tokens_input,
 231                 # tokens_output,
 232                 # attention_matrices,
 233                 # k_top=10,
 234                 ##min_total_attention=0.9,
 235                 # token_gap=12,
 236                 # layer_gap=50,
 237                 # )
 238                 # logger(f"wrote {filename}")
 239
 240
 241 ######################################################################
 242
 243 import picoclvr
 244
 245
 246 class PicoCLVR(Task):
 247     # Make a tensor from a list of strings
 248     def tensorize(self, descr):
 249         token_descr = [s.strip().split(" ") for s in descr]
 250         l = max([len(s) for s in token_descr])
 251         token_descr = [s + ["<nul>"] * (l - len(s)) for s in token_descr]
 252         id_descr = [[self.token2id[u] for u in s] for s in token_descr]
 253         return torch.tensor(id_descr, device=self.device)
 254
 255     # Make a list of strings from a tensor
 256     def detensorize(self, x):
 257         return [" ".join([self.id2token[t.item()] for t in r]) for r in x]
 258
 259     # trim all the tensors in the tuple z to remove as much token from
 260     # left and right in the first tensor. If z is a tuple, all its
 261     # elements are trimed according to the triming for the first
 262     def trim(self, z, token="<nul>"):
 263         n = self.token2id[token]
 264         if type(z) == tuple:
 265             x = z[0]
 266             i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
 267             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
 268             return tuple([t[:, a:b] for t in z])
 269         else:
 270             i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
 271             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
 272             return z[:, a:b]
 273
 274     ######################
 275
 276     def __init__(
 277         self,
 278         nb_train_samples,
 279         nb_test_samples,
 280         batch_size,
 281         height,
 282         width,
 283         nb_colors=5,
 284         logger=None,
 285         device=torch.device("cpu"),
 286         pruner_train=None,
 287         pruner_eval=None,
 288     ):
 289         super().__init__()
 290
 291         def generate_descr(nb, cache_suffix, pruner):
 292             return picoclvr.generate(
 293                 nb,
 294                 height=self.height,
 295                 width=self.width,
 296                 nb_colors=nb_colors,
 297                 pruner=pruner,
 298             )
 299
 300         self.height = height
 301         self.width = width
 302         self.batch_size = batch_size
 303         self.device = device
 304         self.pruner_train = pruner_train
 305         self.pruner_eval = pruner_eval
 306
 307         if logger is not None:
 308             logger(
 309                 f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
 310             )
 311
 312         self.train_descr = generate_descr(
 313             nb_train_samples, "train", pruner=self.pruner_train
 314         )
 315         self.test_descr = generate_descr(nb_test_samples, "test", pruner=None)
 316
 317         # Build the tokenizer
 318         tokens = {"<nul>", "<img>"}
 319         for d in [self.train_descr, self.test_descr]:
 320             for s in d:
 321                 for t in s.strip().split(" "):
 322                     tokens.add(t)
 323         # make this set a sorted list to get the same tensors given
 324         # the same descr
 325         tokens = list(tokens)
 326         tokens.sort()
 327         self.token2id = dict([(t, n) for n, t in enumerate(tokens)])
 328         self.id2token = dict([(n, t) for n, t in enumerate(tokens)])
 329         self.t_img, self.t_nul = self.token2id["<img>"], self.token2id["<nul>"]
 330
 331         # Tokenize the train and test sets
 332         self.train_input = self.tensorize(self.train_descr)
 333         self.test_input = self.tensorize(self.test_descr)
 334
 335     def batches(self, split="train"):
 336         assert split in {"train", "test"}
 337         input = self.train_input if split == "train" else self.test_input
 338         for batch in tqdm.tqdm(
 339             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
 340         ):
 341             yield self.trim(batch)
 342
 343     def vocabulary_size(self):
 344         return len(self.token2id)
 345
 346     def compute_missing_properties(
 347         self, n_epoch, model, logger, deterministic_synthesis, pruner=None
 348     ):
 349         acc_nb_requested_properties = []
 350         acc_nb_missing_properties = []
 351         acc_nb_results = 0
 352
 353         for input in tqdm.tqdm(
 354             self.test_input.split(self.batch_size),
 355             dynamic_ncols=True,
 356             desc=f"test-properties",
 357         ):
 358             result = input.clone()
 359             ar_mask = (result == self.t_img).long().cumsum(dim=1).clamp(max=1)
 360             result = (1 - ar_mask) * result + ar_mask * self.t_nul
 361             masked_inplace_autoregression(
 362                 model,
 363                 self.batch_size,
 364                 result,
 365                 ar_mask,
 366                 deterministic_synthesis,
 367                 progress_bar_desc=None,
 368                 device=self.device,
 369             )
 370
 371             result_descr = self.detensorize(result)
 372             np = picoclvr.nb_properties(
 373                 result_descr,
 374                 height=self.height,
 375                 width=self.width,
 376                 pruner=pruner,
 377             )
 378             nb_requested_properties, _, nb_missing_properties = zip(*np)
 379             acc_nb_requested_properties += nb_requested_properties
 380             acc_nb_missing_properties += nb_missing_properties
 381             acc_nb_results += len(result_descr)
 382
 383         nb_requested_properties = sum(acc_nb_requested_properties)
 384         nb_missing_properties = sum(acc_nb_missing_properties)
 385
 386         prefix = "" if pruner is None else "pruned_"
 387         logger(f"nb_{prefix}samples {n_epoch} {acc_nb_results}")
 388         logger(
 389             f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}"
 390         )
 391         logger(
 392             f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%"
 393         )
 394
 395         logger(
 396             f"main_test_accuracy {n_epoch} {1-nb_missing_properties/nb_requested_properties}"
 397         )
 398
 399     ######################################################################
 400
 401     def produce_results(
 402         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 403     ):
 404         self.compute_missing_properties(n_epoch, model, logger, deterministic_synthesis)
 405
 406         if self.pruner_eval is not None:
 407             self.compute_missing_properties(n_epoch, model, self.pruner_eval)
 408
 409         nb_tokens_to_generate = self.height * self.width + 3
 410         result_descr = []
 411         nb_per_primer = 8
 412         primer = []
 413
 414         for primer_descr in [
 415             "red above green <sep> green top <sep> blue right of red",
 416             "there is red <sep> there is yellow <sep> there is blue",
 417             "red below yellow <sep> yellow below green <sep> green below blue <sep> red right <sep> yellow left <sep> green right <sep> blue left",
 418             "green bottom <sep> yellow bottom <sep> green left of blue <sep> yellow right of blue <sep> blue top",
 419         ]:
 420             primer += [primer_descr + " <img>"] * nb_per_primer
 421
 422         result = self.tensorize(primer)
 423         fill = result.new_full(
 424             result.size()[:-1] + (self.height * self.width + 1,), self.t_nul
 425         )
 426         result = torch.cat((result, fill), 1)
 427         ar_mask = (result == self.t_nul).long()
 428         masked_inplace_autoregression(
 429             model,
 430             self.batch_size,
 431             result,
 432             ar_mask,
 433             deterministic_synthesis,
 434             device=self.device,
 435         )
 436         result_descr = self.detensorize(result)
 437
 438         np = picoclvr.nb_properties(result_descr, height=self.height, width=self.width)
 439
 440         acc_nb_requested_properties, _, acc_nb_missing_properties = zip(*np)
 441         acc_nb_results = len(result_descr)
 442
 443         nb_requested_properties = sum(acc_nb_requested_properties)
 444         nb_missing_properties = sum(acc_nb_missing_properties)
 445
 446         prefix = "demo_"
 447         logger(f"nb_{prefix}samples {n_epoch} {acc_nb_results}")
 448         logger(
 449             f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}"
 450         )
 451         logger(
 452             f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%"
 453         )
 454
 455         img = picoclvr.descr2img(result_descr, height=self.height, width=self.width)
 456
 457         if img.dim() == 5:
 458             if img.size(1) == 1:
 459                 img = F.pad(img.squeeze(1), pad=(1, 1, 1, 1), value=64)
 460             else:
 461                 img = torch.cat(
 462                     [
 463                         torchvision.utils.make_grid(x, padding=1, pad_value=64)[None]
 464                         for x in img
 465                     ],
 466                     0,
 467                 )
 468
 469         image_name = os.path.join(result_dir, f"picoclvr_result_{n_epoch:04d}.png")
 470         torchvision.utils.save_image(
 471             img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=0.0
 472         )
 473         logger(f"wrote {image_name}")
 474
 475
 476 ######################################################################
 477
 478
 479 class MNIST(Task):
 480     def __init__(
 481         self, nb_train_samples, nb_test_samples, batch_size, device=torch.device("cpu")
 482     ):
 483         super().__init__()
 484
 485         self.nb_train_samples = (nb_train_samples,)
 486         self.nb_test_samples = (nb_test_samples,)
 487         self.batch_size = batch_size
 488         self.device = device
 489         data_set = torchvision.datasets.MNIST(root="./data", train=True, download=True)
 490         self.train_input = data_set.data[:nb_train_samples].view(-1, 28 * 28).long()
 491         data_set = torchvision.datasets.MNIST(root="./data", train=False, download=True)
 492         self.test_input = data_set.data[:nb_test_samples].view(-1, 28 * 28).long()
 493
 494     def batches(self, split="train", nb_to_use=-1, desc=None):
 495         assert split in {"train", "test"}
 496         input = self.train_input if split == "train" else self.test_input
 497         if nb_to_use > 0:
 498             input = input[:nb_to_use]
 499         if desc is None:
 500             desc = f"epoch-{split}"
 501         for batch in tqdm.tqdm(
 502             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 503         ):
 504             yield batch
 505
 506     def vocabulary_size(self):
 507         return 256
 508
 509     def produce_results(
 510         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 511     ):
 512         results = torch.empty(64, 28 * 28, device=self.device, dtype=torch.int64)
 513         ar_mask = torch.full_like(results, 1)
 514         masked_inplace_autoregression(
 515             model,
 516             self.batch_size,
 517             results,
 518             ar_mask,
 519             deterministic_synthesis,
 520             device=self.device,
 521         )
 522         image_name = os.path.join(result_dir, f"mnist_result_{n_epoch:04d}.png")
 523         torchvision.utils.save_image(
 524             1 - results.reshape(-1, 1, 28, 28) / 255.0,
 525             image_name,
 526             nrow=16,
 527             pad_value=0.8,
 528         )
 529         logger(f"wrote {image_name}")
 530
 531
 532 ######################################################################
 533
 534 import maze
 535
 536
 537 class Maze(Task):
 538     def map2seq(self, *m):
 539         return torch.cat([x.flatten(1) for x in m], 1)
 540
 541     def seq2map(self, s):
 542         s = s.reshape(s.size(0), -1, self.height, self.width)
 543         return (s[:, k] for k in range(s.size(1)))
 544
 545     def __init__(
 546         self,
 547         nb_train_samples,
 548         nb_test_samples,
 549         batch_size,
 550         height,
 551         width,
 552         nb_walls,
 553         device=torch.device("cpu"),
 554     ):
 555         super().__init__()
 556
 557         self.batch_size = batch_size
 558         self.height = height
 559         self.width = width
 560         self.device = device
 561
 562         train_mazes, train_paths, _ = maze.create_maze_data(
 563             nb_train_samples,
 564             height=height,
 565             width=width,
 566             nb_walls=nb_walls,
 567             progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-train"),
 568         )
 569         self.train_input = self.map2seq(train_mazes.to(device), train_paths.to(device))
 570
 571         test_mazes, test_paths, _ = maze.create_maze_data(
 572             nb_test_samples,
 573             height=height,
 574             width=width,
 575             nb_walls=nb_walls,
 576             progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-test"),
 577         )
 578         self.test_input = self.map2seq(test_mazes.to(device), test_paths.to(device))
 579
 580         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 581
 582     def batches(self, split="train", nb_to_use=-1, desc=None):
 583         assert split in {"train", "test"}
 584         input = self.train_input if split == "train" else self.test_input
 585         if nb_to_use > 0:
 586             input = input[:nb_to_use]
 587         if desc is None:
 588             desc = f"epoch-{split}"
 589         for batch in tqdm.tqdm(
 590             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 591         ):
 592             yield batch
 593
 594     def vocabulary_size(self):
 595         return self.nb_codes
 596
 597     def compute_error(
 598         self, model, split="train", nb_to_use=-1, deterministic_synthesis=False
 599     ):
 600         nb_total, nb_correct = 0, 0
 601         count = torch.zeros(
 602             self.width * self.height,
 603             self.width * self.height,
 604             device=self.device,
 605             dtype=torch.int64,
 606         )
 607
 608         for input in self.batches(split, nb_to_use):
 609             result = input.clone()
 610             ar_mask = result.new_zeros(result.size())
 611             ar_mask[:, self.height * self.width :] = 1
 612             result *= 1 - ar_mask
 613             masked_inplace_autoregression(
 614                 model,
 615                 self.batch_size,
 616                 result,
 617                 ar_mask,
 618                 deterministic_synthesis,
 619                 progress_bar_desc=None,
 620                 device=self.device,
 621             )
 622             mazes, paths = self.seq2map(result)
 623             path_correctness = maze.path_correctness(mazes, paths)
 624             nb_correct += path_correctness.long().sum()
 625             nb_total += mazes.size(0)
 626
 627             optimal_path_lengths = (
 628                 (input[:, self.height * self.width :] == maze.v_path).long().sum(1)
 629             )
 630             predicted_path_lengths = (
 631                 (result[:, self.height * self.width :] == maze.v_path).long().sum(1)
 632             )
 633             optimal_path_lengths = optimal_path_lengths[path_correctness]
 634             predicted_path_lengths = predicted_path_lengths[path_correctness]
 635             count[optimal_path_lengths, predicted_path_lengths] += 1
 636
 637         if count.max() == 0:
 638             count = None
 639         else:
 640             count = count[
 641                 : count.sum(1).nonzero().max() + 1, : count.sum(0).nonzero().max() + 1
 642             ]
 643
 644         return nb_total, nb_correct, count
 645
 646     def produce_results(
 647         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 648     ):
 649         train_nb_total, train_nb_correct, count = self.compute_error(
 650             model,
 651             "train",
 652             nb_to_use=1000,
 653             deterministic_synthesis=deterministic_synthesis,
 654         )
 655         logger(
 656             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}%"
 657         )
 658
 659         test_nb_total, test_nb_correct, count = self.compute_error(
 660             model,
 661             "test",
 662             nb_to_use=1000,
 663             deterministic_synthesis=deterministic_synthesis,
 664         )
 665         logger(
 666             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}%"
 667         )
 668
 669         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 670
 671         if count is not None:
 672             proportion_optimal = count.diagonal().sum().float() / count.sum()
 673             logger(f"proportion_optimal_test {proportion_optimal*100:.02f}%")
 674             with open(
 675                 os.path.join(result_dir, f"maze_result_{n_epoch:04d}.txt"), "w"
 676             ) as f:
 677                 for i in range(count.size(0)):
 678                     for j in range(count.size(1)):
 679                         eol = " " if j < count.size(1) - 1 else "\n"
 680                         f.write(f"{count[i,j]}{eol}")
 681
 682         input = self.test_input[:48]
 683         result = input.clone()
 684         ar_mask = result.new_zeros(result.size())
 685         ar_mask[:, self.height * self.width :] = 1
 686         result *= 1 - ar_mask
 687         masked_inplace_autoregression(
 688             model,
 689             self.batch_size,
 690             result,
 691             ar_mask,
 692             deterministic_synthesis,
 693             device=self.device,
 694         )
 695
 696         mazes, paths = self.seq2map(input)
 697         _, predicted_paths = self.seq2map(result)
 698
 699         filename = os.path.join(result_dir, f"maze_result_{n_epoch:04d}.png")
 700         maze.save_image(
 701             filename,
 702             mazes=mazes,
 703             target_paths=paths,
 704             predicted_paths=predicted_paths,
 705             path_correct=maze.path_correctness(mazes, predicted_paths),
 706             path_optimal=maze.path_optimality(paths, predicted_paths),
 707         )
 708         logger(f"wrote {filename}")
 709
 710
 711 ######################################################################
 712
 713
 714 import snake
 715
 716
 717 class Snake(Task):
 718     def __init__(
 719         self,
 720         nb_train_samples,
 721         nb_test_samples,
 722         batch_size,
 723         height,
 724         width,
 725         nb_colors,
 726         length,
 727         prompt_length,
 728         device=torch.device("cpu"),
 729     ):
 730         super().__init__()
 731
 732         self.batch_size = batch_size
 733         self.height = height
 734         self.width = width
 735         self.device = device
 736         self.prompt_length = prompt_length
 737
 738         self.train_input, self.train_prior_visits, _, _ = snake.generate_sequences(
 739             nb_train_samples,
 740             height,
 741             width,
 742             nb_colors,
 743             length,
 744             prompt_length,
 745             self.device,
 746         )
 747         self.test_input, self.test_prior_visits, _, _ = snake.generate_sequences(
 748             nb_test_samples,
 749             height,
 750             width,
 751             nb_colors,
 752             length,
 753             prompt_length,
 754             self.device,
 755         )
 756
 757         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 758
 759     def batches(self, split="train", nb_to_use=-1, desc=None):
 760         assert split in {"train", "test"}
 761         input = self.train_input if split == "train" else self.test_input
 762         if nb_to_use > 0:
 763             input = input[:nb_to_use]
 764         if desc is None:
 765             desc = f"epoch-{split}"
 766         for batch in tqdm.tqdm(
 767             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 768         ):
 769             yield batch
 770
 771     def vocabulary_size(self):
 772         return self.nb_codes
 773
 774     def produce_results(
 775         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 776     ):
 777         def compute_nb_correct(input, prior_visits):
 778             result = input.clone()
 779             i = torch.arange(result.size(1), device=result.device)[None, :]
 780             ar_mask = (
 781                 torch.logical_and(i >= self.prompt_length * 2, i % 2 == 0)
 782                 .long()
 783                 .expand_as(result)
 784             )
 785             result *= 1 - ar_mask
 786
 787             masked_inplace_autoregression(
 788                 model,
 789                 self.batch_size,
 790                 result,
 791                 ar_mask,
 792                 deterministic_synthesis,
 793                 device=self.device,
 794             )
 795
 796             nb_total = ((prior_visits > 0) * ar_mask).sum()
 797
 798             nb_correct = ((result == input).long() * (prior_visits > 0) * ar_mask).sum()
 799
 800             return nb_total, nb_correct
 801
 802         test_nb_total, test_nb_correct = compute_nb_correct(
 803             self.test_input[:1000], self.test_prior_visits[:1000]
 804         )
 805
 806         logger(
 807             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}%"
 808         )
 809
 810         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 811
 812
 813 ######################################################################
 814
 815
 816 import stack
 817
 818
 819 class Stack(Task):
 820     def __init__(
 821         self,
 822         nb_train_samples,
 823         nb_test_samples,
 824         batch_size,
 825         logger,
 826         nb_steps,
 827         nb_stacks,
 828         nb_digits,
 829         fraction_values_for_train=None,
 830         device=torch.device("cpu"),
 831     ):
 832         super().__init__()
 833
 834         self.batch_size = batch_size
 835         self.nb_steps = nb_steps
 836         self.nb_stacks = nb_stacks
 837         self.nb_digits = nb_digits
 838         self.device = device
 839
 840         if fraction_values_for_train is None:
 841             values_for_train = None
 842             values_for_test = None
 843         else:
 844             all = torch.randperm(10**nb_digits)
 845             nb_for_train = int(all.size(0) * fraction_values_for_train)
 846             values_for_train = all[:nb_for_train]
 847             values_for_test = all[nb_for_train:]
 848
 849         self.train_input, self.train_stack_counts = stack.generate_sequences(
 850             nb_train_samples,
 851             nb_steps,
 852             nb_stacks,
 853             nb_digits,
 854             values_for_train,
 855             self.device,
 856         )
 857
 858         self.test_input, self.test_stack_counts = stack.generate_sequences(
 859             nb_test_samples,
 860             nb_steps,
 861             nb_stacks,
 862             nb_digits,
 863             values_for_test,
 864             self.device,
 865         )
 866
 867         i = torch.logical_and(self.test_input % 2 == 1, self.test_input < 2 * nb_stacks)
 868         counts = self.test_stack_counts.flatten()[i.flatten()]
 869         counts = F.one_hot(counts).sum(0)
 870         logger(f"test_pop_stack_counts {counts}")
 871
 872         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 873
 874     def batches(self, split="train", nb_to_use=-1, desc=None):
 875         assert split in {"train", "test"}
 876         input = self.train_input if split == "train" else self.test_input
 877         if nb_to_use > 0:
 878             input = input[:nb_to_use]
 879         if desc is None:
 880             desc = f"epoch-{split}"
 881         for batch in tqdm.tqdm(
 882             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 883         ):
 884             yield batch
 885
 886     def vocabulary_size(self):
 887         return self.nb_codes
 888
 889     def produce_results(
 890         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 891     ):
 892         def compute_nb_correct(input):
 893             result = input.clone()
 894             stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
 895             ar_mask = (result != input).long()
 896             masked_inplace_autoregression(
 897                 model,
 898                 self.batch_size,
 899                 result,
 900                 ar_mask,
 901                 deterministic_synthesis,
 902                 device=self.device,
 903             )
 904
 905             errors = ((result != input).long() * ar_mask).reshape(
 906                 -1, 1 + self.nb_digits
 907             )
 908             ar_mask = ar_mask.reshape(-1, 1 + self.nb_digits)
 909
 910             nb_total = ar_mask.max(1).values.sum()
 911             nb_correct = nb_total - errors.max(1).values.sum()
 912
 913             return nb_total, nb_correct
 914
 915         test_nb_total, test_nb_correct = compute_nb_correct(self.test_input[:1000])
 916
 917         logger(
 918             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}%"
 919         )
 920
 921         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 922
 923         ##############################################################
 924         # Log a few generated sequences
 925         input = self.test_input[:10, : 12 * (1 + self.nb_digits)]
 926         result = input.clone()
 927         stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
 928         ar_mask = (result != input).long()
 929
 930         # for n in range(result.size(0)):
 931         # logger(
 932         # f"test_before {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
 933         # )
 934
 935         masked_inplace_autoregression(
 936             model,
 937             self.batch_size,
 938             result,
 939             ar_mask,
 940             deterministic_synthesis,
 941             device=self.device,
 942         )
 943
 944         for n in range(result.size(0)):
 945             logger(
 946                 f"test_after  {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
 947             )
 948         ##############################################################
 949
 950
 951 ######################################################################
 952
 953 import rpl
 954
 955
 956 class RPL(Task):
 957     def tensorize(self, sequences):
 958         len_max = max([len(x) for x in sequences])
 959         return torch.cat(
 960             [
 961                 torch.tensor(
 962                     [
 963                         [
 964                             self.token2id[str(c)]
 965                             for c in s + ["<nul>"] * (len_max - len(s))
 966                         ]
 967                         for s in sequences
 968                     ]
 969                 )
 970             ],
 971             0,
 972         )
 973
 974     def seq2str(self, seq):
 975         return " ".join([self.id2token[i] for i in seq])
 976
 977     def __init__(
 978         self,
 979         nb_train_samples,
 980         nb_test_samples,
 981         batch_size,
 982         nb_starting_values=3,
 983         max_input=9,
 984         prog_len=6,
 985         nb_runs=5,
 986         no_prog=False,
 987         logger=None,
 988         device=torch.device("cpu"),
 989     ):
 990         super().__init__()
 991
 992         self.batch_size = batch_size
 993         self.device = device
 994         self.no_prog = no_prog
 995
 996         train_sequences = [
 997             rpl.generate(
 998                 nb_starting_values=nb_starting_values,
 999                 nb_result_values_max=4 * nb_starting_values,
1000                 max_input=max_input,
1001                 prog_len=prog_len,
1002                 nb_runs=nb_runs,
1003             )
1004             for _ in tqdm.tqdm(range(nb_train_samples), desc="train-data")
1005         ]
1006
1007         test_sequences = [
1008             rpl.generate(
1009                 nb_starting_values=nb_starting_values,
1010                 nb_result_values_max=4 * nb_starting_values,
1011                 max_input=max_input,
1012                 prog_len=prog_len,
1013                 nb_runs=nb_runs,
1014             )
1015             for _ in tqdm.tqdm(range(nb_test_samples), desc="test-data")
1016         ]
1017
1018         symbols = list(
1019             set(["<nul>"] + [x for l in train_sequences + test_sequences for x in l])
1020         )
1021         val_max = max([x if type(x) is int else 0 for x in symbols])
1022         symbols = list(filter(lambda x: type(x) is str, symbols))
1023         symbols.sort()
1024         symbols += [str(n) for n in range(val_max + 1)]
1025         self.token2id = dict([(c, n) for n, c in enumerate(symbols)])
1026         self.id2token = dict([(n, c) for c, n in self.token2id.items()])
1027
1028         self.t_nul = self.token2id["<nul>"]
1029         self.t_input = self.token2id["<in>"]
1030         self.t_output = self.token2id["<out>"]
1031         self.t_prog = self.token2id["<prg>"]
1032         self.t_end = self.token2id["<end>"]
1033
1034         self.train_input = self.tensorize(train_sequences)
1035         self.test_input = self.tensorize(test_sequences)
1036
1037         if no_prog:
1038             # Excise the program from every train and test example
1039             k = torch.arange(self.train_input.size(1), device=self.train_input.device)[
1040                 None, :
1041             ]
1042             p = (
1043                 ((self.train_input == self.t_prog).long() * k)
1044                 .max(1, keepdim=True)
1045                 .values
1046             )
1047             self.train_input = (
1048                 self.train_input * (k <= p).long()
1049                 + self.t_end * (k == p + 1).long()
1050                 + self.t_nul * (k > p + 1).long()
1051             )
1052             k = torch.arange(self.test_input.size(1), device=self.test_input.device)[
1053                 None, :
1054             ]
1055             p = (
1056                 ((self.test_input == self.t_prog).long() * k)
1057                 .max(1, keepdim=True)
1058                 .values
1059             )
1060             self.test_input = (
1061                 self.test_input * (k <= p).long()
1062                 + self.t_end * (k == p + 1).long()
1063                 + self.t_nul * (k > p + 1).long()
1064             )
1065
1066         if logger is not None:
1067             logger(f"value_max {val_max}")
1068             for x in self.train_input[:25]:
1069                 end = (x != self.t_nul).nonzero().max().item() + 1
1070                 seq = [self.id2token[i.item()] for i in x[:end]]
1071                 s = " ".join(seq)
1072                 logger(f"example_seq {s}")
1073
1074         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
1075
1076     def batches(self, split="train", nb_to_use=-1, desc=None):
1077         assert split in {"train", "test"}
1078         input = self.train_input if split == "train" else self.test_input
1079         if nb_to_use > 0:
1080             input = input[:nb_to_use]
1081         if desc is None:
1082             desc = f"epoch-{split}"
1083         for batch in tqdm.tqdm(
1084             input.split(self.batch_size), dynamic_ncols=True, desc=desc
1085         ):
1086             last = (batch != self.t_nul).max(0).values.nonzero().max() + 3
1087             batch = batch[:, :last].to(self.device)
1088             yield batch
1089
1090     def vocabulary_size(self):
1091         return self.nb_codes
1092
1093     def produce_results(
1094         self, n_epoch, model, result_dir, logger, deterministic_synthesis
1095     ):
1096         # --------------------------------------------------------------------
1097         def compute_nb_errors_prog(input, nb_to_log=0):
1098             result = input.clone()
1099             s = (result == self.t_prog).long()
1100             ar_mask = (s.cumsum(dim=1) - s).clamp(min=0, max=1)
1101             result = (1 - ar_mask) * result + ar_mask * self.t_nul
1102
1103             masked_inplace_autoregression(
1104                 model,
1105                 self.batch_size,
1106                 result,
1107                 ar_mask,
1108                 deterministic_synthesis,
1109                 device=self.device,
1110             )
1111
1112             sum_nb_total, sum_nb_errors = 0, 0
1113             for one_input, one_result in zip(input, result):
1114                 seq = [self.id2token[i.item()] for i in one_result]
1115                 nb_total, nb_errors, prog, stacks = rpl.compute_nb_errors(seq)
1116                 sum_nb_total += 1
1117                 sum_nb_errors += 0 if nb_errors == 0 else 1
1118                 if nb_to_log > 0:
1119                     gt_seq = [self.id2token[i.item()] for i in one_input]
1120                     _, _, gt_prog, _ = rpl.compute_nb_errors(gt_seq)
1121                     gt_prog = " ".join([str(x) for x in gt_prog])
1122                     prog = " ".join([str(x) for x in prog])
1123                     comment = "*" if nb_errors == 0 else "-"
1124                     logger(f"{comment} PROG [{gt_prog}] PREDICTED [{prog}]")
1125                     for start_stack, target_stack, result_stack, correct in stacks:
1126                         comment = "*" if correct else "-"
1127                         start_stack = " ".join([str(x) for x in start_stack])
1128                         target_stack = " ".join([str(x) for x in target_stack])
1129                         result_stack = " ".join([str(x) for x in result_stack])
1130                         logger(
1131                             f"  {comment} [{start_stack}] -> [{target_stack}] PREDICTED [{result_stack}]"
1132                         )
1133                     nb_to_log -= 1
1134
1135             return sum_nb_total, sum_nb_errors
1136
1137         # --------------------------------------------------------------------
1138         def compute_nb_errors_output(input, nb_to_log=0):
1139             result = input.clone()
1140             k = torch.arange(result.size(1), device=result.device)[None, :]
1141             last_output_idx = (
1142                 ((result == self.t_output) * k).max(dim=1, keepdim=True).values
1143             )
1144             first_prog_idx = (
1145                 ((result == self.t_prog) * k).max(dim=1, keepdim=True).values
1146             )
1147             ar_mask = (k > last_output_idx).long() * (k < first_prog_idx).long()
1148             result = (1 - ar_mask) * result + ar_mask * self.t_nul
1149
1150             masked_inplace_autoregression(
1151                 model,
1152                 self.batch_size,
1153                 result,
1154                 ar_mask,
1155                 deterministic_synthesis,
1156                 device=self.device,
1157             )
1158
1159             sum_nb_total, sum_nb_errors = 0, 0
1160             for one_input, one_result, i, j in zip(
1161                 input, result, last_output_idx, first_prog_idx
1162             ):
1163                 seq = [self.id2token[i.item()] for i in one_result]
1164                 sum_nb_total += 1
1165                 correct = (one_input - one_result).abs().max() == 0
1166                 sum_nb_errors += 0 if correct else 1
1167                 if nb_to_log > 0:
1168                     result_stack = [
1169                         self.id2token[i.item()] for i in one_result[i : j + 1]
1170                     ]
1171                     target_stack = [
1172                         self.id2token[i.item()] for i in one_input[i : j + 1]
1173                     ]
1174                     comment = "*" if correct else "-"
1175                     result_stack = " ".join([str(x) for x in result_stack])
1176                     target_stack = " ".join([str(x) for x in target_stack])
1177                     logger(
1178                         f"output_test {comment} [{target_stack}] PREDICTED [{result_stack}]"
1179                     )
1180                     nb_to_log -= 1
1181
1182             return sum_nb_total, sum_nb_errors
1183
1184         # --------------------------------------------------------------------
1185
1186         if not self.no_prog:
1187             test_nb_total, test_nb_errors = compute_nb_errors_prog(
1188                 self.test_input[:1000].to(self.device), nb_to_log=10
1189             )
1190
1191             logger(
1192                 f"accuracy_prog_test {n_epoch} nb_total {test_nb_total} nb_errors {test_nb_errors} accuracy {100.0*(1-test_nb_errors/test_nb_total):.02f}%"
1193             )
1194
1195             logger(f"main_test_accuracy {n_epoch} {1-test_nb_errors/test_nb_total}")
1196
1197         test_nb_total, test_nb_errors = compute_nb_errors_output(
1198             self.test_input[:1000].to(self.device), nb_to_log=10
1199         )
1200
1201         logger(
1202             f"accuracy_output_test {n_epoch} nb_total {test_nb_total} nb_errors {test_nb_errors} accuracy {100.0*(1-test_nb_errors/test_nb_total):.02f}%"
1203         )
1204
1205         if save_attention_image is None:
1206             logger("no save_attention_image (is pycairo installed?)")
1207         else:
1208             ns = torch.randint(self.test_input.size(0), (1,)).item()
1209             input = self.test_input[ns : ns + 1].clone()
1210             last = (input != self.t_nul).max(0).values.nonzero().max() + 3
1211             input = input[:, :last].to(self.device)
1212
1213             with torch.autograd.no_grad():
1214                 t = model.training
1215                 model.eval()
1216                 model.record_attention(True)
1217                 model(BracketedSequence(input))
1218                 model.train(t)
1219                 ram = model.retrieve_attention()
1220                 model.record_attention(False)
1221
1222             tokens_output = [self.id2token[i.item()] for i in input[0]]
1223             tokens_input = ["n/a"] + tokens_output[:-1]
1224             for n_head in range(ram[0].size(1)):
1225                 filename = os.path.join(
1226                     result_dir, f"rpl_attention_{n_epoch}_h{n_head}.pdf"
1227                 )
1228                 attention_matrices = [m[0, n_head] for m in ram]
1229                 save_attention_image(
1230                     filename,
1231                     tokens_input,
1232                     tokens_output,
1233                     attention_matrices,
1234                     k_top=10,
1235                     # min_total_attention=0.9,
1236                     token_gap=12,
1237                     layer_gap=50,
1238                 )
1239                 logger(f"wrote {filename}")
1240
1241
1242 ######################################################################
1243
1244
1245 import expr
1246
1247
1248 class Expr(Task):
1249     def tensorize(self, sequences):
1250         len_max = max([len(x) for x in sequences])
1251         return torch.cat(
1252             [
1253                 torch.tensor(
1254                     [
1255                         [self.char2id[c] for c in s + "#" * (len_max - len(s))]
1256                         for s in sequences
1257                     ]
1258                 )
1259             ],
1260             0,
1261         ).to(self.device)
1262
1263     def __init__(
1264         self,
1265         nb_train_samples,
1266         nb_test_samples,
1267         nb_variables,
1268         sequence_length,
1269         operand_max,
1270         result_max,
1271         batch_size,
1272         device=torch.device("cpu"),
1273     ):
1274         super().__init__()
1275
1276         self.batch_size = batch_size
1277         self.device = device
1278
1279         train_sequences = expr.generate_sequences(
1280             nb_train_samples,
1281             nb_variables=nb_variables,
1282             length=sequence_length,
1283             operand_max=operand_max,
1284             result_max=result_max,
1285         )
1286
1287         test_sequences = expr.generate_sequences(
1288             nb_test_samples,
1289             nb_variables=nb_variables,
1290             length=sequence_length,
1291             operand_max=operand_max,
1292             result_max=result_max,
1293         )
1294
1295         symbols = list(set("#" + "".join(train_sequences + test_sequences)))
1296         symbols.sort()
1297
1298         self.char2id = dict([(c, n) for n, c in enumerate(symbols)])
1299         self.id2char = dict([(n, c) for c, n in self.char2id.items()])
1300
1301         self.filler, self.space = self.char2id["#"], self.char2id[" "]
1302
1303         self.train_input = self.tensorize(train_sequences)
1304         self.test_input = self.tensorize(test_sequences)
1305
1306         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
1307
1308     def batches(self, split="train", nb_to_use=-1, desc=None):
1309         assert split in {"train", "test"}
1310         input = self.train_input if split == "train" else self.test_input
1311         if nb_to_use > 0:
1312             input = input[:nb_to_use]
1313         if desc is None:
1314             desc = f"epoch-{split}"
1315         for batch in tqdm.tqdm(
1316             input.split(self.batch_size), dynamic_ncols=True, desc=desc
1317         ):
1318             last = (batch != self.filler).max(0).values.nonzero().max() + 3
1319             batch = batch[:, :last]
1320             yield batch
1321
1322     def vocabulary_size(self):
1323         return self.nb_codes
1324
1325     def seq2str(self, s):
1326         return "".join([self.id2char[k.item()] for k in s])
1327
1328     def produce_results(
1329         self,
1330         n_epoch,
1331         model,
1332         result_dir,
1333         logger,
1334         deterministic_synthesis,
1335         input_file=None,
1336     ):
1337         def compute_nb_correct(input):
1338             result = input.clone()
1339             s = (result == self.space).long()
1340             ar_mask = (s.cumsum(dim=1) - s).clamp(min=0, max=1)
1341             result = (1 - ar_mask) * result + ar_mask * self.filler
1342             masked_inplace_autoregression(
1343                 model,
1344                 self.batch_size,
1345                 result,
1346                 ar_mask,
1347                 deterministic_synthesis,
1348                 device=self.device,
1349             )
1350
1351             nb_total = input.size(0)
1352             nb_correct = (input == result).long().min(1).values.sum()
1353
1354             #######################################################################
1355             # Comput predicted vs. true variable values
1356
1357             nb_delta = torch.zeros(5, dtype=torch.int64)
1358             nb_missed = 0
1359
1360             values_input = expr.extract_results([self.seq2str(s) for s in input])
1361             values_result = expr.extract_results([self.seq2str(s) for s in result])
1362
1363             filename = os.path.join(result_dir, f"expr_result_{n_epoch:04d}.txt")
1364
1365             with open(filename, "w") as f:
1366                 for i, r in zip(values_input, values_result):
1367                     for n, vi in i.items():
1368                         vr = r.get(n)
1369                         f.write(f"{vi} {-1 if vr is None else vr}\n")
1370
1371                         if vr is None or vr < 0:
1372                             nb_missed += 1
1373                         else:
1374                             d = abs(vr - vi)
1375                             if d >= nb_delta.size(0):
1376                                 nb_missed += 1
1377                             else:
1378                                 nb_delta[d] += 1
1379
1380             ######################################################################
1381
1382             return nb_total, nb_correct, nb_delta, nb_missed
1383
1384         (
1385             test_nb_total,
1386             test_nb_correct,
1387             test_nb_delta,
1388             test_nb_missed,
1389         ) = compute_nb_correct(self.test_input[:10000])
1390
1391         logger(
1392             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}%"
1393         )
1394
1395         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
1396
1397         nb_total = test_nb_delta.sum() + test_nb_missed
1398         for d in range(test_nb_delta.size(0)):
1399             logger(
1400                 f"error_value {n_epoch} delta {d} {test_nb_delta[d]} {test_nb_delta[d]*100/nb_total:.02f}%"
1401             )
1402         logger(
1403             f"error_value {n_epoch} missed {test_nb_missed} {test_nb_missed*100/nb_total:.02f}%"
1404         )
1405
1406         ##############################################################
1407         # Log a few generated sequences
1408         if input_file is None:
1409             input = self.test_input[:10]
1410         else:
1411             with open(input_file, "r") as f:
1412                 sequences = [e.strip() for e in f.readlines()]
1413                 sequences = [s + " " + "#" * 50 for s in sequences]
1414                 input = self.tensorize(sequences)
1415
1416         result = input.clone()
1417         s = (result == self.space).long()
1418         ar_mask = (s.cumsum(dim=1) - s).clamp(min=0, max=1)
1419         result = (1 - ar_mask) * result + ar_mask * self.filler
1420
1421         for n in range(result.size(0)):
1422             logger(f"test_before {self.seq2str(result[n])}")
1423
1424         masked_inplace_autoregression(
1425             model,
1426             self.batch_size,
1427             result,
1428             ar_mask,
1429             deterministic_synthesis,
1430             device=self.device,
1431         )
1432
1433         correct = (1 - ar_mask) * self.space + ar_mask * input
1434         for n in range(result.size(0)):
1435             comment = "GOOD" if (result[n] - input[n]).abs().max() == 0 else ""
1436             logger(f"test_after  {self.seq2str(result[n])} {comment}")
1437             logger(f"truth       {self.seq2str(correct[n])}")
1438         ##############################################################
1439
1440
1441 ######################################################################
1442
1443 import grid
1444
1445
1446 class Grid(Task):
1447     # Make a tensor from a list of strings
1448     def str2tensor(self, descr):
1449         token_descr = [s.strip().split(" ") for s in descr]
1450         l = max([len(s) for s in token_descr])
1451         token_descr = [s + ["#"] * (l - len(s)) for s in token_descr]
1452         id_descr = [[self.token2id[u] for u in s] for s in token_descr]
1453         return torch.tensor(id_descr, device=self.device)
1454
1455     # Make a list of strings from a tensor
1456     def tensor2str(self, x):
1457         return [" ".join([self.id2token[t.item()] for t in r]) for r in x]
1458
1459     # trim all the tensors in the tuple z to remove as much token from
1460     # left and right in the first tensor. If z is a tuple, all its
1461     # elements are trimed according to the triming for the first
1462     def trim(self, z, token="#"):
1463         n = self.token2id[token]
1464         if type(z) == tuple:
1465             x = z[0]
1466             i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
1467             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
1468             return tuple([t[:, a:b] for t in z])
1469         else:
1470             i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
1471             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
1472             return z[:, a:b]
1473
1474     ######################
1475
1476     def __init__(
1477         self,
1478         nb_train_samples,
1479         nb_test_samples,
1480         batch_size,
1481         size,
1482         logger=None,
1483         device=torch.device("cpu"),
1484     ):
1485         super().__init__()
1486
1487         self.device = device
1488         self.batch_size = batch_size
1489         self.grid_factory = grid.GridFactory(size=size)
1490
1491         if logger is not None:
1492             logger(
1493                 f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
1494             )
1495
1496         self.train_descr = self.grid_factory.generate_samples(
1497             nb_train_samples, lambda r: tqdm.tqdm(r)
1498         )
1499         self.test_descr = self.grid_factory.generate_samples(
1500             nb_test_samples, lambda r: tqdm.tqdm(r)
1501         )
1502
1503         # Build the tokenizer
1504         tokens = set()
1505         for d in [self.train_descr, self.test_descr]:
1506             for s in d:
1507                 for t in s.strip().split(" "):
1508                     tokens.add(t)
1509         # make this set a sorted list to get the same tensors given
1510         # the same descr
1511         tokens = list(tokens)
1512         tokens.sort()
1513         tokens = ["#"] + tokens
1514         self.token2id = dict([(t, n) for n, t in enumerate(tokens)])
1515         self.id2token = dict([(n, t) for n, t in enumerate(tokens)])
1516         self.t_nul = self.token2id["#"]
1517         self.t_true = self.token2id["true"]
1518         self.t_false = self.token2id["false"]
1519
1520         # Tokenize the train and test sets
1521         self.train_input = self.str2tensor(self.train_descr)
1522         self.test_input = self.str2tensor(self.test_descr)
1523
1524     def batches(self, split="train"):
1525         assert split in {"train", "test"}
1526         input = self.train_input if split == "train" else self.test_input
1527         for batch in tqdm.tqdm(
1528             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
1529         ):
1530             yield self.trim(batch)
1531
1532     def vocabulary_size(self):
1533         return len(self.token2id)
1534
1535     def produce_results(
1536         self, n_epoch, model, result_dir, logger, deterministic_synthesis
1537     ):
1538         correct = self.test_input[:1000]
1539         result = correct.clone()
1540         ar_mask = torch.logical_or(result == self.t_true, result == self.t_false).long()
1541         result *= 1 - ar_mask  # paraaaaanoiaaaaaaa
1542
1543         logger(f"----------------------------------------------------------")
1544
1545         for e in self.tensor2str(result[:10]):
1546             logger(f"test_before {e}")
1547
1548         masked_inplace_autoregression(
1549             model,
1550             self.batch_size,
1551             result,
1552             ar_mask,
1553             deterministic_synthesis,
1554             device=self.device,
1555         )
1556
1557         logger(f"----------------------------------------------------------")
1558
1559         for e in self.tensor2str(result[:10]):
1560             logger(f"test_after  {e}")
1561
1562         logger(f"----------------------------------------------------------")
1563
1564         nb_total = ar_mask.sum().item()
1565         nb_correct = ((correct == result).long() * ar_mask).sum().item()
1566
1567         logger(f"test_performance {n_epoch} {nb_total=} {nb_correct=}")
1568         logger(f"main_test_accuracy {n_epoch} {nb_correct / nb_total}")
1569
1570
1571 ######################################################################
1572
1573 import qmlp
1574
1575
1576 class QMLP(Task):
1577     ######################
1578
1579     def __init__(
1580         self,
1581         nb_train_samples,
1582         nb_test_samples,
1583         batch_size,
1584         result_dir,
1585         logger=None,
1586         device=torch.device("cpu"),
1587     ):
1588         super().__init__()
1589
1590         self.device = device
1591         self.batch_size = batch_size
1592         self.nb_samples_per_mlp = 256
1593
1594         if logger is not None:
1595             logger(
1596                 f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
1597             )
1598
1599         seq, q_test_set, test_error = qmlp.generate_sequence_and_test_set(
1600             nb_mlps=nb_train_samples + nb_test_samples,
1601             nb_samples=self.nb_samples_per_mlp,
1602             device=self.device,
1603             batch_size=64,
1604             nb_epochs=250,
1605             nb_mlps_per_batch=1024,
1606         )
1607
1608         self.train_input = seq[:nb_train_samples]
1609         self.train_q_test_set = q_test_set[:nb_train_samples]
1610         self.train_ref_test_errors = test_error[:nb_train_samples]
1611         self.test_input = seq[nb_train_samples:]
1612         self.test_q_test_set = q_test_set[nb_train_samples:]
1613         self.test_ref_test_errors = test_error[nb_train_samples:]
1614
1615         filename = os.path.join(result_dir, f"train_errors_ref.dat")
1616         with open(filename, "w") as f:
1617             for e in self.train_ref_test_errors:
1618                 f.write(f"{e}\n")
1619
1620         filename = os.path.join(result_dir, f"test_errors_ref.dat")
1621         with open(filename, "w") as f:
1622             for e in self.test_ref_test_errors:
1623                 f.write(f"{e}\n")
1624
1625         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
1626
1627     def batches(self, split="train"):
1628         assert split in {"train", "test"}
1629         input = self.train_input if split == "train" else self.test_input
1630         for batch in tqdm.tqdm(
1631             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
1632         ):
1633             yield batch
1634
1635     def vocabulary_size(self):
1636         return self.nb_codes
1637
1638     def produce_results(
1639         self, n_epoch, model, result_dir, logger, deterministic_synthesis
1640     ):
1641         correct = self.test_input[:1000]
1642         result = correct.clone()
1643         ar_mask = (
1644             torch.arange(result.size(1), device=result.device)
1645             > self.nb_samples_per_mlp * 3 + 1
1646         ).long()[None, :]
1647         ar_mask = ar_mask.expand_as(result)
1648         result *= 1 - ar_mask  # paraaaaanoiaaaaaaa
1649
1650         masked_inplace_autoregression(
1651             model,
1652             self.batch_size,
1653             result,
1654             ar_mask,
1655             deterministic_synthesis,
1656             device=self.device,
1657         )
1658
1659         q_train_set = result[:, : self.nb_samples_per_mlp * 3]
1660         q_params = result[:, self.nb_samples_per_mlp * 3 + 1 :]
1661         error_test = qmlp.evaluate_q_params(q_params, self.test_q_test_set)
1662
1663         filename = os.path.join(result_dir, f"test_errors_{n_epoch:04d}.dat")
1664         with open(filename, "w") as f:
1665             for e in error_test:
1666                 f.write(f"{e}\n")
1667
1668
1669 ######################################################################