("cyan", [0, 255, 255]),
("violet", [255, 0, 255]),
("lightgreen", [192, 255, 192]),
- ("pink", [255, 192, 192]),
+ ("brown", [165, 42, 42]),
("lightblue", [192, 192, 255]),
- ("gray", [192, 192, 192]),
+ ("gray", [128, 128, 128]),
]
def __init__(self, device=torch.device("cpu")):
X[i1:i2, j1:j2] = c[n]
f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
- def task_move(self, A, f_A, B, f_B):
- di, dj = torch.randint(2, (2,)) * 2 - 1
+ def task_translate(self, A, f_A, B, f_B):
+ di, dj = torch.randint(3, (2,)) - 1
nb_rec = 3
c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
for X, f_X in [(A, f_A), (B, f_B)]:
if n < nb_rec - 1:
f_X[i1, j1] = c[-1]
+ def task_count(self, A, f_A, B, f_B):
+ N = torch.randint(4, (1,)) + 2
+ c = torch.randperm(len(self.colors) - 1)[:N] + 1
+
+ for X, f_X in [(A, f_A), (B, f_B)]:
+
+ def contact(i, j, q):
+ nq, nq_diag = 0, 0
+ no = 0
+
+ for ii, jj in [
+ (i - 1, j - 1),
+ (i - 1, j),
+ (i - 1, j + 1),
+ (i, j - 1),
+ (i, j + 1),
+ (i + 1, j - 1),
+ (i + 1, j),
+ (i + 1, j + 1),
+ ]:
+ if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
+ if X[ii, jj] != 0 and X[ii, jj] != q:
+ no += 1
+
+ for ii, jj in [
+ (i - 1, j - 1),
+ (i - 1, j + 1),
+ (i + 1, j - 1),
+ (i + 1, j + 1),
+ ]:
+ if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
+ if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
+ nq_diag += 1
+
+ for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
+ if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
+ if X[ii, jj] == q:
+ nq += 1
+
+ return no, nq, nq_diag
+
+ nb = torch.zeros(N, dtype=torch.int64)
+ q = torch.randint(N, (self.height * self.width,))
+ k = torch.randperm(self.height * self.width)
+ for p in range(self.height * self.width):
+ i, j = k[p] % self.height, k[p] // self.height
+ no, nq, nq_diag = contact(i, j, c[q[p]])
+ if no == 0 and nq_diag == 0:
+ if nq == 0:
+ if nb[q[p]] < self.width:
+ X[i, j] = c[q[p]]
+ nb[q[p]] += 1
+ if nq == 1:
+ X[i, j] = c[q[p]]
+
+ for n in range(N):
+ for j in range(nb[n]):
+ f_X[n, j] = c[n]
+
+ def task_count_(self, A, f_A, B, f_B):
+ N = torch.randint(3, (1,)) + 1
+ c = torch.randperm(len(self.colors) - 1)[:N] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ nb = torch.randint(self.width, (3,)) + 1
+ k = torch.randperm(self.height * self.width)[: nb.sum()]
+ p = 0
+ for n in range(N):
+ for m in range(nb[n]):
+ i, j = k[p] % self.height, k[p] // self.height
+ X[i, j] = c[n]
+ f_X[n, m] = c[n]
+ p += 1
+
######################################################################
def generate_prompts_and_answers(self, nb, device="cpu"):
tasks = [
self.task_replace_color,
- self.task_move,
+ self.task_translate,
self.task_grow,
self.task_color_grow,
self.task_frame,
self.task_detect,
+ self.task_count,
]
prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
delay = time.perf_counter() - start_time
print(f"{prompts.size(0)/delay:02f} seq/s")
- # predicted_prompts = torch.rand(prompts.size(0)) < 0.5
- # predicted_answers = torch.logical_not(predicted_prompts)
+ predicted_prompts = torch.rand(prompts.size(0)) < 0.5
+ predicted_answers = torch.logical_not(predicted_prompts)
reasoning.save_quizzes(
"/tmp",
prompts[:64],
answers[:64],
# You can add a bool to put a frame around the predicted parts
- # predicted_prompts, predicted_answers
+ # predicted_prompts[:64],
+ # predicted_answers[:64],
)