Update.

diff --git a/quizz_machine.py b/quizz_machine.py
index 153317c..90f288e 100755 (executable)
--- a/quizz_machine.py
+++ b/quizz_machine.py
@@ -105,10 +105,64 @@ def masked_inplace_autoregression(
 
 
 class QuizzMachine:
 
 
 class QuizzMachine:

-    def make_ar_mask(self, input, first, nb):
-        i = torch.arange(input.size(1), device=input.device)
-        b = torch.logical_and(i >= first, i < first + nb)
-        return b.long()[None, :].expand_as(input)
+    def indices_forward_and_backward(self, quizzes):
+        i_forward = quizzes[:, 0] == self.token_forward
+        j_forward = quizzes[:, 1 + self.prompt_len] == self.token_forward
+        i_backward = quizzes[:, 0] == self.token_backward
+        j_backward = quizzes[:, 1 + self.answer_len] == self.token_backward
+        assert torch.logical_or(
+            torch.logical_and(i_forward, j_forward),
+            torch.logical_and(i_backward, j_backward),
+        ).all()
+        return i_forward, i_backward
+
+    def reverse_time(self, quizzes):
+        i_forward, i_backward = self.indices_forward_and_backward(quizzes)
+
+        forward_to_backward = torch.cat(
+            [
+                quizzes[:, 0:1],
+                quizzes[:, 2 + self.prompt_len :],
+                quizzes[:, 1 + self.prompt_len : 2 + self.prompt_len],
+                quizzes[:, 1 : 1 + self.prompt_len],
+            ],
+            dim=1,
+        )
+        forward_to_backward[:, 0] = self.token_backward
+        forward_to_backward[:, 1 + self.answer_len] = self.token_backward
+
+        backward_to_forward = torch.cat(
+            [
+                quizzes[:, 0:1],
+                quizzes[:, 2 + self.answer_len :],
+                quizzes[:, 1 + self.answer_len : 2 + self.answer_len],
+                quizzes[:, 1 : 1 + self.answer_len],
+            ],
+            dim=1,
+        )
+
+        backward_to_forward[:, 0] = self.token_forward
+        backward_to_forward[:, 1 + self.prompt_len] = self.token_forward
+
+        m = i_forward.long()[:, None]
+
+        return m * forward_to_backward + (1 - m) * backward_to_forward
+
+    def make_ar_mask(self, quizzes, first=False):
+        i_forward, i_backward = self.indices_forward_and_backward(quizzes)
+
+        t = torch.arange(quizzes.size(1), device=quizzes.device)
+
+        if first:
+            m_forward = (t >= 1).long() * (t < 1 + self.prompt_len).long()
+            m_backward = (t >= 1).long() * (t < 1 + self.answer_len).long()
+        else:
+            m_forward = (t >= 2 + self.prompt_len).long()
+            m_backward = (t >= 2 + self.answer_len).long()
+
+        m = i_forward.long()[:, None]
+
+        return m * m_forward + (1 - m) * m_backward

 
     def generate_token_sequences(self, nb):
         prompts, answers = self.problem.generate_prompts_and_answers(nb)
 
     def generate_token_sequences(self, nb):
         prompts, answers = self.problem.generate_prompts_and_answers(nb)


@@ -181,19 +235,20 @@ class QuizzMachine:
                 result_dir, "culture_w_quizzes", self.train_w_quizzes[:72]
             )
 
                 result_dir, "culture_w_quizzes", self.train_w_quizzes[:72]
             )
 

+            # toto = self.reverse_time(self.train_w_quizzes[:72])
+            # self.save_quizzes(result_dir, "toto", toto)
+            # exit(0)
+

     def save_quizzes(self, result_dir, filename_prefix, quizzes, prediction=False):
     def save_quizzes(self, result_dir, filename_prefix, quizzes, prediction=False):

