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     def batches(self, split="train", nb_to_use=-1, desc=None):
 129         assert split in {"train", "test"}
 130         input = self.train_input if split == "train" else self.test_input
 131         if nb_to_use > 0:
 132             input = input[:nb_to_use]
 133         if desc is None:
 134             desc = f"epoch-{split}"
 135         for batch in tqdm.tqdm(
 136             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 137         ):
 138             yield batch
 139
 140     def vocabulary_size(self):
 141         return self.nb_codes
 142
 143     def produce_results(
 144         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
 145     ):
 146         def compute_accuracy(input, ar_mask, logger=None):
 147             input, ar_mask = input[:nmax], ar_mask[:nmax]
 148             result = input.clone() * (1 - ar_mask)
 149
 150             masked_inplace_autoregression(
 151                 model,
 152                 self.batch_size,
 153                 result,
 154                 ar_mask,
 155                 deterministic_synthesis,
 156                 progress_bar_desc=None,
 157                 device=self.device,
 158             )
 159
 160             log_ground_truth = ar_mask.min() == 0
 161
 162             if logger is not None:
 163                 for sp, st in zip(result[:10], input[:10]):
 164                     logger(
 165                         f"test_sequences {n_epoch} prediction   {self.problem.seq2str(sp)}"
 166                     )
 167                     if log_ground_truth:
 168                         logger(
 169                             f"               {n_epoch} ground truth {self.problem.seq2str(st)}"
 170                         )
 171
 172             nb_total, nb_correct = self.problem.compute_nb_correct(
 173                 input, ar_mask, result
 174             )
 175
 176             # nb_total = ar_mask.sum().item()
 177             # nb_correct = ((result == input).long() * ar_mask).sum().item()
 178
 179             return nb_total, nb_correct
 180
 181         train_nb_total, train_nb_correct = compute_accuracy(
 182             self.train_input, self.train_ar_mask
 183         )
 184
 185         logger(
 186             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}%"
 187         )
 188
 189         test_nb_total, test_nb_correct = compute_accuracy(
 190             self.test_input, self.test_ar_mask, logger
 191         )
 192
 193         logger(
 194             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}%"
 195         )
 196
 197         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 198
 199         if save_attention_image is None:
 200             logger("no save_attention_image (is pycairo installed?)")
 201         else:
 202             for k in range(10):
 203                 ns = torch.randint(self.test_input.size(0), (1,)).item()
 204                 input = self.test_input[ns : ns + 1].clone()
 205
 206                 with torch.autograd.no_grad():
 207                     t = model.training
 208                     model.eval()
 209                     model.record_attention(True)
 210                     model(BracketedSequence(input))
 211                     model.train(t)
 212                     ram = model.retrieve_attention()
 213                     model.record_attention(False)
 214
 215                 tokens_output = [c for c in self.problem.seq2str(input[0])]
 216                 tokens_input = ["n/a"] + tokens_output[:-1]
 217                 for n_head in range(ram[0].size(1)):
 218                     filename = os.path.join(
 219                         result_dir, f"sandbox_attention_{k}_h{n_head}.pdf"
 220                     )
 221                     attention_matrices = [m[0, n_head] for m in ram]
 222                     save_attention_image(
 223                         filename,
 224                         tokens_input,
 225                         tokens_output,
 226                         attention_matrices,
 227                         k_top=10,
 228                         # min_total_attention=0.9,
 229                         token_gap=12,
 230                         layer_gap=50,
 231                     )
 232                     logger(f"wrote {filename}")
 233
 234
 235 ######################################################################
 236
 237 import picoclvr
 238
 239
 240 class PicoCLVR(Task):
 241     # Make a tensor from a list of strings
 242     def tensorize(self, descr):
 243         token_descr = [s.strip().split(" ") for s in descr]
 244         l = max([len(s) for s in token_descr])
 245         token_descr = [s + ["<nul>"] * (l - len(s)) for s in token_descr]
 246         id_descr = [[self.token2id[u] for u in s] for s in token_descr]
 247         return torch.tensor(id_descr, device=self.device)
 248
 249     # Make a list of strings from a tensor
 250     def detensorize(self, x):
 251         return [" ".join([self.id2token[t.item()] for t in r]) for r in x]
 252
 253     # trim all the tensors in the tuple z to remove as much token from
 254     # left and right in the first tensor. If z is a tuple, all its
 255     # elements are trimed according to the triming for the first
 256     def trim(self, z, token="<nul>"):
 257         n = self.token2id[token]
 258         if type(z) == tuple:
 259             x = z[0]
 260             i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
 261             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
 262             return tuple([t[:, a:b] for t in z])
 263         else:
 264             i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
 265             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
 266             return z[:, a:b]
 267
 268     ######################
 269
 270     def __init__(
 271         self,
 272         nb_train_samples,
 273         nb_test_samples,
 274         batch_size,
 275         height,
 276         width,
 277         nb_colors=5,
 278         logger=None,
 279         device=torch.device("cpu"),
 280         pruner_train=None,
 281         pruner_eval=None,
 282     ):
 283         super().__init__()
 284
 285         def generate_descr(nb, cache_suffix, pruner):
 286             return picoclvr.generate(
 287                 nb,
 288                 height=self.height,
 289                 width=self.width,
 290                 nb_colors=nb_colors,
 291                 pruner=pruner,
 292             )
 293
 294         self.height = height
 295         self.width = width
 296         self.batch_size = batch_size
 297         self.device = device
 298         self.pruner_train = pruner_train
 299         self.pruner_eval = pruner_eval
 300
 301         if logger is not None:
 302             logger(
 303                 f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
 304             )
 305
 306         self.train_descr = generate_descr(
 307             nb_train_samples, "train", pruner=self.pruner_train
 308         )
 309         self.test_descr = generate_descr(nb_test_samples, "test", pruner=None)
 310
 311         # Build the tokenizer
 312         tokens = {"<nul>", "<img>"}
 313         for d in [self.train_descr, self.test_descr]:
 314             for s in d:
 315                 for t in s.strip().split(" "):
 316                     tokens.add(t)
 317         # make this set a sorted list to get the same tensors given
 318         # the same descr
 319         tokens = list(tokens)
 320         tokens.sort()
 321         self.token2id = dict([(t, n) for n, t in enumerate(tokens)])
 322         self.id2token = dict([(n, t) for n, t in enumerate(tokens)])
 323         self.t_img, self.t_nul = self.token2id["<img>"], self.token2id["<nul>"]
 324
 325         # Tokenize the train and test sets
 326         self.train_input = self.tensorize(self.train_descr)
 327         self.test_input = self.tensorize(self.test_descr)
 328
 329     def batches(self, split="train"):
 330         assert split in {"train", "test"}
 331         input = self.train_input if split == "train" else self.test_input
 332         for batch in tqdm.tqdm(
 333             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
 334         ):
 335             yield self.trim(batch)
 336
 337     def vocabulary_size(self):
 338         return len(self.token2id)
 339
 340     def compute_missing_properties(
 341         self, n_epoch, model, logger, deterministic_synthesis, pruner=None
 342     ):
 343         acc_nb_requested_properties = []
 344         acc_nb_missing_properties = []
 345         acc_nb_results = 0
 346
 347         for input in tqdm.tqdm(
 348             self.test_input.split(self.batch_size),
 349             dynamic_ncols=True,
 350             desc=f"test-properties",
 351         ):
 352             result = input.clone()
 353             ar_mask = (result == self.t_img).long().cumsum(dim=1).clamp(max=1)
 354             result = (1 - ar_mask) * result + ar_mask * self.t_nul
 355             masked_inplace_autoregression(
 356                 model,
 357                 self.batch_size,
 358                 result,
 359                 ar_mask,
 360                 deterministic_synthesis,
 361                 progress_bar_desc=None,
 362                 device=self.device,
 363             )
 364
 365             result_descr = self.detensorize(result)
 366             np = picoclvr.nb_properties(
 367                 result_descr,
 368                 height=self.height,
 369                 width=self.width,
 370                 pruner=pruner,
 371             )
 372             nb_requested_properties, _, nb_missing_properties = zip(*np)
 373             acc_nb_requested_properties += nb_requested_properties
 374             acc_nb_missing_properties += nb_missing_properties
 375             acc_nb_results += len(result_descr)
 376
 377         nb_requested_properties = sum(acc_nb_requested_properties)
 378         nb_missing_properties = sum(acc_nb_missing_properties)
 379
 380         prefix = "" if pruner is None else "pruned_"
 381         logger(f"nb_{prefix}samples {n_epoch} {acc_nb_results}")
 382         logger(
 383             f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}"
 384         )
 385         logger(
 386             f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%"
 387         )
 388
 389         logger(
 390             f"main_test_accuracy {n_epoch} {1-nb_missing_properties/nb_requested_properties}"
 391         )
 392
 393     ######################################################################
 394
 395     def produce_results(
 396         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 397     ):
 398         self.compute_missing_properties(n_epoch, model, logger, deterministic_synthesis)
 399
 400         if self.pruner_eval is not None:
 401             self.compute_missing_properties(n_epoch, model, self.pruner_eval)
 402
 403         nb_tokens_to_generate = self.height * self.width + 3
 404         result_descr = []
 405         nb_per_primer = 8
 406         primer = []
 407
 408         for primer_descr in [
 409             "red above green <sep> green top <sep> blue right of red",
 410             "there is red <sep> there is yellow <sep> there is blue",
 411             "red below yellow <sep> yellow below green <sep> green below blue <sep> red right <sep> yellow left <sep> green right <sep> blue left",
 412             "green bottom <sep> yellow bottom <sep> green left of blue <sep> yellow right of blue <sep> blue top",
 413         ]:
 414             primer += [primer_descr + " <img>"] * nb_per_primer
 415
 416         result = self.tensorize(primer)
 417         fill = result.new_full(
 418             result.size()[:-1] + (self.height * self.width + 1,), self.t_nul
 419         )
 420         result = torch.cat((result, fill), 1)
 421         ar_mask = (result == self.t_nul).long()
 422         masked_inplace_autoregression(
 423             model,
 424             self.batch_size,
 425             result,
 426             ar_mask,
 427             deterministic_synthesis,
 428             device=self.device,
 429         )
 430         result_descr = self.detensorize(result)
 431
 432         np = picoclvr.nb_properties(result_descr, height=self.height, width=self.width)
 433
 434         acc_nb_requested_properties, _, acc_nb_missing_properties = zip(*np)
 435         acc_nb_results = len(result_descr)
 436
 437         nb_requested_properties = sum(acc_nb_requested_properties)
 438         nb_missing_properties = sum(acc_nb_missing_properties)
 439
 440         prefix = "demo_"
 441         logger(f"nb_{prefix}samples {n_epoch} {acc_nb_results}")
 442         logger(
 443             f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}"
 444         )
 445         logger(
 446             f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%"
 447         )
 448
 449         img = picoclvr.descr2img(result_descr, height=self.height, width=self.width)
 450
 451         if img.dim() == 5:
 452             if img.size(1) == 1:
 453                 img = F.pad(img.squeeze(1), pad=(1, 1, 1, 1), value=64)
 454             else:
 455                 img = torch.cat(
 456                     [
 457                         torchvision.utils.make_grid(x, padding=1, pad_value=64)[None]
 458                         for x in img
 459                     ],
 460                     0,
 461                 )
 462
 463         image_name = os.path.join(result_dir, f"picoclvr_result_{n_epoch:04d}.png")
 464         torchvision.utils.save_image(
 465             img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=0.0
 466         )
 467         logger(f"wrote {image_name}")
 468
 469
 470 ######################################################################
 471
 472
 473 class MNIST(Task):
 474     def __init__(
 475         self, nb_train_samples, nb_test_samples, batch_size, device=torch.device("cpu")
 476     ):
 477         super().__init__()
 478
 479         self.nb_train_samples = (nb_train_samples,)
 480         self.nb_test_samples = (nb_test_samples,)
 481         self.batch_size = batch_size
 482         self.device = device
 483         data_set = torchvision.datasets.MNIST(root="./data", train=True, download=True)
 484         self.train_input = data_set.data[:nb_train_samples].view(-1, 28 * 28).long()
 485         data_set = torchvision.datasets.MNIST(root="./data", train=False, download=True)
 486         self.test_input = data_set.data[:nb_test_samples].view(-1, 28 * 28).long()
 487
 488     def batches(self, split="train", nb_to_use=-1, desc=None):
 489         assert split in {"train", "test"}
 490         input = self.train_input if split == "train" else self.test_input
 491         if nb_to_use > 0:
 492             input = input[:nb_to_use]
 493         if desc is None:
 494             desc = f"epoch-{split}"
 495         for batch in tqdm.tqdm(
 496             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 497         ):
 498             yield batch
 499
 500     def vocabulary_size(self):
 501         return 256
 502
 503     def produce_results(
 504         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 505     ):
 506         results = torch.empty(64, 28 * 28, device=self.device, dtype=torch.int64)
 507         ar_mask = torch.full_like(results, 1)
 508         masked_inplace_autoregression(
 509             model,
 510             self.batch_size,
 511             results,
 512             ar_mask,
 513             deterministic_synthesis,
 514             device=self.device,
 515         )
 516         image_name = os.path.join(result_dir, f"mnist_result_{n_epoch:04d}.png")
 517         torchvision.utils.save_image(
 518             1 - results.reshape(-1, 1, 28, 28) / 255.0,
 519             image_name,
 520             nrow=16,
 521             pad_value=0.8,
 522         )
 523         logger(f"wrote {image_name}")
 524
 525
 526 ######################################################################
 527
 528 import maze
 529
 530
 531 class Maze(Task):
 532     def map2seq(self, *m):
 533         return torch.cat([x.flatten(1) for x in m], 1)
 534
 535     def seq2map(self, s):
 536         s = s.reshape(s.size(0), -1, self.height, self.width)
 537         return (s[:, k] for k in range(s.size(1)))
 538
 539     def __init__(
 540         self,
 541         nb_train_samples,
 542         nb_test_samples,
 543         batch_size,
 544         height,
 545         width,
 546         nb_walls,
 547         device=torch.device("cpu"),
 548     ):
 549         super().__init__()
 550
 551         self.batch_size = batch_size
 552         self.height = height
 553         self.width = width
 554         self.device = device
 555
 556         train_mazes, train_paths, _ = maze.create_maze_data(
 557             nb_train_samples,
 558             height=height,
 559             width=width,
 560             nb_walls=nb_walls,
 561             progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-train"),
 562         )
 563         self.train_input = self.map2seq(train_mazes.to(device), train_paths.to(device))
 564
 565         test_mazes, test_paths, _ = maze.create_maze_data(
 566             nb_test_samples,
 567             height=height,
 568             width=width,
 569             nb_walls=nb_walls,
 570             progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-test"),
 571         )
 572         self.test_input = self.map2seq(test_mazes.to(device), test_paths.to(device))
 573
 574         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 575
 576     def batches(self, split="train", nb_to_use=-1, desc=None):
 577         assert split in {"train", "test"}
 578         input = self.train_input if split == "train" else self.test_input
 579         if nb_to_use > 0:
 580             input = input[:nb_to_use]
 581         if desc is None:
 582             desc = f"epoch-{split}"
 583         for batch in tqdm.tqdm(
 584             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 585         ):
 586             yield batch
 587
 588     def vocabulary_size(self):
 589         return self.nb_codes
 590
 591     def compute_error(
 592         self, model, split="train", nb_to_use=-1, deterministic_synthesis=False
 593     ):
 594         nb_total, nb_correct = 0, 0
 595         count = torch.zeros(
 596             self.width * self.height,
 597             self.width * self.height,
 598             device=self.device,
 599             dtype=torch.int64,
 600         )
 601
 602         for input in self.batches(split, nb_to_use):
 603             result = input.clone()
 604             ar_mask = result.new_zeros(result.size())
 605             ar_mask[:, self.height * self.width :] = 1
 606             result *= 1 - ar_mask
 607             masked_inplace_autoregression(
 608                 model,
 609                 self.batch_size,
 610                 result,
 611                 ar_mask,
 612                 deterministic_synthesis,
 613                 progress_bar_desc=None,
 614                 device=self.device,
 615             )
 616             mazes, paths = self.seq2map(result)
 617             path_correctness = maze.path_correctness(mazes, paths)
 618             nb_correct += path_correctness.long().sum()
 619             nb_total += mazes.size(0)
 620
 621             optimal_path_lengths = (
 622                 (input[:, self.height * self.width :] == maze.v_path).long().sum(1)
 623             )
 624             predicted_path_lengths = (
 625                 (result[:, self.height * self.width :] == maze.v_path).long().sum(1)
 626             )
 627             optimal_path_lengths = optimal_path_lengths[path_correctness]
 628             predicted_path_lengths = predicted_path_lengths[path_correctness]
 629             count[optimal_path_lengths, predicted_path_lengths] += 1
 630
 631         if count.max() == 0:
 632             count = None
 633         else:
 634             count = count[
 635                 : count.sum(1).nonzero().max() + 1, : count.sum(0).nonzero().max() + 1
 636             ]
 637
 638         return nb_total, nb_correct, count
 639
 640     def produce_results(
 641         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 642     ):
 643         train_nb_total, train_nb_correct, count = self.compute_error(
 644             model,
 645             "train",
 646             nb_to_use=1000,
 647             deterministic_synthesis=deterministic_synthesis,
 648         )
 649         logger(
 650             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}%"
 651         )
 652
 653         test_nb_total, test_nb_correct, count = self.compute_error(
 654             model,
 655             "test",
 656             nb_to_use=1000,
 657             deterministic_synthesis=deterministic_synthesis,
 658         )
 659         logger(
 660             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}%"
 661         )
 662
 663         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 664
 665         if count is not None:
 666             proportion_optimal = count.diagonal().sum().float() / count.sum()
 667             logger(f"proportion_optimal_test {proportion_optimal*100:.02f}%")
 668             with open(
 669                 os.path.join(result_dir, f"maze_result_{n_epoch:04d}.txt"), "w"
 670             ) as f:
 671                 for i in range(count.size(0)):
 672                     for j in range(count.size(1)):
 673                         eol = " " if j < count.size(1) - 1 else "\n"
 674                         f.write(f"{count[i,j]}{eol}")
 675
 676         input = self.test_input[:48]
 677         result = input.clone()
 678         ar_mask = result.new_zeros(result.size())
 679         ar_mask[:, self.height * self.width :] = 1
 680         result *= 1 - ar_mask
 681         masked_inplace_autoregression(
 682             model,
 683             self.batch_size,
 684             result,
 685             ar_mask,
 686             deterministic_synthesis,
 687             device=self.device,
 688         )
 689
 690         mazes, paths = self.seq2map(input)
 691         _, predicted_paths = self.seq2map(result)
 692
 693         filename = os.path.join(result_dir, f"maze_result_{n_epoch:04d}.png")
 694         maze.save_image(
 695             filename,
 696             mazes=mazes,
 697             target_paths=paths,
 698             predicted_paths=predicted_paths,
 699             path_correct=maze.path_correctness(mazes, predicted_paths),
 700             path_optimal=maze.path_optimality(paths, predicted_paths),
 701         )
 702         logger(f"wrote {filename}")
 703
 704
 705 ######################################################################
 706
 707
 708 import snake
 709
 710
 711 class Snake(Task):
 712     def __init__(
 713         self,
 714         nb_train_samples,
 715         nb_test_samples,
 716         batch_size,
 717         height,
 718         width,
 719         nb_colors,
 720         length,
 721         prompt_length,
 722         device=torch.device("cpu"),
 723     ):
 724         super().__init__()
 725
 726         self.batch_size = batch_size
 727         self.height = height
 728         self.width = width
 729         self.device = device
 730         self.prompt_length = prompt_length
 731
 732         self.train_input, self.train_prior_visits, _, _ = snake.generate_sequences(
 733             nb_train_samples,
 734             height,
 735             width,
 736             nb_colors,
 737             length,
 738             prompt_length,
 739             self.device,
 740         )
 741         self.test_input, self.test_prior_visits, _, _ = snake.generate_sequences(
 742             nb_test_samples,
 743             height,
 744             width,
 745             nb_colors,
 746             length,
 747             prompt_length,
 748             self.device,
 749         )
 750
 751         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 752
 753     def batches(self, split="train", nb_to_use=-1, desc=None):
 754         assert split in {"train", "test"}
 755         input = self.train_input if split == "train" else self.test_input
 756         if nb_to_use > 0:
 757             input = input[:nb_to_use]
 758         if desc is None:
 759             desc = f"epoch-{split}"
 760         for batch in tqdm.tqdm(
 761             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 762         ):
 763             yield batch
 764
 765     def vocabulary_size(self):
 766         return self.nb_codes
 767
 768     def produce_results(
 769         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 770     ):
 771         def compute_nb_correct(input, prior_visits):
 772             result = input.clone()
 773             i = torch.arange(result.size(1), device=result.device)[None, :]
 774             ar_mask = (
 775                 torch.logical_and(i >= self.prompt_length * 2, i % 2 == 0)
 776                 .long()
 777                 .expand_as(result)
 778             )
 779             result *= 1 - ar_mask
 780
 781             masked_inplace_autoregression(
 782                 model,
 783                 self.batch_size,
 784                 result,
 785                 ar_mask,
 786                 deterministic_synthesis,
 787                 device=self.device,
 788             )
 789
 790             nb_total = ((prior_visits > 0) * ar_mask).sum()
 791
 792             nb_correct = ((result == input).long() * (prior_visits > 0) * ar_mask).sum()
 793
 794             return nb_total, nb_correct
 795
 796         test_nb_total, test_nb_correct = compute_nb_correct(
 797             self.test_input[:1000], self.test_prior_visits[:1000]
 798         )
 799
 800         logger(
 801             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}%"
 802         )
 803
 804         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 805
 806
 807 ######################################################################
 808
 809
 810 import stack
 811
 812
 813 class Stack(Task):
 814     def __init__(
 815         self,
 816         nb_train_samples,
 817         nb_test_samples,
 818         batch_size,
 819         logger,
 820         nb_steps,
 821         nb_stacks,
 822         nb_digits,
 823         fraction_values_for_train=None,
 824         device=torch.device("cpu"),
 825     ):
 826         super().__init__()
 827
 828         self.batch_size = batch_size
 829         self.nb_steps = nb_steps
 830         self.nb_stacks = nb_stacks
 831         self.nb_digits = nb_digits
 832         self.device = device
 833
 834         if fraction_values_for_train is None:
 835             values_for_train = None
 836             values_for_test = None
 837         else:
 838             all = torch.randperm(10**nb_digits)
 839             nb_for_train = int(all.size(0) * fraction_values_for_train)
 840             values_for_train = all[:nb_for_train]
 841             values_for_test = all[nb_for_train:]
 842
 843         self.train_input, self.train_stack_counts = stack.generate_sequences(
 844             nb_train_samples,
 845             nb_steps,
 846             nb_stacks,
 847             nb_digits,
 848             values_for_train,
 849             self.device,
 850         )
 851
 852         self.test_input, self.test_stack_counts = stack.generate_sequences(
 853             nb_test_samples,
 854             nb_steps,
 855             nb_stacks,
 856             nb_digits,
 857             values_for_test,
 858             self.device,
 859         )
 860
 861         i = torch.logical_and(self.test_input % 2 == 1, self.test_input < 2 * nb_stacks)
 862         counts = self.test_stack_counts.flatten()[i.flatten()]
 863         counts = F.one_hot(counts).sum(0)
 864         logger(f"test_pop_stack_counts {counts}")
 865
 866         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 867
 868     def batches(self, split="train", nb_to_use=-1, desc=None):
 869         assert split in {"train", "test"}
 870         input = self.train_input if split == "train" else self.test_input
 871         if nb_to_use > 0:
 872             input = input[:nb_to_use]
 873         if desc is None:
 874             desc = f"epoch-{split}"
 875         for batch in tqdm.tqdm(
 876             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 877         ):
 878             yield batch
 879
 880     def vocabulary_size(self):
 881         return self.nb_codes
 882
 883     def produce_results(
 884         self, n_epoch, model, result_dir, logger, deterministic_synthesis
 885     ):
 886         def compute_nb_correct(input):
 887             result = input.clone()
 888             stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
 889             ar_mask = (result != input).long()
 890             masked_inplace_autoregression(
 891                 model,
 892                 self.batch_size,
 893                 result,
 894                 ar_mask,
 895                 deterministic_synthesis,
 896                 device=self.device,
 897             )
 898
 899             errors = ((result != input).long() * ar_mask).reshape(
 900                 -1, 1 + self.nb_digits
 901             )
 902             ar_mask = ar_mask.reshape(-1, 1 + self.nb_digits)
 903
 904             nb_total = ar_mask.max(1).values.sum()
 905             nb_correct = nb_total - errors.max(1).values.sum()
 906
 907             return nb_total, nb_correct
 908
 909         test_nb_total, test_nb_correct = compute_nb_correct(self.test_input[:1000])
 910
 911         logger(
 912             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}%"
 913         )
 914
 915         logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}")
 916
 917         ##############################################################
 918         # Log a few generated sequences
 919         input = self.test_input[:10, : 12 * (1 + self.nb_digits)]
 920         result = input.clone()
 921         stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
 922         ar_mask = (result != input).long()
 923
 924         # for n in range(result.size(0)):
 925         # logger(
 926         # f"test_before {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
 927         # )
 928
 929         masked_inplace_autoregression(
 930             model,
 931             self.batch_size,
 932             result,
 933             ar_mask,
 934             deterministic_synthesis,
 935             device=self.device,
 936         )
 937
 938         for n in range(result.size(0)):
 939             logger(
 940                 f"test_after  {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
 941             )
 942         ##############################################################
 943
 944
 945 ######################################################################
 946
 947 import rpl
 948
 949
 950 class RPL(Task):
 951     def tensorize(self, sequences):
 952         len_max = max([len(x) for x in sequences])
 953         return torch.cat(
 954             [
 955                 torch.tensor(
 956                     [
 957                         [
 958                             self.token2id[str(c)]
 959                             for c in s + ["<nul>"] * (len_max - len(s))
 960                         ]
 961                         for s in sequences
 962                     ]
 963                 )
 964             ],
 965             0,
 966         )
 967
 968     def seq2str(self, seq):
 969         return " ".join([self.id2token[i] for i in seq])
 970
 971     def __init__(
 972         self,
 973         nb_train_samples,
 974         nb_test_samples,
 975         batch_size,
 976         nb_starting_values=3,
 977         max_input=9,
 978         prog_len=6,
 979         nb_runs=5,
 980         no_prog=False,
 981         logger=None,
 982         device=torch.device("cpu"),
 983     ):
 984         super().__init__()
 985
 986         self.batch_size = batch_size
 987         self.device = device
 988         self.no_prog = no_prog
 989
 990         train_sequences = [
 991             rpl.generate(
 992                 nb_starting_values=nb_starting_values,
 993                 nb_result_values_max=4 * nb_starting_values,
 994                 max_input=max_input,
 995                 prog_len=prog_len,
 996                 nb_runs=nb_runs,
 997             )
 998             for _ in tqdm.tqdm(range(nb_train_samples), desc="train-data")
 999         ]
