- logger(
- 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(
- f"error_value {n_epoch} delta {d} {test_nb_delta[d]} {test_nb_delta[d]*100/nb_total:.02f}%"
- )
- logger(
- f"error_value {n_epoch} missed {test_nb_missed} {test_nb_missed*100/nb_total:.02f}%"
- )
-
- ##############################################################
- # Log a few generated sequences
- if input_file is None:
- input = self.test_input[:10]
- else:
- with open(input_file, "r") as f:
- sequences = [e.strip() for e in f.readlines()]
- sequences = [s + " " + "#" * 50 for s in sequences]
- input = self.tensorize(sequences)
-
- result = input.clone()
- s = (result == self.space).long()
- ar_mask = (s.cumsum(dim=1) - s).clamp(min=0, max=1)
- result = (1 - ar_mask) * result + ar_mask * self.filler
-
- for n in range(result.size(0)):
- logger(f"test_before {self.seq2str(result[n])}")
-
- masked_inplace_autoregression(
- model,
- self.batch_size,
- result,
- ar_mask,
- deterministic_synthesis,
- device=self.device,
- )
-
- correct = (1 - ar_mask) * self.space + ar_mask * input
- for n in range(result.size(0)):
- comment = "GOOD" if (result[n] - input[n]).abs().max() == 0 else ""
- logger(f"test_after {self.seq2str(result[n])} {comment}")
- logger(f"truth {self.seq2str(correct[n])}")
- ##############################################################
-
-
-######################################################################
-
-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]
-
- ######################
-
- def __init__(
- self,
- nb_train_samples,
- nb_test_samples,
- batch_size,
- size,
- fraction_play=0.0,
- logger=None,
- device=torch.device("cpu"),
- ):
- super().__init__()
-
- self.device = device
- self.batch_size = batch_size
- self.grid_factory = grid.GridFactory(size=size)
- self.fraction_play = fraction_play
-
- if logger is not None:
- logger(
- f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
- )
-
- self.train_descr = self.grid_factory.generate_samples(
- nb=nb_train_samples,
- fraction_play=fraction_play,
- progress_bar=lambda r: tqdm.tqdm(r),
- )
-
- self.test_descr = self.grid_factory.generate_samples(
- nb=nb_test_samples, fraction_play=0.0, progress_bar=lambda r: tqdm.tqdm(r)
- )
-
- 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 = []
-
- # 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["|"]
-
- # 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"):
- 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 self.trim(batch)
-
- def vocabulary_size(self):
- return len(self.token2id)
-
- def produce_results(
- self, n_epoch, model, result_dir, logger, deterministic_synthesis
- ):
- 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
-
- logger(f"----------------------------------------------------------")
-
- for e in self.tensor2str(result[:10]):
- logger(f"test_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"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)"
- )
-
- 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,
- )
-
- 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,
- )
-
- 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")
-
-
-######################################################################
-
-import escape
-
-
-class Escape(Task):
- def __init__(
- self,
- nb_train_samples,
- nb_test_samples,
- batch_size,
- height,
- width,
- T,
- nb_walls,
- logger=None,
- device=torch.device("cpu"),
- ):
- super().__init__()
-
- self.batch_size = batch_size
- self.device = device
- self.height = height
- self.width = width
-
- states, actions, rewards = escape.generate_episodes(
- nb_train_samples + nb_test_samples, height, width, T, nb_walls
- )
- seq = escape.episodes2seq(states, actions, rewards, lookahead_delta=T)
- # seq = seq[:, seq.size(1) // 3 : 2 * seq.size(1) // 3]
- self.train_input = seq[:nb_train_samples].to(self.device)
- self.test_input = seq[nb_train_samples:].to(self.device)
-
- self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
-
- 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 batch
-
- def vocabulary_size(self):
- return self.nb_codes
-
- def thinking_autoregression(
- self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
- ):
- result = self.test_input[:250].clone()
- t = torch.arange(result.size(1), device=result.device)[None, :]
-
- state_len = self.height * self.width
- index_action = state_len
- index_reward = state_len + 1
- index_lookahead_reward = state_len + 2
- it_len = state_len + 3 # state / action / reward / lookahead_reward
-
- def ar(result, ar_mask, logit_biases=None):
- ar_mask = ar_mask.expand_as(result)
- result *= 1 - ar_mask
- masked_inplace_autoregression(
- model,
- self.batch_size,
- result,
- ar_mask,
- deterministic_synthesis=deterministic_synthesis,
- logit_biases=logit_biases,
- device=self.device,
- progress_bar_desc=None,
- )
-
- # Generate iteration after iteration
-
- optimistic_bias = result.new_zeros(self.nb_codes, device=result.device)
- optimistic_bias[(-1) + escape.first_lookahead_rewards_code + 1] = math.log(1e-1)
- optimistic_bias[(1) + escape.first_lookahead_rewards_code + 1] = math.log(1e1)
-
- for u in tqdm.tqdm(
- range(it_len, result.size(1) - it_len + 1, it_len), desc="thinking"
- ):
- # Generate the lookahead_reward pessimistically
- ar_mask = (t < u).long() * (t % it_len == index_lookahead_reward).long()
- ar(result, ar_mask, logit_biases=-optimistic_bias)
-
- # Generate the state
- ar_mask = (t >= u).long() * (t < u + state_len).long()
- ar(result, ar_mask)
-
- # Generate the lookahead_reward optimistically
- ar_mask = (t < u).long() * (t % it_len == index_lookahead_reward).long()
- ar(result, ar_mask, logit_biases=optimistic_bias)
-
- # Generate the action and reward
- ar_mask = (t >= u + index_action).long() * (t <= u + index_reward).long()
- ar(result, ar_mask)
-
- # Saving the generated sequences
-
- s, a, r, lr = escape.seq2episodes(
- result, self.height, self.width, lookahead=True
- )
- str = escape.episodes2str(
- s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
- )
-
- filename = os.path.join(result_dir, f"test_thinking_seq_{n_epoch:04d}.txt")
- with open(filename, "w") as f:
- f.write(str)
- logger(f"wrote {filename}")
-
- def produce_results(
- self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
- ):
- result = self.test_input[:250].clone()
-
- # Saving the ground truth
-
- s, a, r, lr = escape.seq2episodes(
- result, self.height, self.width, lookahead=True
- )
- str = escape.episodes2str(
- s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
- )
-
- filename = os.path.join(result_dir, f"test_true_seq_{n_epoch:04d}.txt")
- with open(filename, "w") as f:
- f.write(str)
- logger(f"wrote {filename}")
-
- # Re-generating from the first frame
-
- ar_mask = (
- torch.arange(result.size(1), device=result.device)
- >= self.height * self.width + 3
- ).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,
- )
-
- # Saving the generated sequences
-
- s, a, r, lr = escape.seq2episodes(
- result, self.height, self.width, lookahead=True
- )
- str = escape.episodes2str(
- s, a, r, lookahead_rewards=lr, unicode=True, ansi_colors=True
- )
-
- filename = os.path.join(result_dir, f"test_seq_{n_epoch:04d}.txt")
- with open(filename, "w") as f:
- f.write(str)
- logger(f"wrote {filename}")
-
- self.thinking_autoregression(
- n_epoch, model, result_dir, logger, deterministic_synthesis, nmax
- )
-
-
-######################################################################