-        l = (quizzes.size(1) - 1) // 2
-        forward = (quizzes[:, 0] == self.token_forward).long()
-        backward = (quizzes[:, 0] == self.token_backward).long()
-        assert forward.equal(1 - backward)
-        first = quizzes[:, 1 : 1 + l]
-        second = quizzes[:, 1 + l : 1 + 2 * l]
-        prompts = forward[:, None] * first + backward[:, None] * second
-        answers = forward[:, None] * second + backward[:, None] * first
+        forward = quizzes[quizzes[:, 0] == self.token_forward]
+        ib = quizzes[:, 0] == self.token_backward
+        backward = quizzes[ib]
+        assert forward.size(0) + backward.size(0) == quizzes.size(0)
+        quizzes[ib] = self.reverse_time(quizzes[ib])

 
         if prediction:
 
         if prediction:

-            predicted_prompts = backward
-            predicted_answers = forward
+            predicted_prompts = ib
+            predicted_answers = torch.logical_not(ib)

         else:
             predicted_prompts = None
             predicted_answers = None
         else:
             predicted_prompts = None
             predicted_answers = None


@@ -201,8 +256,8 @@ class QuizzMachine:
         self.problem.save_quizzes(
             result_dir,
             filename_prefix,
         self.problem.save_quizzes(
             result_dir,
             filename_prefix,

-            prompts,
-            answers,
+            quizzes[:, 1 : 1 + self.prompt_len],
+            quizzes[:, 2 + self.prompt_len :],

             predicted_prompts,
             predicted_answers,
         )
             predicted_prompts,
             predicted_answers,
         )


@@ -254,7 +309,7 @@ class QuizzMachine:
     ):
         def compute_accuracy(input):
             input = input[:nmax]
     ):
         def compute_accuracy(input):
             input = input[:nmax]

-            ar_mask = self.make_ar_mask(input, 2 + self.prompt_len, self.answer_len)
+            ar_mask = self.make_ar_mask(input)

             result = input.clone() * (1 - ar_mask)
             seq_logproba = torch.empty(input.size(0), device=self.device)
 
             result = input.clone() * (1 - ar_mask)
             seq_logproba = torch.empty(input.size(0), device=self.device)
 


@@ -293,7 +348,7 @@ class QuizzMachine:
         ##############################
 
         input = self.test_w_quizzes[:96]
         ##############################
 
         input = self.test_w_quizzes[:96]

-        ar_mask = self.make_ar_mask(input, 2 + self.prompt_len, self.answer_len)
+        ar_mask = self.make_ar_mask(input)

         result = input.clone() * (1 - ar_mask)
         seq_logproba = torch.empty(input.size(0), device=self.device)
 
         result = input.clone() * (1 - ar_mask)
         seq_logproba = torch.empty(input.size(0), device=self.device)
 


@@ -330,32 +385,6 @@ class QuizzMachine:
         else:
             self.test_c_quizzes.append(new_c_quizzes)
 
         else:
             self.test_c_quizzes.append(new_c_quizzes)
 

