Update.
[picoclvr.git] / problems.py
1 #!/usr/bin/env python
2
3 import math
4
5 import torch, torchvision
6
7 from torch import nn
8 from torch.nn import functional as F
9
10 ######################################################################
11
12
13 class Problem:
14     def generate_sequences(self, nb):
15         pass
16
17     def seq2str(self, seq):
18         return "[NOT IMPLEMENTED]"
19
20     def compute_nb_correct(self, input, ar_mask, result):
21         nb_total = ar_mask.sum().item()
22         nb_correct = ((result == input).long() * ar_mask).sum().item()
23         return nb_total, nb_correct
24
25
26 ####################
27
28
29 class ProblemDegradation(Problem):
30     def __init__(self, nb_state_tokens=5, nb_time_steps=12, value_max=25, hard=False):
31         assert value_max // nb_state_tokens >= 2
32         self.nb_state_tokens = nb_state_tokens
33         self.nb_time_steps = nb_time_steps
34         self.value_max = value_max
35         self.hard = hard
36
37     def generate_sequences(self, nb):
38         x = (
39             torch.rand(nb, self.nb_state_tokens).sort(dim=-1).indices == 0
40         ).long() * self.value_max
41         seq = [x]
42
43         for t in range(self.nb_time_steps - 1):
44             v = (torch.rand(x.size()).sort(dim=-1).indices + 1) * (x >= 2).long()
45             u = (v.max(dim=-1, keepdim=True).values == v).long()
46             n = (
47                 (u * x)
48                 .minimum(2 + torch.randint(self.value_max // 4 - 2, x.size()))
49                 .sum(dim=-1, keepdim=True)
50             )
51             m = 1 + ((n - 1) * torch.rand(n.size())).long()
52             x = (
53                 x
54                 + m * u.roll(shifts=-1, dims=-1)
55                 - n * u
56                 + (n - m) * u.roll(shifts=1, dims=-1)
57             )
58             seq.append(x)
59
60         if self.hard:
61             seq.reverse()
62
63         seq = torch.cat(seq, dim=1)
64         return seq, seq.new_full(seq.size(), 1, dtype=torch.int64)
65
66     def compute_nb_correct(self, input, ar_mask, result):
67         nb_total = result.size(0)
68         nb_correct = 0
69         e = result.new_zeros(self.nb_state_tokens)
70
71         for seq in result:
72             states = list(seq.split(self.nb_state_tokens))
73             if self.hard:
74                 states.reverse()
75
76             d = states[0]
77             j = d.sort(descending=True).indices[0]
78             e.zero_()
79             e[j] = self.value_max
80             if (d - e).abs().sum() == 0:
81                 nb_errors = 0
82                 for k in range(len(states) - 1):
83                     d = states[k + 1] - states[k]
84                     j = d.sort(descending=False).indices[0]
85                     if (
86                         d[j] == 0
87                         or d[j] > self.value_max // 4
88                         or d[(j + 1) % e.size(0)] <= 0
89                         or d[(j + 1) % e.size(0)] >= -d[j]
90                     ):
91                         nb_errors += 1
92                     else:
93                         e.zero_()
94                         e[j] = d[j]
95                         e[(j + 1) % e.size(0)] = d[(j + 1) % e.size(0)]
96                         e[(j - 1) % e.size(0)] = -d[(j + 1) % e.size(0)] - d[j]
97                         if (d - e).abs().sum() > 0:
98                             nb_errors += 1
99                 if nb_errors == 0:
100                     nb_correct += 1
101
102         return nb_total, nb_correct
103
104     def seq2str(self, seq):
105         return " | ".join(
106             [" ".join([f"{x:02d}" for x in s]) for s in seq.split(self.nb_state_tokens)]
107         )
108
109
110 ####################
111
112
113 class ProblemMemory(Problem):
114     def __init__(self, len_total=25):
115         self.len_total = len_total
116         self.max_len_pattern = 5
117         self.nb_noise_tokens = 10
118         self.start_pattern_token = 0
119         self.end_pattern_token = 1
120         self.start_result_token = 2
121         self.end_result_token = 3
122         self.token_string = "[]<>" + "".join(
123             [chr(ord("a") + k) for k in range(self.nb_noise_tokens)]
124         )
125
126     def generate_sequences(self, nb):
127         sequences = (
128             torch.randint(self.nb_noise_tokens, (nb, self.len_total))
129             + self.end_result_token
130             + 1
131         )
132         len_patterns = torch.randint(self.max_len_pattern, (nb,)) + 1
133         pattern_positions = torch.randint(
134             self.len_total - (5 + 2 * self.max_len_pattern), (nb,)
135         )
136         k = self.len_total - (3 + self.max_len_pattern)
137         for i in range(nb):
138             l = len_patterns[i]
139             j = pattern_positions[i]
140             sequences[i, j] = self.start_pattern_token
141             sequences[i, j + l + 2] = self.end_pattern_token
142             sequences[i, k] = self.start_result_token
143             sequences[i, k + l + 2] = self.end_result_token
144             sequences[i, k + 1 : k + 2 + l] = sequences[i, j + 1 : j + 2 + l]
145
146         j = torch.arange(self.len_total)[None, :]
147         ar_mask = (j > k).long() * (j <= k + 1 + len_patterns[:, None]).long()
148
149         return sequences, ar_mask
150
151     def seq2str(self, seq):
152         return "".join(self.token_string[x.item()] for x in seq)
153
154
155 class ProblemTwoTargets(Problem):
156     def __init__(self, len_total=10, len_targets=3):
157         assert len_targets >= 3
158         assert len_total >= 3 * len_targets - 1
159         self.len_total = len_total
160         self.len_targets = len_targets
161
162     def generate_sequences(self, nb):
163         k = torch.arange(self.len_total)[None, :]
164         s = torch.randint(10, (nb, self.len_total))
165         l = torch.rand(nb, self.len_total)
166         l = l * (k <= self.len_total - self.len_targets).long()
167         k1 = l.argmax(dim=1, keepdim=True)
168         m = (k != k1).long() * (k != k1 + self.len_targets - 1).long()
169         s = s * m + 10 * (1 - m)
170         l = l * (
171             1
172             - (k + self.len_targets - 1 >= k1).long()
173             * (k < k1 + self.len_targets).long()
174         )
175         k2 = l.argmax(dim=1, keepdim=True)
176         m = (k != k2).long() * (k != k2 + self.len_targets - 1).long()
177         s = s * m + 11 * (1 - m)
178         a1 = s.gather(dim=1, index=k1 + 1 + torch.arange(self.len_targets - 2)[None, :])
179         a2 = s.gather(dim=1, index=k2 + 1 + torch.arange(self.len_targets - 2)[None, :])
180         sequences = torch.cat(
181             (
182                 s,
183                 torch.full((nb, 1), 12),
184                 a1,
185                 torch.full((nb, 1), 12),
186                 a2,
187                 torch.full((nb, 1), 12),
188             ),
189             1,
190         )
191         ar_mask = (sequences == 12).long()
192         ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
193         return sequences, ar_mask
194
195     def seq2str(self, seq):
196         return "".join("0123456789-+|"[x.item()] for x in seq)
197
198
199 ####################
200
201
202 class ProblemByHeart(Problem):
203     def __init__(self, nb_sentences=100, len_prompt=8, len_result=8, separation=1):
204         self.seq = torch.randint(
205             10, (nb_sentences, len_prompt + separation + len_result)
206         )
207         self.seq[:, len_prompt : len_prompt + separation] = 10
208
209     def generate_sequences(self, nb):
210         sequences = self.seq[torch.randint(self.seq.size(0), (nb,))]
211         ar_mask = (sequences == 10).long()
212         ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
213         return sequences, ar_mask
214
215     def seq2str(self, seq):
216         return "".join("0123456789|"[x.item()] for x in seq)
217
218
219 ####################
220
221
222 class ProblemLearnOperator(Problem):
223     def __init__(self, nb_operators=100, len_source=6, len_result=9):
224         self.len_source = len_source
225         self.len_result = len_result
226         self.len_nb_operator = int(math.log(nb_operators) / math.log(10)) + 1
227         self.operators = F.one_hot(
228             torch.rand(nb_operators, len_result, len_source).argmax(-1),
229             num_classes=len_source,
230         )
231
232     def generate_sequences(self, nb):
233         nb_operators = torch.randint(self.operators.size(0), (nb,))
234         operators = self.operators[nb_operators]
235         nb_operators = (
236             nb_operators[:, None]
237             // 10 ** torch.arange(self.len_nb_operator - 1, -1, -1)
238         ) % 10
239         marker1 = torch.full((nb, 1), 10)
240         source = torch.rand(nb, 10).sort(dim=1).indices[:, : self.len_source]
241         marker2 = torch.full((nb, 1), 11)
242         result = operators.bmm(source[:, :, None]).squeeze(-1)
243         sequences = torch.cat((nb_operators, marker1, source, marker2, result), 1)
244         ar_mask = (sequences == 11).long()
245         ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
246         return sequences, ar_mask
247
248     def seq2str(self, seq):
249         return "".join("0123456789|>"[x.item()] for x in seq)
250
251
252 ####################
253
254
255 class ProblemGuessOperator(Problem):
256     def __init__(self, len_source=5, len_result=8):
257         self.len_source = len_source
258         self.len_result = len_result
259
260     def generate_sequences(self, nb):
261         operators = F.one_hot(
262             torch.rand(nb, self.len_result, self.len_source).argmax(-1),
263             num_classes=self.len_source,
264         )
265         source1 = torch.rand(nb, 10).sort(dim=1).indices[:, : self.len_source]
266         marker1 = torch.full((nb, 1), 10)
267         result1 = operators.bmm(source1[:, :, None]).squeeze(-1)
268         marker2 = torch.full((nb, 1), 11)
269         source2 = torch.randint(10, (nb, self.len_source))
270         marker3 = torch.full((nb, 1), 12)
271         result2 = operators.bmm(source2[:, :, None]).squeeze(-1)
272
273         sequences = torch.cat(
274             (source1, marker1, result1, marker2, source2, marker3, result2), 1
275         )
276         ar_mask = (sequences == 12).long()
277         ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
278         return sequences, ar_mask
279
280     def seq2str(self, seq):
281         return "".join("0123456789>|~"[x.item()] for x in seq)
282
283
284 ####################
285
286
287 class ProblemAddition(Problem):
288     def __init__(self, nb_digits=10, zero_padded=False, inverted_result=False):
289         self.nb_digits = nb_digits
290         self.zero_padded = zero_padded
291         self.inverted_result = inverted_result
292         self.char2id = dict([(c, n) for n, c in enumerate("0123456789+=$")])
293         self.id2char = dict([(n, c) for c, n in self.char2id.items()])
294
295     def tensorize(self, strings):
296         len_max = max([len(x) for x in strings])
297         return torch.cat(
298             [
299                 torch.tensor(
300                     [
301                         [self.char2id[c] for c in s + "$" * (len_max - len(s))]
302                         for s in strings
303                     ]
304                 )
305             ],
306             0,
307         )
308
309     def generate_sequences(self, nb):
310         sequences = []
311         for k in range(nb):
312             a, b = torch.randint(10**self.nb_digits, (2,))
313             c = a + b
314             a, b, c = str(a.item()), str(b.item()), str(c.item())
315             if self.zero_padded:
316                 a = "0" * (self.nb_digits - len(a)) + a
317                 b = "0" * (self.nb_digits - len(b)) + b
318                 c = "0" * (self.nb_digits + 1 - len(c)) + c
319             if self.inverted_result:
320                 c = c[::-1]
321             sequences.append(f"{a}+{b}={c}$")
322
323         sequences = self.tensorize(sequences)
324         ar_mask = (sequences == self.char2id["="]).long()
325         ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
326         return sequences, ar_mask
327
328     def seq2str(self, seq):
329         return "".join(self.id2char[x.item()] for x in seq)
330
331
332 ####################
333
334
335 class ProblemMixing(Problem):
336     def __init__(
337         self, height=4, width=4, nb_time_steps=9, hard=False, random_start=True
338     ):
339         self.height = height
340         self.width = width
341         self.nb_time_steps = nb_time_steps
342         self.hard = hard
343         self.random_start = random_start
344
345     def start_random(self, nb):
346         y = torch.arange(self.height * self.width).reshape(1, -1).expand(nb, -1)
347
348         if self.random_start:
349             i = (
350                 torch.arange(self.height)
351                 .reshape(1, -1, 1)
352                 .expand(nb, self.height, self.width)
353             )
354             j = (
355                 torch.arange(self.width)
356                 .reshape(1, 1, -1)
357                 .expand(nb, self.height, self.width)
358             )
359
360             ri = torch.randint(self.height, (nb,)).reshape(nb, 1, 1)
361             rj = torch.randint(self.width, (nb,)).reshape(nb, 1, 1)
362
363             m = 1 - torch.logical_or(i == ri, j == rj).long().flatten(1)
364
365             y = y * m + self.height * self.width * (1 - m)
366
367         y = y.reshape(nb, self.height, self.width)
368
369         return y
370
371     def start_error(self, x):
372         if self.random_start:
373             i = (
374                 torch.arange(self.height, device=x.device)
375                 .reshape(1, -1, 1)
376                 .expand_as(x)
377             )
378             j = torch.arange(self.width, device=x.device).reshape(1, 1, -1).expand_as(x)
379
380             ri = (
381                 (x == self.height * self.width)
382                 .long()
383                 .sum(dim=-1)
384                 .argmax(-1)
385                 .view(-1, 1, 1)
386             )
387             rj = (
388                 (x == self.height * self.width)
389                 .long()
390                 .sum(dim=-2)
391                 .argmax(-1)
392                 .view(-1, 1, 1)
393             )
394
395             m = 1 - torch.logical_or(i == ri, j == rj).long().flatten(1)
396         else:
397             m = 1
398
399         x = x.flatten(1)
400         u = torch.arange(self.height * self.width, device=x.device).reshape(1, -1)
401
402         d = (x - (m * u + (1 - m) * self.height * self.width)).abs().sum(-1)
403
404         return d
405
406     def moves(self, x):
407         y = (
408             x[:, None, :, :]
409             .expand(-1, self.height * 2 + self.width * 2, -1, -1)
410             .clone()
411         )
412         k = 0
413
414         for i in range(self.height):
415             y[:, k, i, :] = y[:, k, i, :].roll(dims=-1, shifts=-1)
416             k += 1
417             y[:, k, i, :] = y[:, k, i, :].roll(dims=-1, shifts=1)
418             k += 1
419
420         for j in range(self.width):
421             y[:, k, :, j] = y[:, k, :, j].roll(dims=-1, shifts=-1)
422             k += 1
423             y[:, k, :, j] = y[:, k, :, j].roll(dims=-1, shifts=1)
424             k += 1
425
426         return y
427
428     def generate_sequences(self, nb):
429         x = self.start_random(nb)
430
431         seq = [x.flatten(1)]
432
433         for t in range(self.nb_time_steps - 1):
434             y = self.moves(x)
435             x = y[torch.arange(nb), torch.randint(y.size(1), (nb,))]
436             seq.append(x.flatten(1))
437
438         if self.hard:
439             seq.reverse()
440
441         seq = torch.cat(seq, dim=1)
442         return seq, seq.new_full(seq.size(), 1, dtype=torch.int64)
443
444     def compute_nb_correct(self, input, ar_mask, result):
445         a = [
446             x.reshape(result.size(0), self.height, self.width)
447             for x in result.split(self.height * self.width, dim=1)
448         ]
449         if self.hard:
450             a.reverse()
451
452         x = a[0]
453
454         d = self.start_error(x)
455
456         for t in range(self.nb_time_steps - 1):
457             x0, x = a[t], a[t + 1]
458             y = self.moves(x0)
459             d = d + (x[:, None] - y).abs().sum((-1, -2)).min(dim=-1).values
460
461         nb_total, nb_correct = result.size(0), (d == 0).long().sum().item()
462
463         return nb_total, nb_correct
464
465     def seq2str(self, seq):
466         return " | ".join(
467             [
468                 " ".join(
469                     [
470                         "-".join(
471                             [
472                                 f"{x:02d}" if x < self.height * self.width else "**"
473                                 for x in s
474                             ]
475                         )
476                         for s in r.split(self.width)
477                     ]
478                 )
479                 for r in seq.split(self.height * self.width)
480             ]
481         )
482
483
484 ####################
485
486 if __name__ == "__main__":
487     p = ProblemMixing(height=3, width=3, random_start=False)
488
489     s, m = p.generate_sequences(10000)
490     for x in s[:5]:
491         print(p.seq2str(x))
492     print(p.compute_nb_correct(None, None, s))