X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=c5418b4bb616386d8f34e038ff96460dfafff585;hb=e781d77071fa26f393f50451f91c70f4a0850ca5;hp=706e1d913c20be0bf9449f94bf3346658a5a1bdc;hpb=8d2ebe29b48e3cf2f0a3937ab1e44d0e12a4924e;p=culture.git

diff --git a/tasks.py b/tasks.py
index 706e1d9..c5418b4 100755
--- a/tasks.py
+++ b/tasks.py
@@ -76,7 +76,7 @@ class Problem:
 
 class ProblemLevel0(Problem):
     def __init__(self, nb_sentences=100, len_prompt=5, len_result=5):
-        self.seq = torch.randint(10, (nb_seq, len_prompt + 1 + len_result))
+        self.seq = torch.randint(10, (nb_sentences, len_prompt + 1 + len_result))
         self.seq[:, len_prompt] = 10
 
     def generate_sequences(self, nb):
@@ -96,18 +96,20 @@ class ProblemLevel1(Problem):
             num_classes=len_source,
         )
 
-
-
     def generate_sequences(self, nb):
         nb_operators = torch.randint(self.operators.size(0), (nb,))
         operators = self.operators[nb_operators]
-        nb_operators = (nb_operators[:, None] // 10 ** torch.arange(self.len_nb_operator-1,-1,-1)) % 10
-        marker1 = torch.full((nb,1),10)
-        source = torch.randint(10, (nb, self.len_source))
-        marker2 = torch.full((nb,1),11)
+        nb_operators = (
+            nb_operators[:, None]
+            // 10 ** torch.arange(self.len_nb_operator - 1, -1, -1)
+        ) % 10
+        marker1 = torch.full((nb, 1), 10)
+        # source = torch.randint(10, (nb, self.len_source))
+        source = torch.rand(nb, 10).sort(dim=1).indices[:, : self.len_source]
+        marker2 = torch.full((nb, 1), 11)
         result = operators.bmm(source[:, :, None]).squeeze(-1)
         print(f"{nb_operators.dtype=} {marker1.dtype=}")
-        sequences = torch.cat((nb_operators, marker1, source,marker2,result),1)
+        sequences = torch.cat((nb_operators, marker1, source, marker2, result), 1)
         print(f"{sequences.size()=}")
         ar_mask = (sequences == 11).long()
         ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
@@ -117,6 +119,36 @@ class ProblemLevel1(Problem):
         return "".join("0123456789|>"[x.item()] for x in seq)
 
 
+class ProblemLevel2(Problem):
+    def __init__(self, len_source=5, len_result=8):
+        self.len_source = len_source
+        self.len_result = len_result
+
+    def generate_sequences(self, nb):
+        operators = F.one_hot(
+            torch.rand(nb, self.len_result, self.len_source).argmax(-1),
+            num_classes=self.len_source,
+        )
+        source1 = torch.rand(nb, 10).sort(dim=1).indices[:, : self.len_source]
+        # source1 = torch.randint(10, (nb, self.len_source))
+        marker1 = torch.full((nb, 1), 10)
+        result1 = operators.bmm(source1[:, :, None]).squeeze(-1)
+        marker2 = torch.full((nb, 1), 11)
+        source2 = torch.randint(10, (nb, self.len_source))
+        marker3 = torch.full((nb, 1), 12)
+        result2 = operators.bmm(source2[:, :, None]).squeeze(-1)
+
+        sequences = torch.cat(
+            (source1, marker1, result1, marker2, source2, marker3, result2), 1
+        )
+        ar_mask = (sequences == 12).long()
+        ar_mask = (ar_mask.cumsum(1) - ar_mask).clamp(max=1)
+        return sequences, ar_mask
+
+    def seq2str(self, seq):
+        return "".join("0123456789>|~"[x.item()] for x in seq)
+
+
 ####################