From a85b2f79ae551da3454daa7b90554be5ca3c5bf6 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sun, 24 Mar 2024 09:57:15 +0100 Subject: [PATCH] Update. --- escape.py | 15 +++++++++------ tasks.py | 31 ++++++++++++++++++++++++++++--- 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/escape.py b/escape.py index 7fcafaa..aa87cdd 100755 --- a/escape.py +++ b/escape.py @@ -150,14 +150,16 @@ def episodes2str(states, actions, rewards, unicode=False, ansi_colors=False): result = hline for n in range(states.size(0)): + + def state_symbol(v): + v = v.item() + return "?" if v < 0 or v >= len(symbols) else symbols[v] + for i in range(states.size(2)): result += ( vert + vert.join( - [ - "".join([symbols[v.item()] for v in row]) - for row in states[n, :, i] - ] + ["".join([state_symbol(v) for v in row]) for row in states[n, :, i]] ) + vert + "\n" @@ -166,8 +168,9 @@ def episodes2str(states, actions, rewards, unicode=False, ansi_colors=False): result += (vert + thin_hori * states.size(-1)) * states.size(1) + vert + "\n" def status_bar(a, r): - a = "ISNEW"[a.item()] - r = "" if r == 0 else f"{r.item()}" + a = a.item() + a = "ISNEW"[a] if a >= 0 and a < 5 else "?" + r = "?" if r < -1 or r > 2 else ("" if r == 0 else f"{r.item()}") return a + " " * (states.size(-1) - len(a) - len(r)) + r result += ( diff --git a/tasks.py b/tasks.py index 8b22cad..dba6e13 100755 --- a/tasks.py +++ b/tasks.py @@ -1881,13 +1881,15 @@ class Escape(Task): 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 ) seq = escape.episodes2seq(states, actions, rewards) - self.train_input = seq[:nb_train_samples] - self.test_input = seq[nb_train_samples:] + 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 @@ -1915,7 +1917,30 @@ class Escape(Task): def produce_results( self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000 ): - pass + result = self.test_input[:100].clone() + ar_mask = ( + torch.arange(result.size(1), device=result.device) + > self.height * self.width + 2 + ).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, + ) + + s, a, r = escape.seq2episodes(result, self.height, self.width) + str = escape.episodes2str(s, a, r, 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}") ###################################################################### -- 2.39.5