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

diff --git a/tasks.py b/tasks.py
index 332d6c5..706e1d9 100755
--- a/tasks.py
+++ b/tasks.py
@@ -87,35 +87,34 @@ class ProblemLevel0(Problem):
 
 
 class ProblemLevel1(Problem):
-    def __init__(self, nb_operators=100, len_prompt=5, len_result=8):
-        self.len_prompt = len_prompt
+    def __init__(self, nb_operators=100, len_source=5, len_result=8):
+        self.len_source = len_source
         self.len_result = len_result
         self.len_nb_operator = int(math.log(nb_operators) / math.log(10)) + 1
         self.operators = F.one_hot(
-            torch.rand(nb_operators, len_result, len_prompt).argmax(-1),
-            num_classes=len_prompt,
+            torch.rand(nb_operators, len_result, len_source).argmax(-1),
+            num_classes=len_source,
         )
 
+
+
     def generate_sequences(self, nb):
-        a = self.len_nb_operator
-        b = a + 1 + self.len_prompt
-        sequences = torch.empty(nb, b + 1 + self.len_result, dtype=torch.int64)
         nb_operators = torch.randint(self.operators.size(0), (nb,))
-        sequences[:, :a] = (nb_operators[:, None] / 10 ** torch.arange(a)) % 10
-        sequences[:, a] = 10
-        sequences[:, a + 1 : b] = torch.randint(10, (nb, b - a - 1))
-        sequences[:, b] = 11
-
-        o = self.operators[nb_operators]
-        p = sequences[:, a + 1 : b]
-        print(f"{o.size()=} {p.size()=} {sequences[:,b+1:].size()=}")
-        sequences[:, b + 1 :] = o.bmm(p[:, :, None]).squeeze(-1)
+        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)
+        result = operators.bmm(source[:, :, None]).squeeze(-1)
+        print(f"{nb_operators.dtype=} {marker1.dtype=}")
+        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)
         return sequences, ar_mask
 
     def seq2str(self, seq):
-        return "".join(self.id2char[x.item()] for x in seq)
+        return "".join("0123456789|>"[x.item()] for x in seq)
 
 
 ####################