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