("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")):
######################################################################
def frame2img(self, x, scale=15):
+ x = x.reshape(x.size(0), self.height, -1)
+ m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
+ x = self.colors[x * m].permute(0, 3, 1, 2)
+ s = x.shape
+ x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
+ x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
+
+ x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
+ x[:, :, torch.arange(0, x.size(2), scale), :] = 0
+ x = x[:, :, 1:, 1:]
+
+ for n in range(m.size(0)):
+ for i in range(m.size(1)):
+ for j in range(m.size(2)):
+ if m[n, i, j] == 0:
+ for k in range(2, scale - 2):
+ for l in [0, 1]:
+ x[n, :, i * scale + k, j * scale + k - l] = 0
+ x[
+ n, :, i * scale + scale - 1 - k, j * scale + k - l
+ ] = 0
+
+ return x
+
+ def frame2img_(self, x, scale=15):
x = x.reshape(x.size(0), self.height, -1)
x = self.colors[x].permute(0, 3, 1, 2)
s = x.shape
def nb_token_values(self):
return len(self.colors)
- def rec_coo(self, x, n, min_height=3, min_width=3):
- K = 3
- N = 4000
+ # That's quite a tensorial spaghetti mess to sample
+ # non-overlapping rectangles quickly, but made the generation of
+ # 100k samples go from 1h50 with a lame pure python code to 3min30s
+ # with this one.
+ def rec_coo(self, nb_rec, min_height=3, min_width=3):
+ nb_trials = 200
while True:
v = (
(
- torch.rand(N * K, self.height + 1, device=self.device)
+ torch.rand(nb_trials * nb_rec, self.height + 1, device=self.device)
.sort(dim=-1)
.indices
< 2
h = (
(
- torch.rand(N * K, self.width + 1, device=self.device)
+ torch.rand(nb_trials * nb_rec, self.width + 1, device=self.device)
.sort(dim=-1)
.indices
< 2
)
v, h = v[i], h[i]
- v = v[: v.size(0) - v.size(0) % K]
- h = h[: h.size(0) - h.size(0) % K]
- v = v.reshape(v.size(0) // K, K, -1)
- h = h.reshape(h.size(0) // K, K, -1)
+ v = v[: v.size(0) - v.size(0) % nb_rec]
+ h = h[: h.size(0) - h.size(0) % nb_rec]
+ v = v.reshape(v.size(0) // nb_rec, nb_rec, -1)
+ h = h.reshape(h.size(0) // nb_rec, nb_rec, -1)
r = v[:, :, :, None] * h[:, :, None, :]
######################################################################
def task_replace_color(self, A, f_A, B, f_B):
- N = 3
- c = torch.randperm(len(self.colors) - 1)[: N + 1] + 1
+ nb_rec = 3
+ c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
for X, f_X in [(A, f_A), (B, f_B)]:
- r = self.rec_coo(X, N)
- for n in range(N):
+ r = self.rec_coo(nb_rec)
+ for n in range(nb_rec):
i1, j1, i2, j2 = r[n]
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
- N = 3
- c = torch.randperm(len(self.colors) - 1)[:N] + 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)]:
while True:
- r = self.rec_coo(X, N)
- i1, j1, i2, j2 = r[N - 1]
+ r = self.rec_coo(nb_rec)
+ i1, j1, i2, j2 = r[nb_rec - 1]
if (
i1 + di >= 0
and i2 + di < X.size(0)
):
break
- for n in range(N):
+ for n in range(nb_rec):
i1, j1, i2, j2 = r[n]
X[i1:i2, j1:j2] = c[n]
- if n == N - 1:
+ if n == nb_rec - 1:
f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
else:
f_X[i1:i2, j1:j2] = c[n]
def task_grow(self, A, f_A, B, f_B):
di, dj = torch.randint(2, (2,)) * 2 - 1
- N = 3
- c = torch.randperm(len(self.colors) - 1)[:N] + 1
+ nb_rec = 3
+ c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
direction = torch.randint(2, (1,))
for X, f_X in [(A, f_A), (B, f_B)]:
while True:
- r = self.rec_coo(X, N)
- i1, j1, i2, j2 = r[N - 1]
+ r = self.rec_coo(nb_rec)
+ i1, j1, i2, j2 = r[nb_rec - 1]
if i1 + 3 < i2 and j1 + 3 < j2:
break
- for n in range(N):
+ for n in range(nb_rec):
i1, j1, i2, j2 = r[n]
- if n == N - 1:
+ if n == nb_rec - 1:
if direction == 0:
X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
f_X[i1:i2, j1:j2] = c[n]
def task_color_grow(self, A, f_A, B, f_B):
di, dj = torch.randint(2, (2,)) * 2 - 1
- N = 3
- c = torch.randperm(len(self.colors) - 1)[: 2 * N] + 1
- direction = torch.randint(2, (1,))
+ nb_rec = 3
+ c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
+ direction = torch.randint(4, (1,))
for X, f_X in [(A, f_A), (B, f_B)]:
- r = self.rec_coo(X, N)
- for n in range(N):
+ r = self.rec_coo(nb_rec)
+ for n in range(nb_rec):
i1, j1, i2, j2 = r[n]
- i = (i1 + i2) // 2
X[i1:i2, j1:j2] = c[2 * n]
- X[i : i + 1, j1:j2] = c[2 * n + 1]
f_X[i1:i2, j1:j2] = c[2 * n]
- if n == N - 1:
- f_X[i:i2, j1:j2] = c[2 * n + 1]
- else:
- f_X[i : i + 1, j1:j2] = c[2 * n + 1]
+ # Not my proudest moment
+ if direction == 0:
+ i = (i1 + i2) // 2
+ X[i : i + 1, j1:j2] = c[2 * n + 1]
+ if n == nb_rec - 1:
+ f_X[i:i2, j1:j2] = c[2 * n + 1]
+ else:
+ f_X[i : i + 1, j1:j2] = c[2 * n + 1]
+ elif direction == 1:
+ i = (i1 + i2 - 1) // 2
+ X[i : i + 1, j1:j2] = c[2 * n + 1]
+ if n == nb_rec - 1:
+ f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
+ else:
+ f_X[i : i + 1, j1:j2] = c[2 * n + 1]
+ elif direction == 2:
+ j = (j1 + j2) // 2
+ X[i1:i2, j : j + 1] = c[2 * n + 1]
+ if n == nb_rec - 1:
+ f_X[i1:i2, j:j2] = c[2 * n + 1]
+ else:
+ f_X[i1:i2, j : j + 1] = c[2 * n + 1]
+ elif direction == 3:
+ j = (j1 + j2 - 1) // 2
+ X[i1:i2, j : j + 1] = c[2 * n + 1]
+ if n == nb_rec - 1:
+ f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
+ else:
+ f_X[i1:i2, j : j + 1] = c[2 * n + 1]
def task_frame(self, A, f_A, B, f_B):
- N = 3
- c = torch.randperm(len(self.colors) - 1)[: N + 1] + 1
+ nb_rec = 3
+ c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
for X, f_X in [(A, f_A), (B, f_B)]:
- r = self.rec_coo(X, N)
- for n in range(N):
+ r = self.rec_coo(nb_rec)
+ for n in range(nb_rec):
i1, j1, i2, j2 = r[n]
X[i1:i2, j1:j2] = c[n]
f_X[i1:i2, j1:j2] = c[n]
- if n == N - 1:
+ if n == nb_rec - 1:
f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
def task_detect(self, A, f_A, B, f_B):
- N = 3
- c = torch.randperm(len(self.colors) - 1)[: N + 1] + 1
+ nb_rec = 3
+ c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
for X, f_X in [(A, f_A), (B, f_B)]:
- r = self.rec_coo(X, N)
- for n in range(N):
+ r = self.rec_coo(nb_rec)
+ for n in range(nb_rec):
i1, j1, i2, j2 = r[n]
X[i1:i2, j1:j2] = c[n]
- f_X[i1, j1] = c[-1]
+ 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_trajectory(self, A, f_A, B, f_B):
+ c = torch.randperm(len(self.colors) - 1)[:2] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ while True:
+ di, dj = torch.randint(7, (2,)) - 3
+ i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
+ if (
+ abs(di) + abs(dj) > 0
+ and i + 2 * di >= 0
+ and i + 2 * di < self.height
+ and j + 2 * dj >= 0
+ and j + 2 * dj < self.width
+ ):
+ break
+
+ k = 0
+ while (
+ i + k * di >= 0
+ and i + k * di < self.height
+ and j + k * dj >= 0
+ and j + k * dj < self.width
+ ):
+ if k < 2:
+ X[i + k * di, j + k * dj] = c[k]
+ f_X[i + k * di, j + k * dj] = c[min(k, 1)]
+ k += 1
+
+ def task_bounce(self, A, f_A, B, f_B):
+ c = torch.randperm(len(self.colors) - 1)[:3] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+
+ def free(i, j):
+ return (
+ i >= 0
+ and i < self.height
+ and j >= 0
+ and j < self.width
+ and f_X[i, j] == 0
+ )
+
+ while True:
+ f_X[...] = 0
+ X[...] = 0
+
+ for _ in range((self.height * self.width) // 10):
+ i, j = torch.randint(self.height, (1,)), torch.randint(
+ self.width, (1,)
+ )
+ X[i, j] = c[0]
+ f_X[i, j] = c[0]
+
+ while True:
+ di, dj = torch.randint(7, (2,)) - 3
+ if abs(di) + abs(dj) == 1:
+ break
+
+ i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
+
+ X[i, j] = c[1]
+ f_X[i, j] = c[1]
+ l = 0
+
+ while True:
+ l += 1
+ if free(i + di, j + dj):
+ pass
+ elif free(i - dj, j + di):
+ di, dj = -dj, di
+ if free(i + dj, j - di):
+ if torch.rand(1) < 0.5:
+ di, dj = -di, -dj
+ elif free(i + dj, j - di):
+ di, dj = dj, -di
+ else:
+ break
+
+ i, j = i + di, j + dj
+ f_X[i, j] = c[2]
+ if l <= 1:
+ X[i, j] = c[2]
+
+ if l >= self.width:
+ break
+
+ f_X[i, j] = c[1]
+ X[i, j] = c[1]
+
+ if l > 3:
+ break
######################################################################
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,
+ self.task_trajectory,
+ self.task_bounce,
]
- prompts = torch.zeros(
- nb, self.height, self.width * 3, dtype=torch.int64, device=self.device
- )
- answers = torch.zeros(
- nb, self.height, self.width, dtype=torch.int64, device=self.device
- )
+ prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
+ answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
w = self.width
for prompt, answer in tqdm.tqdm(
f_B = answer
task = tasks[torch.randint(len(tasks), (1,))]
task(A, f_A, B, f_B)
+
return prompts.flatten(1), answers.flatten(1)
def save_quizzes(
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",
"test",
- prompts[:36],
- answers[:36],
+ 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],
)
projects / culture.git / blobdiff