X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=78910a06d7a24977b462af0836a9219073381507;hb=b35745d09b33aed20670ecb96726f89206487a24;hp=da39a830e3b4d899e2f3521f73444fa9cdd8c36b;hpb=0c6d29f73e35adbbaab1263de439f73efa98d99e;p=picoclvr.git diff --git a/tasks.py b/tasks.py index da39a83..78910a0 100755 --- a/tasks.py +++ b/tasks.py @@ -12,6 +12,11 @@ import torch, torchvision from torch import nn from torch.nn import functional as F +from mygpt import BracketedSequence + +# from graph import save_attention_image +save_attention_image = None + ###################################################################### @@ -65,159 +70,141 @@ class Task: pass -###################################################################### +class TaskFromFile(Task): + def tensorize(self, pairs): + len_max = max([len(x[0]) for x in pairs]) + input = torch.cat( + [ + torch.tensor( + [ + [self.char2id[c] for c in s[0] + "#" * (len_max - len(s[0]))] + for s in pairs + ] + ) + ], + 0, + ).to("cpu") -class Problem: - def generate_sequences(self, nb): - pass - - def seq2str(self, seq): - return "[NOT IMPLEMENTED]" - + pred_mask = torch.cat( + [ + torch.tensor( + [ + [int(c) for c in s[1] + "0" * (len_max - len(s[1]))] + for s in pairs + ] + ) + ], + 0, + ).to("cpu") -#################### + return input, pred_mask + # trim all the tensors in the tuple z to remove as much token from + # left and right in the first tensor. If z is a tuple, all its + # elements are trimed according to the triming for the first + def trim(self, z, token="#"): + n = self.char2id[token] + if type(z) == tuple: + x = z[0] + i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0) + a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() + return tuple([t[:, a:b] for t in z]) + else: + i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0) + a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() + return z[:, a:b] -class ProblemLevel0(Problem): - def __init__(self, nb_sentences=100, len_prompt=5, len_result=5): - self.seq = torch.randint(10, (nb_sentences, len_prompt + 1 + len_result)) - self.seq[:, len_prompt] = 10 + def __init__( + self, + filename, + nb_train_samples, + nb_test_samples, + batch_size, + device=torch.device("cpu"), + ): + self.batch_size = batch_size + self.device = device - def generate_sequences(self, nb): - sequences = self.seq[torch.randint(self.seq.size(0), (nb,))] - ar_mask = (sequences == 10).long() - ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1) - return sequences, ar_mask + pairs = [] + with open(filename, "r") as f: + for _ in range(nb_train_samples + nb_test_samples): + sequence = f.readline().strip() + pred_mask = f.readline().strip() + assert len(sequence) == len(pred_mask) + assert set(pred_mask).issubset({"0", "1", "2"}), f"{set(pred_mask)}" + pairs.append((sequence, pred_mask)) + symbols = ["#"] + list(set("".join([x[0] for x in pairs])) - set(["#"])) + self.char2id = dict([(c, n) for n, c in enumerate(symbols)]) + self.id2char = dict([(n, c) for c, n in self.char2id.items()]) -class ProblemLevel1(Problem): - def __init__(self, nb_operators=100, len_source=5, len_result=8): - self.len_source = len_source - self.len_result = len_result - self.len_nb_operator = int(math.log(nb_operators) / math.log(10)) + 1 - self.operators = F.one_hot( - torch.rand(nb_operators, len_result, len_source).argmax(-1), - num_classes=len_source, + self.train_input, self.train_pred_masks = self.tensorize( + pairs[:nb_train_samples] ) + self.test_input, self.test_pred_masks = self.tensorize(pairs[nb_train_samples:]) - def generate_sequences(self, nb): - nb_operators = torch.randint(self.operators.size(0), (nb,)) - operators = self.operators[nb_operators] - nb_operators = ( - nb_operators[:, None] - // 10 ** torch.arange(self.len_nb_operator - 1, -1, -1) - ) % 10 - marker1 = torch.full((nb, 1), 10) - # source = torch.randint(10, (nb, self.len_source)) - source = torch.rand(nb, 10).sort(dim=1).indices[:, : self.len_source] - marker2 = torch.full((nb, 1), 11) - result = operators.bmm(source[:, :, None]).squeeze(-1) - sequences = torch.cat((nb_operators, marker1, source, marker2, result), 1) - ar_mask = (sequences == 11).long() - ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1) - return sequences, ar_mask - - def seq2str(self, seq): - return "".join("0123456789|>"[x.item()] for x in seq) - - -class ProblemLevel2(Problem): - def __init__(self, len_source=5, len_result=8): - self.len_source = len_source - self.len_result = len_result - - def generate_sequences(self, nb): - operators = F.one_hot( - torch.rand(nb, self.len_result, self.len_source).argmax(-1), - num_classes=self.len_source, - ) - source1 = torch.rand(nb, 10).sort(dim=1).indices[:, : self.len_source] - # source1 = torch.randint(10, (nb, self.len_source)) - marker1 = torch.full((nb, 1), 10) - result1 = operators.bmm(source1[:, :, None]).squeeze(-1) - marker2 = torch.full((nb, 1), 11) - source2 = torch.randint(10, (nb, self.len_source)) - marker3 = torch.full((nb, 1), 12) - result2 = operators.bmm(source2[:, :, None]).squeeze(-1) - - sequences = torch.cat( - (source1, marker1, result1, marker2, source2, marker3, result2), 1 - ) - ar_mask = (sequences == 12).long() - ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1) - return sequences, ar_mask + def batches(self, split="train", nb_to_use=-1, desc=None): + assert split in {"train", "test"} + input = self.train_input if split == "train" else self.test_input + if nb_to_use > 0: + input = input[:nb_to_use] + if desc is None: + desc = f"epoch-{split}" + for batch in tqdm.tqdm( + input.split(self.batch_size), dynamic_ncols=True, desc=desc + ): + yield self.trim(batch).to(self.device) - def seq2str(self, seq): - return "".join("0123456789>|~"[x.item()] for x in seq) + def vocabulary_size(self): + return len(self.char2id) + def tensor2str(self, t): + return ["".join([self.id2char[x.item()] for x in s]) for s in t] -#################### + def produce_results( + self, n_epoch, model, result_dir, logger, deterministic_synthesis + ): + correct = self.trim(self.test_input[:1000]).to(self.device) + result = correct.clone() + pred_mask = self.test_pred_masks[:1000, : result.size(1)].to(self.device) + ar_mask = (pred_mask > 0).long() + result *= 1 - ar_mask # paraaaaanoiaaaaaaa + logger(f"----------------------------------------------------------") -class ProblemAddition(Problem): - def __init__(self, nb_digits=10, zero_padded=False, inverted_result=False): - self.nb_digits = nb_digits - self.zero_padded = zero_padded - self.inverted_result = inverted_result - self.char2id = dict([(c, n) for n, c in enumerate("0123456789+=$")]) - self.id2char = dict([(n, c) for c, n in self.char2id.items()]) + for e in self.tensor2str(result[:10]): + logger(f"test_before {e}") - def tensorize(self, strings): - len_max = max([len(x) for x in strings]) - return torch.cat( - [ - torch.tensor( - [ - [self.char2id[c] for c in s + "$" * (len_max - len(s))] - for s in strings - ] - ) - ], - 0, + masked_inplace_autoregression( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis, + device=self.device, ) - def generate_sequences(self, nb): - sequences = [] - for k in range(nb): - a, b = torch.randint(10**self.nb_digits, (2,)) - c = a + b - a, b, c = str(a.item()), str(b.item()), str(c.item()) - if self.zero_padded: - a = "0" * (self.nb_digits - len(a)) + a - b = "0" * (self.nb_digits - len(b)) + b - c = "0" * (self.nb_digits + 1 - len(c)) + c - if self.inverted_result: - c = c[::-1] - sequences.append(f"{a}+{b}={c}$") - - sequences = self.tensorize(sequences) - ar_mask = (sequences == self.char2id["="]).long() - ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1) - return sequences, ar_mask + logger(f"----------------------------------------------------------") - def seq2str(self, seq): - return "".join(self.id2char[x.item()] for x in seq) + for e, c in zip(self.tensor2str(result[:10]), self.tensor2str(correct[:10])): + logger(f"test_after {e}") + logger(f"correct {c}") + logger(f"----------------------------------------------------------") -# class ProblemUnion(Problem): -# problems = [ProblemByheart()] -# nb_common_codes = 100 + err_mask = (pred_mask == 2).long() + nb_total = err_mask.sum().item() + nb_correct = ((correct == result).long() * err_mask).sum().item() -# def generate_sequences(nb_samples): -# problem_indexes = torch.randint(len(problems), (nb_samples,)) -# nb_samples_per_problem = torch.one_hot(problem_indexes).sum(0) -# print(f"{nb_samples_per_problem}") -# all_seq = [] -# for nb, p in zip(nb_samples_per_problem, problems): -# all_seq.append(p.generate_sequences(nb_samples_per_problem[nb])) -# return all_seq + logger(f"test_performance {n_epoch} {nb_total=} {nb_correct=}") + logger(f"main_test_accuracy {n_epoch} {nb_correct / nb_total}") -# for strain, stest in zip(train_seq, test_seq): -# s = torch.cat((strain, stest), 0) #################### +import problems + class SandBox(Task): def __init__( @@ -253,13 +240,25 @@ class SandBox(Task): self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 # A bit of paranoia never hurts - assert ( - self.nb_codes <= max_nb_codes - and self.train_input.min() >= 0 - and self.test_input.min() >= 0 - and tuple(self.train_ar_mask.unique()) == (0, 1) - and tuple(self.test_ar_mask.unique()) == (0, 1) - ) + assert self.nb_codes <= max_nb_codes + assert self.train_input.min() >= 0 + assert self.test_input.min() >= 0 + assert tuple(x.item() for x in self.train_ar_mask.unique()) in { + (0,), + (1,), + (0, 1), + } + assert tuple(x.item() for x in self.test_ar_mask.unique()) in { + (0,), + (1,), + (0, 1), + } + + if logger is not None: + for s, a in zip(self.train_input[:100], self.train_ar_mask[:100]): + logger(f"train_sequences {self.problem.seq2str(s)}") + a = "".join(["01"[x.item()] for x in a]) + logger(f" {a}") def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} @@ -293,17 +292,24 @@ class SandBox(Task): device=self.device, ) + log_ground_truth = ar_mask.min() == 0 + if logger is not None: for sp, st in zip(result[:10], input[:10]): logger( f"test_sequences {n_epoch} prediction {self.problem.seq2str(sp)}" ) - logger( - f" {n_epoch} ground truth {self.problem.seq2str(st)}" - ) + if log_ground_truth: + logger( + f" {n_epoch} ground truth {self.problem.seq2str(st)}" + ) + + nb_total, nb_correct = self.problem.compute_nb_correct( + input, ar_mask, result + ) - nb_total = ar_mask.sum().item() - nb_correct = ((result == input).long() * ar_mask).sum().item() + # nb_total = ar_mask.sum().item() + # nb_correct = ((result == input).long() * ar_mask).sum().item() return nb_total, nb_correct @@ -323,6 +329,41 @@ class SandBox(Task): 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}%" ) + logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}") + + if save_attention_image is not None: + for k in range(10): + ns = torch.randint(self.test_input.size(0), (1,)).item() + input = self.test_input[ns : ns + 1].clone() + + with torch.autograd.no_grad(): + t = model.training + model.eval() + # model.record_attention(True) + model(BracketedSequence(input)) + model.train(t) + # ram = model.retrieve_attention() + # model.record_attention(False) + + # tokens_output = [c for c in self.problem.seq2str(input[0])] + # tokens_input = ["n/a"] + tokens_output[:-1] + # for n_head in range(ram[0].size(1)): + # filename = os.path.join( + # result_dir, f"sandbox_attention_{k}_h{n_head}.pdf" + # ) + # attention_matrices = [m[0, n_head] for m in ram] + # save_attention_image( + # filename, + # tokens_input, + # tokens_output, + # attention_matrices, + # k_top=10, + ##min_total_attention=0.9, + # token_gap=12, + # layer_gap=50, + # ) + # logger(f"wrote {filename}") + ###################################################################### @@ -478,6 +519,10 @@ class PicoCLVR(Task): f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%" ) + logger( + f"main_test_accuracy {n_epoch} {1-nb_missing_properties/nb_requested_properties}" + ) + ###################################################################### def produce_results( @@ -748,6 +793,8 @@ class Maze(Task): 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}%" ) + logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}") + if count is not None: proportion_optimal = count.diagonal().sum().float() / count.sum() logger(f"proportion_optimal_test {proportion_optimal*100:.02f}%") @@ -887,6 +934,8 @@ class Snake(Task): 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}%" ) + logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}") + ###################################################################### @@ -996,6 +1045,8 @@ class Stack(Task): 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}%" ) + logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}") + ############################################################## # Log a few generated sequences input = self.test_input[:10, : 12 * (1 + self.nb_digits)] @@ -1059,6 +1110,7 @@ class RPL(Task): max_input=9, prog_len=6, nb_runs=5, + no_prog=False, logger=None, device=torch.device("cpu"), ): @@ -1066,6 +1118,7 @@ class RPL(Task): self.batch_size = batch_size self.device = device + self.no_prog = no_prog train_sequences = [ rpl.generate( @@ -1100,13 +1153,43 @@ class RPL(Task): self.id2token = dict([(n, c) for c, n in self.token2id.items()]) self.t_nul = self.token2id[""] - self.t_prog = self.token2id[""] - self.t_input = self.token2id[""] - self.t_output = self.token2id[""] + self.t_input = self.token2id[""] + self.t_output = self.token2id[""] + self.t_prog = self.token2id[""] + self.t_end = self.token2id[""] self.train_input = self.tensorize(train_sequences) self.test_input = self.tensorize(test_sequences) + if no_prog: + # Excise the program from every train and test example + k = torch.arange(self.train_input.size(1), device=self.train_input.device)[ + None, : + ] + p = ( + ((self.train_input == self.t_prog).long() * k) + .max(1, keepdim=True) + .values + ) + self.train_input = ( + self.train_input * (k <= p).long() + + self.t_end * (k == p + 1).long() + + self.t_nul * (k > p + 1).long() + ) + k = torch.arange(self.test_input.size(1), device=self.test_input.device)[ + None, : + ] + p = ( + ((self.test_input == self.t_prog).long() * k) + .max(1, keepdim=True) + .values + ) + self.test_input = ( + self.test_input * (k <= p).long() + + self.t_end * (k == p + 1).long() + + self.t_nul * (k > p + 1).long() + ) + if logger is not None: logger(f"value_max {val_max}") for x in self.train_input[:25]: @@ -1154,13 +1237,13 @@ class RPL(Task): ) sum_nb_total, sum_nb_errors = 0, 0 - for x, y in zip(input, result): - seq = [self.id2token[i.item()] for i in y] + for one_input, one_result in zip(input, result): + seq = [self.id2token[i.item()] for i in one_result] nb_total, nb_errors, prog, stacks = rpl.compute_nb_errors(seq) sum_nb_total += 1 sum_nb_errors += 0 if nb_errors == 0 else 1 if nb_to_log > 0: - gt_seq = [self.id2token[i.item()] for i in x] + gt_seq = [self.id2token[i.item()] for i in one_input] _, _, gt_prog, _ = rpl.compute_nb_errors(gt_seq) gt_prog = " ".join([str(x) for x in gt_prog]) prog = " ".join([str(x) for x in prog]) @@ -1201,14 +1284,20 @@ class RPL(Task): ) sum_nb_total, sum_nb_errors = 0, 0 - for x, y, i, j in zip(input, result, last_output_idx, first_prog_idx): - seq = [self.id2token[i.item()] for i in y] + for one_input, one_result, i, j in zip( + input, result, last_output_idx, first_prog_idx + ): + seq = [self.id2token[i.item()] for i in one_result] sum_nb_total += 1 - correct = (x - y).abs().max() == 0 + correct = (one_input - one_result).abs().max() == 0 sum_nb_errors += 0 if correct else 1 if nb_to_log > 0: - result_stack = [self.id2token[i.item()] for i in y[i : j + 1]] - target_stack = [self.id2token[i.item()] for i in x[i : j + 1]] + result_stack = [ + self.id2token[i.item()] for i in one_result[i : j + 1] + ] + target_stack = [ + self.id2token[i.item()] for i in one_input[i : j + 1] + ] comment = "*" if correct else "-" result_stack = " ".join([str(x) for x in result_stack]) target_stack = " ".join([str(x) for x in target_stack]) @@ -1221,13 +1310,16 @@ class RPL(Task): # -------------------------------------------------------------------- - test_nb_total, test_nb_errors = compute_nb_errors_prog( - self.test_input[:1000].to(self.device), nb_to_log=10 - ) + if not self.no_prog: + test_nb_total, test_nb_errors = compute_nb_errors_prog( + self.test_input[:1000].to(self.device), nb_to_log=10 + ) - logger( - 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}%" - ) + logger( + 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}%" + ) + + logger(f"main_test_accuracy {n_epoch} {1-test_nb_errors/test_nb_total}") test_nb_total, test_nb_errors = compute_nb_errors_output( self.test_input[:1000].to(self.device), nb_to_log=10 @@ -1237,6 +1329,42 @@ class RPL(Task): 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}%" ) + if save_attention_image is None: + logger("no save_attention_image (is pycairo installed?)") + else: + ns = torch.randint(self.test_input.size(0), (1,)).item() + input = self.test_input[ns : ns + 1].clone() + last = (input != self.t_nul).max(0).values.nonzero().max() + 3 + input = input[:, :last].to(self.device) + + with torch.autograd.no_grad(): + t = model.training + model.eval() + model.record_attention(True) + model(BracketedSequence(input)) + model.train(t) + ram = model.retrieve_attention() + model.record_attention(False) + + tokens_output = [self.id2token[i.item()] for i in input[0]] + tokens_input = ["n/a"] + tokens_output[:-1] + for n_head in range(ram[0].size(1)): + filename = os.path.join( + result_dir, f"rpl_attention_{n_epoch}_h{n_head}.pdf" + ) + attention_matrices = [m[0, n_head] for m in ram] + save_attention_image( + filename, + tokens_input, + tokens_output, + attention_matrices, + k_top=10, + # min_total_attention=0.9, + token_gap=12, + layer_gap=50, + ) + logger(f"wrote {filename}") + ###################################################################### @@ -1391,6 +1519,8 @@ class Expr(Task): 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}%" ) + logger(f"main_test_accuracy {n_epoch} {test_nb_correct/test_nb_total}") + nb_total = test_nb_delta.sum() + test_nb_missed for d in range(test_nb_delta.size(0)): logger( @@ -1437,95 +1567,126 @@ class Expr(Task): ###################################################################### -import world +import grid + + +class Grid(Task): + # Make a tensor from a list of strings + def str2tensor(self, descr): + token_descr = [s.strip().split(" ") for s in descr] + l = max([len(s) for s in token_descr]) + token_descr = [s + ["#"] * (l - len(s)) for s in token_descr] + id_descr = [[self.token2id[u] for u in s] for s in token_descr] + return torch.tensor(id_descr, device=self.device) + + # Make a list of strings from a tensor + def tensor2str(self, x): + return [" ".join([self.id2token[t.item()] for t in r]) for r in x] + + # trim all the tensors in the tuple z to remove as much token from + # left and right in the first tensor. If z is a tuple, all its + # elements are trimed according to the triming for the first + def trim(self, z, token="#"): + n = self.token2id[token] + if type(z) == tuple: + x = z[0] + i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0) + a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() + return tuple([t[:, a:b] for t in z]) + else: + i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0) + a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() + return z[:, a:b] + ###################### -class World(Task): def __init__( self, nb_train_samples, nb_test_samples, batch_size, - vqae_nb_epochs, + size, + fraction_play=0.0, logger=None, device=torch.device("cpu"), - device_storage=torch.device("cpu"), ): super().__init__() - self.batch_size = batch_size self.device = device + self.batch_size = batch_size + self.grid_factory = grid.GridFactory(size=size) + self.fraction_play = fraction_play - ( - train_frames, - train_action_seq, - test_frames, - test_action_seq, - self.frame2seq, - self.seq2frame, - ) = world.create_data_and_processors( - nb_train_samples, - nb_test_samples, - mode="first_last", - nb_steps=30, - nb_epochs=vqae_nb_epochs, - logger=logger, - device=device, - device_storage=device_storage, - ) - - train_frame_seq = self.frame2seq(train_frames).to(device_storage) - test_frame_seq = self.frame2seq(test_frames).to(device_storage) + if logger is not None: + logger( + f"generating {nb_train_samples+nb_test_samples} samples (can take some time)" + ) - nb_frame_codes = max(train_frame_seq.max(), test_frame_seq.max()) + 1 - nb_action_codes = max(train_action_seq.max(), test_action_seq.max()) + 1 + self.train_descr = self.grid_factory.generate_samples( + nb=nb_train_samples, + fraction_play=fraction_play, + progress_bar=lambda r: tqdm.tqdm(r), + ) - self.len_frame_seq = train_frame_seq.size(1) - self.len_action_seq = train_action_seq.size(1) - self.nb_codes = nb_frame_codes + nb_action_codes + self.test_descr = self.grid_factory.generate_samples( + nb=nb_test_samples, fraction_play=0.0, progress_bar=lambda r: tqdm.tqdm(r) + ) - train_frame_seq = train_frame_seq.reshape(train_frame_seq.size(0) // 2, 2, -1) + if fraction_play > 0: + self.play_descr = self.grid_factory.generate_samples( + nb=25, fraction_play=1.0, progress_bar=lambda r: tqdm.tqdm(r) + ) + else: + self.play_descr = [] - train_action_seq += nb_frame_codes - self.train_input = torch.cat( - (train_frame_seq[:, 0, :], train_action_seq, train_frame_seq[:, 1, :]), 1 - ) + # Build the tokenizer + tokens = set() + for d in [self.train_descr, self.test_descr, self.play_descr]: + for s in d: + for t in s.strip().split(" "): + tokens.add(t) + # make this set a sorted list to get the same tensors given + # the same descr + tokens = list(tokens) + tokens.sort() + tokens = ["#"] + tokens + self.token2id = dict([(t, n) for n, t in enumerate(tokens)]) + self.id2token = dict([(n, t) for n, t in enumerate(tokens)]) + self.t_nul = self.token2id["#"] + self.t_true = self.token2id["true"] + self.t_false = self.token2id["false"] + self.t_pipe = self.token2id["|"] - test_frame_seq = test_frame_seq.reshape(test_frame_seq.size(0) // 2, 2, -1) - test_action_seq += nb_frame_codes - self.test_input = torch.cat( - (test_frame_seq[:, 0, :], test_action_seq, test_frame_seq[:, 1, :]), 1 + # Tokenize the train and test sets + self.train_input = self.str2tensor(self.train_descr) + self.test_input = self.str2tensor(self.test_descr) + self.play_input = ( + None if len(self.play_descr) == 0 else self.str2tensor(self.play_descr) ) - def batches(self, split="train", nb_to_use=-1, desc=None): + def batches(self, split="train"): assert split in {"train", "test"} input = self.train_input if split == "train" else self.test_input - if nb_to_use > 0: - input = input[:nb_to_use] - if desc is None: - desc = f"epoch-{split}" for batch in tqdm.tqdm( - input.split(self.batch_size), dynamic_ncols=True, desc=desc + input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}" ): - yield batch.to(self.device) + yield self.trim(batch) def vocabulary_size(self): - return self.nb_codes + return len(self.token2id) def produce_results( self, n_epoch, model, result_dir, logger, deterministic_synthesis ): - k = torch.arange( - 2 * self.len_frame_seq + self.len_action_seq, device=self.device - )[None, :] + correct = self.test_input[:1000] + result = correct.clone() + ar_mask = torch.logical_or(result == self.t_true, result == self.t_false).long() + result *= 1 - ar_mask # paraaaaanoiaaaaaaa - input = self.test_input[:64].to(self.device) - result = input.clone() + logger(f"----------------------------------------------------------") - ar_mask = ( - (k >= self.len_frame_seq + self.len_action_seq).long().expand_as(result) - ) - result *= 1 - ar_mask + for e in self.tensor2str(result[:10]): + logger(f"test_before {e}") masked_inplace_autoregression( model, @@ -1536,25 +1697,142 @@ class World(Task): device=self.device, ) - seq_start = input[:, : self.len_frame_seq] - seq_end = input[:, self.len_frame_seq + self.len_action_seq :] - seq_predicted = result[:, self.len_frame_seq + self.len_action_seq :] + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): + logger(f"test_after {e}") + + logger(f"----------------------------------------------------------") + + nb_total = ar_mask.sum().item() + nb_correct = ((correct == result).long() * ar_mask).sum().item() + + logger(f"test_performance {n_epoch} {nb_total=} {nb_correct=}") + logger(f"main_test_accuracy {n_epoch} {nb_correct / nb_total}") + + if self.play_input is not None: + result = self.play_input.clone() + ar_mask = (result == self.t_pipe).long().cumsum(dim=1).clamp(max=1) + result *= 1 - ar_mask # paraaaaanoiaaaaaaa + + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): + logger(f"play_before {e}") + + masked_inplace_autoregression( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis, + device=self.device, + ) + + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): + logger(f"play_after {e}") + + logger(f"----------------------------------------------------------") + + +###################################################################### + +import qmlp + + +class QMLP(Task): + ###################### + + def __init__( + self, + nb_train_samples, + nb_test_samples, + batch_size, + result_dir, + logger=None, + device=torch.device("cpu"), + ): + super().__init__() + + self.device = device + self.batch_size = batch_size + self.nb_samples_per_mlp = 256 + + if logger is not None: + logger( + f"generating {nb_train_samples+nb_test_samples} samples (can take some time)" + ) - result = torch.cat( - (seq_start[:, None, :], seq_end[:, None, :], seq_predicted[:, None, :]), 1 + seq, q_test_set, test_error = qmlp.generate_sequence_and_test_set( + nb_mlps=nb_train_samples + nb_test_samples, + nb_samples=self.nb_samples_per_mlp, + device=self.device, + batch_size=64, + nb_epochs=250, + nb_mlps_per_batch=1024, ) - result = result.reshape(-1, result.size(-1)) - frames = self.seq2frame(result) - image_name = os.path.join(result_dir, f"world_result_{n_epoch:04d}.png") - torchvision.utils.save_image( - frames.float() / (world.Box.nb_rgb_levels - 1), - image_name, - nrow=12, - padding=1, - pad_value=0.0, + self.train_input = seq[:nb_train_samples] + self.train_q_test_set = q_test_set[:nb_train_samples] + self.train_ref_test_errors = test_error[:nb_train_samples] + self.test_input = seq[nb_train_samples:] + self.test_q_test_set = q_test_set[nb_train_samples:] + self.test_ref_test_errors = test_error[nb_train_samples:] + + filename = os.path.join(result_dir, f"train_errors_ref.dat") + with open(filename, "w") as f: + for e in self.train_ref_test_errors: + f.write(f"{e}\n") + + filename = os.path.join(result_dir, f"test_errors_ref.dat") + with open(filename, "w") as f: + for e in self.test_ref_test_errors: + f.write(f"{e}\n") + + self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 + + def batches(self, split="train"): + assert split in {"train", "test"} + input = self.train_input if split == "train" else self.test_input + for batch in tqdm.tqdm( + input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}" + ): + yield batch + + def vocabulary_size(self): + return self.nb_codes + + def produce_results( + self, n_epoch, model, result_dir, logger, deterministic_synthesis + ): + correct = self.test_input[:1000] + result = correct.clone() + ar_mask = ( + torch.arange(result.size(1), device=result.device) + > self.nb_samples_per_mlp * 3 + 1 + ).long()[None, :] + ar_mask = ar_mask.expand_as(result) + result *= 1 - ar_mask # paraaaaanoiaaaaaaa + + masked_inplace_autoregression( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis, + device=self.device, ) - logger(f"wrote {image_name}") + + q_train_set = result[:, : self.nb_samples_per_mlp * 3] + q_params = result[:, self.nb_samples_per_mlp * 3 + 1 :] + error_test = qmlp.evaluate_q_params(q_params, self.test_q_test_set) + + filename = os.path.join(result_dir, f"test_errors_{n_epoch:04d}.dat") + with open(filename, "w") as f: + for e in error_test: + f.write(f"{e}\n") ######################################################################