1000
1001         test_sequences = [
1002             rpl.generate(
1003                 nb_starting_values=nb_starting_values,
1004                 nb_result_values_max=4 * nb_starting_values,
1005                 max_input=max_input,
1006                 prog_len=prog_len,
1007                 nb_runs=nb_runs,
1008             )
1009             for _ in tqdm.tqdm(range(nb_test_samples), desc="test-data")
1010         ]
1011
1012         symbols = list(
1013             set(["<nul>"] + [x for l in train_sequences + test_sequences for x in l])
1014         )
1015         val_max = max([x if type(x) is int else 0 for x in symbols])
1016         symbols = list(filter(lambda x: type(x) is str, symbols))
1017         symbols.sort()
1018         symbols += [str(n) for n in range(val_max + 1)]
1019         self.token2id = dict([(c, n) for n, c in enumerate(symbols)])
1020         self.id2token = dict([(n, c) for c, n in self.token2id.items()])
1021
1022         self.t_nul = self.token2id["<nul>"]
1023         self.t_input = self.token2id["<in>"]
1024         self.t_output = self.token2id["<out>"]
1025         self.t_prog = self.token2id["<prg>"]
1026         self.t_end = self.token2id["<end>"]
1027
1028         self.train_input = self.tensorize(train_sequences)
1029         self.test_input = self.tensorize(test_sequences)
1030
1031         if no_prog:
1032             # Excise the program from every train and test example
1033             k = torch.arange(self.train_input.size(1), device=self.train_input.device)[
1034                 None, :
1035             ]
1036             p = (
1037                 ((self.train_input == self.t_prog).long() * k)
1038                 .max(1, keepdim=True)
1039                 .values
1040             )
1041             self.train_input = (
1042                 self.train_input * (k <= p).long()
1043                 + self.t_end * (k == p + 1).long()
1044                 + self.t_nul * (k > p + 1).long()
1045             )
1046             k = torch.arange(self.test_input.size(1), device=self.test_input.device)[
1047                 None, :
1048             ]
1049             p = (
1050                 ((self.test_input == self.t_prog).long() * k)
1051                 .max(1, keepdim=True)
1052                 .values
1053             )
1054             self.test_input = (
1055                 self.test_input * (k <= p).long()
1056                 + self.t_end * (k == p + 1).long()
1057                 + self.t_nul * (k > p + 1).long()
1058             )
1059
1060         if logger is not None:
1061             logger(f"value_max {val_max}")
1062             for x in self.train_input[:25]:
1063                 end = (x != self.t_nul).nonzero().max().item() + 1
1064                 seq = [self.id2token[i.item()] for i in x[:end]]
1065                 s = " ".join(seq)
1066                 logger(f"example_seq {s}")
1067
1068         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
1069
1070     def batches(self, split="train", nb_to_use=-1, desc=None):
1071         assert split in {"train", "test"}
1072         input = self.train_input if split == "train" else self.test_input
1073         if nb_to_use > 0:
1074             input = input[:nb_to_use]
1075         if desc is None:
1076             desc = f"epoch-{split}"
1077         for batch in tqdm.tqdm(
1078             input.split(self.batch_size), dynamic_ncols=True, desc=desc
1079         ):
1080             last = (batch != self.t_nul).max(0).values.nonzero().max() + 3
1081             batch = batch[:, :last].to(self.device)
1082             yield batch
1083
1084     def vocabulary_size(self):
1085         return self.nb_codes
1086
1087     def produce_results(
1088         self, n_epoch, model, result_dir, logger, deterministic_synthesis
1089     ):
1090         # --------------------------------------------------------------------
1091         def compute_nb_errors_prog(input, nb_to_log=0):
1092             result = input.clone()
1093             s = (result == self.t_prog).long()
1094             ar_mask = (s.cumsum(dim=1) - s).clamp(min=0, max=1)
1095             result = (1 - ar_mask) * result + ar_mask * self.t_nul
1096
1097             masked_inplace_autoregression(
1098                 model,
1099                 self.batch_size,
1100                 result,
1101                 ar_mask,
1102                 deterministic_synthesis,
1103                 device=self.device,
1104             )
1105
1106             sum_nb_total, sum_nb_errors = 0, 0
1107             for one_input, one_result in zip(input, result):
1108                 seq = [self.id2token[i.item()] for i in one_result]
1109                 nb_total, nb_errors, prog, stacks = rpl.compute_nb_errors(seq)
1110                 sum_nb_total += 1
1111                 sum_nb_errors += 0 if nb_errors == 0 else 1
1112                 if nb_to_log > 0:
1113                     gt_seq = [self.id2token[i.item()] for i in one_input]
1114                     _, _, gt_prog, _ = rpl.compute_nb_errors(gt_seq)
1115                     gt_prog = " ".join([str(x) for x in gt_prog])
1116                     prog = " ".join([str(x) for x in prog])
1117                     comment = "*" if nb_errors == 0 else "-"
1118                     logger(f"{comment} PROG [{gt_prog}] PREDICTED [{prog}]")
1119                     for start_stack, target_stack, result_stack, correct in stacks:
1120                         comment = "*" if correct else "-"
1121                         start_stack = " ".join([str(x) for x in start_stack])
1122                         target_stack = " ".join([str(x) for x in target_stack])
1123                         result_stack = " ".join([str(x) for x in result_stack])
1124                         logger(
1125                             f"  {comment} [{start_stack}] -> [{target_stack}] PREDICTED [{result_stack}]"
1126                         )
1127                     nb_to_log -= 1
1128
1129             return sum_nb_total, sum_nb_errors
1130
1131         # --------------------------------------------------------------------
1132         def compute_nb_errors_output(input, nb_to_log=0):
1133             result = input.clone()
1134             k = torch.arange(result.size(1), device=result.device)[None, :]
1135             last_output_idx = (
1136                 ((result == self.t_output) * k).max(dim=1, keepdim=True).values
1137             )
1138             first_prog_idx = (
1139                 ((result == self.t_prog) * k).max(dim=1, keepdim=True).values
1140             )
1141             ar_mask = (k > last_output_idx).long() * (k < first_prog_idx).long()
1142             result = (1 - ar_mask) * result + ar_mask * self.t_nul
1143
1144             masked_inplace_autoregression(
1145                 model,
1146                 self.batch_size,
1147                 result,
1148                 ar_mask,
1149                 deterministic_synthesis,
1150                 device=self.device,
1151             )
1152
1153             sum_nb_total, sum_nb_errors = 0, 0
1154             for one_input, one_result, i, j in zip(
1155                 input, result, last_output_idx, first_prog_idx
1156             ):
1157                 seq = [self.id2token[i.item()] for i in one_result]
1158                 sum_nb_total += 1
1159                 correct = (one_input - one_result).abs().max() == 0
1160                 sum_nb_errors += 0 if correct else 1
1161                 if nb_to_log > 0:
1162                     result_stack = [
1163                         self.id2token[i.item()] for i in one_result[i : j + 1]
1164                     ]
1165                     target_stack = [
1166                         self.id2token[i.item()] for i in one_input[i : j + 1]
1167                     ]
1168                     comment = "*" if correct else "-"
1169                     result_stack = " ".join([str(x) for x in result_stack])
1170                     target_stack = " ".join([str(x) for x in target_stack])
1171                     logger(
1172                         f"output_test {comment} [{target_stack}] PREDICTED [{result_stack}]"
1173                     )
1174                     nb_to_log -= 1
1175
1176             return sum_nb_total, sum_nb_errors
1177
1178         # --------------------------------------------------------------------
1179
1180         if not self.no_prog:
1181             test_nb_total, test_nb_errors = compute_nb_errors_prog(
1182                 self.test_input[:1000].to(self.device), nb_to_log=10
1183             )
1184
1185             logger(
1186                 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}%"
1187             )
1188
1189             logger(f"main_test_accuracy {n_epoch} {1-test_nb_errors/test_nb_total}")
1190
1191         test_nb_total, test_nb_errors = compute_nb_errors_output(
1192             self.test_input[:1000].to(self.device), nb_to_log=10
1193         )
1194
1195         logger(
1196             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}%"
1197         )
1198
1199         if save_attention_image is None:
1200             logger("no save_attention_image (is pycairo installed?)")
1201         else:
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         logger=None,
1477         device=torch.device("cpu"),
1478     ):
1479         super().__init__()
1480
1481         self.device = device
1482         self.batch_size = batch_size
1483         self.grid_factory = grid.GridFactory(size=size)
1484
1485         if logger is not None:
1486             logger(
1487                 f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
1488             )
1489
1490         self.train_descr = self.grid_factory.generate_samples(
1491             nb_train_samples, lambda r: tqdm.tqdm(r)
1492         )
1493         self.test_descr = self.grid_factory.generate_samples(
1494             nb_test_samples, lambda r: tqdm.tqdm(r)
1495         )
1496
1497         # Build the tokenizer
1498         tokens = set()
1499         for d in [self.train_descr, self.test_descr]:
1500             for s in d:
1501                 for t in s.strip().split(" "):
1502                     tokens.add(t)
1503         # make this set a sorted list to get the same tensors given
1504         # the same descr
1505         tokens = list(tokens)
1506         tokens.sort()
1507         tokens = ["#"] + tokens
1508         self.token2id = dict([(t, n) for n, t in enumerate(tokens)])
1509         self.id2token = dict([(n, t) for n, t in enumerate(tokens)])
1510         self.t_nul = self.token2id["#"]
1511         self.t_true = self.token2id["true"]
1512         self.t_false = self.token2id["false"]
1513
1514         # Tokenize the train and test sets
1515         self.train_input = self.str2tensor(self.train_descr)
1516         self.test_input = self.str2tensor(self.test_descr)
1517
1518     def batches(self, split="train"):
1519         assert split in {"train", "test"}
1520         input = self.train_input if split == "train" else self.test_input
1521         for batch in tqdm.tqdm(
1522             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
1523         ):
1524             yield self.trim(batch)
1525
1526     def vocabulary_size(self):
1527         return len(self.token2id)
1528
1529     def produce_results(
1530         self, n_epoch, model, result_dir, logger, deterministic_synthesis
1531     ):
1532         correct = self.test_input[:1000]
1533         result = correct.clone()
1534         ar_mask = torch.logical_or(result == self.t_true, result == self.t_false).long()
1535         result *= 1 - ar_mask  # paraaaaanoiaaaaaaa
1536
1537         logger(f"----------------------------------------------------------")
1538
1539         for e in self.tensor2str(result[:10]):
1540             logger(f"test_before {e}")
1541
1542         masked_inplace_autoregression(
1543             model,
1544             self.batch_size,
1545             result,
1546             ar_mask,
1547             deterministic_synthesis,
1548             device=self.device,
1549         )
1550
1551         logger(f"----------------------------------------------------------")
1552
1553         for e in self.tensor2str(result[:10]):
1554             logger(f"test_after  {e}")
1555
1556         logger(f"----------------------------------------------------------")
1557
1558         nb_total = ar_mask.sum().item()
1559         nb_correct = ((correct == result).long() * ar_mask).sum().item()
1560
1561         logger(f"test_performance {n_epoch} {nb_total=} {nb_correct=}")
1562         logger(f"main_test_accuracy {n_epoch} {nb_correct / nb_total}")
1563
1564
1565 ######################################################################
1566
1567 import qmlp
1568
1569
1570 class QMLP(Task):
1571     ######################
1572
1573     def __init__(
1574         self,
1575         nb_train_samples,
1576         nb_test_samples,
1577         batch_size,
1578         result_dir,
1579         logger=None,
1580         device=torch.device("cpu"),
1581     ):
1582         super().__init__()
1583
1584         self.device = device
1585         self.batch_size = batch_size
1586         self.nb_samples_per_mlp = 256
1587
1588         if logger is not None:
1589             logger(
1590                 f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
1591             )
1592
1593         seq, q_test_set, test_error = qmlp.generate_sequence_and_test_set(
1594             nb_mlps=nb_train_samples + nb_test_samples,
1595             nb_samples=self.nb_samples_per_mlp,
1596             device=self.device,
1597             batch_size=64,
1598             nb_epochs=250,
1599             nb_mlps_per_batch=1024,
1600         )
1601
1602         self.train_input = seq[:nb_train_samples]
1603         self.train_q_test_set = q_test_set[:nb_train_samples]
1604         self.train_ref_test_errors = test_error[:nb_train_samples]
1605         self.test_input = seq[nb_train_samples:]
1606         self.test_q_test_set = q_test_set[nb_train_samples:]
1607         self.test_ref_test_errors = test_error[nb_train_samples:]
1608
1609         filename = os.path.join(result_dir, f"train_errors_ref.dat")
1610         with open(filename, "w") as f:
1611             for e in self.train_ref_test_errors:
1612                 f.write(f"{e}\n")
1613
1614         filename = os.path.join(result_dir, f"test_errors_ref.dat")
1615         with open(filename, "w") as f:
1616             for e in self.test_ref_test_errors:
1617                 f.write(f"{e}\n")
1618
1619         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
1620
1621     def batches(self, split="train"):
1622         assert split in {"train", "test"}
1623         input = self.train_input if split == "train" else self.test_input
1624         for batch in tqdm.tqdm(
1625             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
1626         ):
1627             yield batch
1628
1629     def vocabulary_size(self):
1630         return self.nb_codes
1631
1632     def produce_results(
1633         self, n_epoch, model, result_dir, logger, deterministic_synthesis
1634     ):
1635         correct = self.test_input[:1000]
1636         result = correct.clone()
1637         ar_mask = (
1638             torch.arange(result.size(1), device=result.device)
1639             > self.nb_samples_per_mlp * 3 + 1
1640         ).long()[None, :]
1641         ar_mask = ar_mask.expand_as(result)
1642         result *= 1 - ar_mask  # paraaaaanoiaaaaaaa
1643
1644         masked_inplace_autoregression(
1645             model,
1646             self.batch_size,
1647             result,
1648             ar_mask,
1649             deterministic_synthesis,
1650             device=self.device,
1651         )
1652
1653         q_train_set = result[:, : self.nb_samples_per_mlp * 3]
1654         q_params = result[:, self.nb_samples_per_mlp * 3 + 1 :]
1655         error_test = qmlp.evaluate_q_params(q_params, self.test_q_test_set)
1656
1657         filename = os.path.join(result_dir, f"test_errors_{n_epoch:04d}.dat")
1658         with open(filename, "w") as f:
1659             for e in error_test:
1660                 f.write(f"{e}\n")
1661
1662
1663 ######################################################################