-    def forward_to_backward(self, c_quizzes):
-        prompts = c_quizzes[:, 1 : 1 + self.prompt_len]
-        answers = c_quizzes[:, 2 + self.prompt_len :]
-        return torch.cat(
-            [
-                c_quizzes.new_full((c_quizzes, 1), self.token_backward),
-                answers,
-                c_quizzes.new_full((c_quizzes, 1), self.token_backward),
-                prompts,
-            ],
-            dim=1,
-        )
-
-    def backward_to_forward(self, c_quizzes):
-        answers = c_quizzes[:, 1 : 1 + self.answer_len :]
-        prompts = c_quizzes[:, 2 + self.answer_len :]
-        return torch.cat(
-            [
-                c_quizzes.new_full((c_quizzes.size(0), 1), self.token_forward),
-                prompts,
-                c_quizzes.new_full((c_quizzes.size(0), 1), self.token_forward),
-                answers,
-            ],
-            dim=1,
-        )
-

     def compute_correctness(
         self,
         c_quizzes,
     def compute_correctness(
         self,
         c_quizzes,


@@ -379,7 +408,7 @@ class QuizzMachine:
 
             seq_logproba[...] = 0.0
 
 
             seq_logproba[...] = 0.0
 

-            ar_mask = self.make_ar_mask(result, 2 + self.prompt_len, self.answer_len)
+            ar_mask = self.make_ar_mask(result)

 
             masked_inplace_autoregression(
                 model=model,
 
             masked_inplace_autoregression(
                 model=model,


@@ -398,9 +427,7 @@ class QuizzMachine:
             if bidirectional_validation:
                 backward_result = backward_c_quizzes.clone()
 
             if bidirectional_validation:
                 backward_result = backward_c_quizzes.clone()
 

-                ar_mask = self.make_ar_mask(
-                    backward_result, 2 + self.answer_len, self.prompt_len
-                )
+                ar_mask = self.make_ar_mask(backward_result)

 
                 masked_inplace_autoregression(
                     model=model,
 
                 masked_inplace_autoregression(
                     model=model,


@@ -433,23 +460,18 @@ class QuizzMachine:
             nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64
         )
 
             nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64
         )
 

-        ar_mask_first = torch.zeros(c_quizzes.size(), device=self.device)
-        ar_mask_first[:, : ar_mask_first.size(1) // 2 + 1] = 1
-        ar_mask_second = 1 - ar_mask_first
-        ar_mask_first[:, 0] = 0
-        ar_mask_second[:, 0] = 0
-
-        seq_logproba = torch.zeros(ar_mask_first.size(0), device=self.device)
+        seq_logproba = torch.zeros(nb, device=self.device)

 
         # First, we generate the answer at high temperature
 
         c_quizzes[:, 0] = self.token_backward
 
         # First, we generate the answer at high temperature
 
         c_quizzes[:, 0] = self.token_backward

+        c_quizzes[:, 1 + self.answer_len] = self.token_backward

 
         masked_inplace_autoregression(
             model=model_for_generation,
             batch_size=self.batch_size,
             input=c_quizzes,
 
         masked_inplace_autoregression(
             model=model_for_generation,
             batch_size=self.batch_size,
             input=c_quizzes,

-            ar_mask=ar_mask_first,
+            ar_mask=self.make_ar_mask(c_quizzes, first=True),

             seq_logproba=seq_logproba,
             temperature=temperature,
             deterministic_synthesis=False,
             seq_logproba=seq_logproba,
             temperature=temperature,
             deterministic_synthesis=False,


@@ -462,7 +484,7 @@ class QuizzMachine:
             model=model_for_generation,
             batch_size=self.batch_size,
             input=c_quizzes,
             model=model_for_generation,
             batch_size=self.batch_size,
             input=c_quizzes,

-            ar_mask=ar_mask_second,
+            ar_mask=self.make_ar_mask(c_quizzes),

             seq_logproba=seq_logproba,
             temperature=1 / temperature,
             deterministic_synthesis=False,
             seq_logproba=seq_logproba,
             temperature=1 / temperature,
             deterministic_synthesis=False,


@@ -472,13 +494,13 @@ class QuizzMachine:
         # Then we return the quizz, and re-generate the response, now
         # at low temperature
 
         # Then we return the quizz, and re-generate the response, now
         # at low temperature
 

-        c_quizzes = self.backward_to_forward(c_quizzes)
+        c_quizzes = self.reverse_time(c_quizzes)

 
         masked_inplace_autoregression(
             model=model_for_generation,
             batch_size=self.batch_size,
             input=c_quizzes,
 
         masked_inplace_autoregression(
             model=model_for_generation,
             batch_size=self.batch_size,
             input=c_quizzes,

-            ar_mask=ar_mask_second,
+            ar_mask=self.make_ar_mask(c_quizzes),

             seq_logproba=seq_logproba,
             temperature=1 / temperature,
             deterministic_synthesis=False,
             seq_logproba=seq_logproba,
             temperature=1 / temperature,
             deterministic_synthesis=False,

diff --git a/sky.py b/sky.py
index 040ec67..e509fb7 100755 (executable)
--- a/sky.py
+++ b/sky.py
@@ -5,7 +5,7 @@
 
 # Written by Francois Fleuret <francois@fleuret.org>
 
 
 # Written by Francois Fleuret <francois@fleuret.org>
 

-import math, sys, tqdm, os
+import math, sys, tqdm, os, warnings

 
 import torch, torchvision
 
 
 import torch, torchvision
 


@@ -200,38 +200,66 @@ class Sky(problem.Problem):
         if predicted_answers is None:
             predicted_answers = 255
 
         if predicted_answers is None:
             predicted_answers = 255
 

-        def add_frame(x, c, margin):
-            y = x.new_full(
-                (x.size(0), x.size(1), x.size(2) + 2 * margin, x.size(3) + 2 * margin),
-                0,
-            )
+        def add_frame(x, c, margin, bottom=False):
+            if bottom:
+                h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
+            else:
+                h, w, di, dj = (
+                    x.size(2) + 2 * margin,
+                    x.size(3) + 2 * margin,
+                    margin,
+                    margin,
+                )
+
+            y = x.new_full((x.size(0), x.size(1), h, w), 0)
+

             if type(c) is int:
                 y[...] = c
             else:
                 c = c.long()[:, None]
             if type(c) is int:
                 y[...] = c
             else:
                 c = c.long()[:, None]

-                c = c * torch.tensor([192, 192, 192], device=c.device) + (
+                c = c * torch.tensor([0, 0, 0], device=c.device) + (

                     1 - c
                 ) * torch.tensor([255, 255, 255], device=c.device)
                 y[...] = c[:, :, None, None]
                     1 - c
                 ) * torch.tensor([255, 255, 255], device=c.device)
                 y[...] = c[:, :, None, None]

-            y[:, :, margin:-margin, margin:-margin] = x
+
+            y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
+

             return y
 
         margin = 4
 
             return y
 
         margin = 4
 

-        img_prompts = add_frame(self.frame2img(prompts.to("cpu")), 0, 1)
-        img_answers = add_frame(self.frame2img(answers.to("cpu")), 0, 1)
+        img_prompts = add_frame(self.frame2img(prompts.to("cpu")), c=0, margin=1)
+        img_answers = add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1)

 
 

-        # img_prompts = add_frame(img_prompts, 255, margin)
-        # img_answers = add_frame(img_answers, 255, margin)
+        img_prompts = add_frame(img_prompts, c=255, margin=margin, bottom=True)
+        img_answers = add_frame(img_answers, c=255, margin=margin, bottom=True)

 
 

-        img_prompts = add_frame(img_prompts, predicted_prompts, margin)
-        img_answers = add_frame(img_answers, predicted_answers, margin)
+        img_prompts = add_frame(
+            img_prompts, c=predicted_prompts, margin=margin, bottom=True
+        )
+        img_answers = add_frame(
+            img_answers, c=predicted_answers, margin=margin, bottom=True
+        )
+
+        marker_size = 8

 
         separator = img_prompts.new_full(
 
         separator = img_prompts.new_full(

-            (img_prompts.size(0), img_prompts.size(1), img_prompts.size(2), margin), 255
+            (
+                img_prompts.size(0),
+                img_prompts.size(1),
+                img_prompts.size(2),
+                marker_size,
+            ),
+            255,

         )
 
         )
 

-        img = torch.cat([img_prompts, img_answers], dim=3)
+        for k in range(2, 2 * marker_size - 3):
+            i = k + 1 - marker_size
+            j = marker_size - 2 - abs(k - marker_size + 1)
+            separator[:, :, separator.size(2) // 2 + i, j] = 0
+            separator[:, :, separator.size(2) // 2 + i + 1, j] = 0
+
+        img = torch.cat([img_prompts, separator, img_answers], dim=3)

 
         image_name = os.path.join(result_dir, filename)
         torchvision.utils.save_image(
 
         image_name = os.path.join(result_dir, filename)
         torchvision.utils.save_image(


@@ -246,8 +274,19 @@ class Sky(problem.Problem):
     def generate_prompts_and_answers(self, nb):
         frame_sequences = self.generate_frame_sequences(nb)
         frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0)
     def generate_prompts_and_answers(self, nb):
         frame_sequences = self.generate_frame_sequences(nb)
         frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0)

+

         prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1)
         prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1)

+

         answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1)
         answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1)

+        # warnings.warn("dirty test with longer answer", RuntimeWarning)
+        # answers = torch.cat(
+        # [
+        # frame_sequences[:, frame_sequences.size(1) // 2 :],
+        # frame_sequences[:, frame_sequences.size(1) // 2 :],
+        # ],
+        # dim=3,
+        # ).flatten(1)
+

         return prompts, answers
 
     def save_quizzes(
         return prompts, answers
 
     def save_quizzes(