X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=quizz_machine.py;h=bf36d0bbe7a4f47cc867fe426c2d6ff7c5931d18;hb=784dfe9667c9cc326533a1081eb11571fe33e113;hp=f799bf1c52a133dd6c2970ae80ac1f835143d45c;hpb=9f787901b2c7591a323f843ab973fe6abcf6b8ce;p=culture.git

diff --git a/quizz_machine.py b/quizz_machine.py
index f799bf1..bf36d0b 100755
--- a/quizz_machine.py
+++ b/quizz_machine.py
@@ -17,6 +17,88 @@ from mygpt import BracketedSequence
 ######################################################################
 
 
+class Gang(nn.Module):
+    def __init__(self, models, nb_models_for_generation, mode="groupthink"):
+        super().__init__()
+        self.models = models
+        self.nb_models_for_generation = nb_models_for_generation
+        self.mode = mode
+
+    def forward(self, bs):
+        # If first = 0, we are re-starting an auto-regressive process,
+        # that's the right moment to randomize who gonna do it
+        if bs.first == 0:
+            self.models_to_use = [
+                self.models[k]
+                for k in torch.randperm(len(self.models))[
+                    : self.nb_models_for_generation
+                ]
+            ]
+
+        all_the_logits = torch.cat(
+            [model(bs).x[None] for model in self.models_to_use], dim=0
+        )
+
+        if self.mode == "groupthink":
+            y = all_the_logits.mean(dim=0)
+        elif self.mode == "groupwork":
+            m = torch.rand(all_the_logits.size(), device=all_the_logits.device)
+            m = (m.sort(dim=0).indices == 0).long()
+            y = (y * m).sum(dim=0)
+        else:
+            raise ValueError(f"Invalid mode {self.mode}")
+
+        return BracketedSequence(y, bs.first, bs.nb)
+
+
+######################################################################
+
+# ar_mask is a tensor with 0s and 1s, of same shape as input, with
+# 1s where tokens should be generated. The others are kept
+# unchanged.
+
+
+def one_batch_masked_inplace_autoregression(
+    model,
+    input,
+    ar_mask,
+    seq_logproba,
+    temperature=1.0,
+    deterministic_synthesis=False,
+    forbidden_tokens=None,
+    forced_biases=None,
+):
+    to_generate = (ar_mask.sum(0) > 0).nonzero()
+
+    if to_generate.min() > 0:
+        model(
+            BracketedSequence(input, 0, to_generate.min())
+        )  # Needed to initialize the model's cache
+    for s in range(to_generate.min(), to_generate.max() + 1):
+        output = model(BracketedSequence(input, s, 1)).x
+
+        logits = output[:, s]
+
+        logits = (logits / temperature).log_softmax(dim=-1)
+
+        if forbidden_tokens is not None:
+            logits = logits.masked_fill(forbidden_tokens, float("-inf"))
+
+        if forced_biases is not None:
+            logits = logits + forced_biases[None, :]
+
+        if deterministic_synthesis:
+            t_next = logits.argmax(-1)
+        else:
+            dist = torch.distributions.categorical.Categorical(logits=logits)
+            t_next = dist.sample()
+
+        all_n = torch.arange(t_next.size(0))
+        seq_logproba += logits[all_n, t_next].sum(dim=-1)
+
+        input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s]
+
+
 def masked_inplace_autoregression(
     model,
     batch_size,
@@ -51,7 +133,8 @@ def masked_inplace_autoregression(
         model.eval()
 
         for input, ar_mask, seq_logproba in batches:
-            model.masked_inplace_autoregression(
+            one_batch_masked_inplace_autoregression(
+                model=model,
                 input=input,
                 ar_mask=ar_mask,
                 seq_logproba=seq_logproba,
@@ -88,8 +171,12 @@ class QuizzMachine:
         self.batch_size = batch_size
         self.device = device
 
-        self.train_w_quizzes = self.problem.generate_seq(nb_train_samples).to(device)
-        self.test_w_quizzes = self.problem.generate_seq(nb_test_samples).to(device)
+        self.train_w_quizzes = self.problem.generate_token_sequences(
+            nb_train_samples
+        ).to(device)
+        self.test_w_quizzes = self.problem.generate_token_sequences(nb_test_samples).to(
+            device
+        )
 
         self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1
 
@@ -113,7 +200,7 @@ class QuizzMachine:
         if len(c_quizzes) > 0:
             c_quizzes = torch.cat(c_quizzes, dim=0)
             if c_quizzes.size(0) > w_quizzes.size(0) // 2:
-                i = torch.randperm(w_quizzes.size(0))[: w_quizzes.size(0) // 2]
+                i = torch.randperm(c_quizzes.size(0))[: w_quizzes.size(0) // 2]
                 c_quizzes = c_quizzes[i]
 
             i = torch.randperm(w_quizzes.size(0))[
@@ -215,7 +302,7 @@ class QuizzMachine:
         input = self.train_w_quizzes if for_train else self.test_w_quizzes
         nb = min(nb, input.size(0))
         input[:-nb] = input[nb:].clone()
-        input[-nb:] = self.problem.generate_seq(nb).to(self.device)
+        input[-nb:] = self.problem.generate_token_sequences(nb).to(self.device)
 
     def store_c_quizzes(self, new_c_quizzes, for_train=True):
         if for_train:
@@ -223,64 +310,7 @@ class QuizzMachine:
         else:
             self.test_c_quizzes.append(new_c_quizzes)
 
-    def create_c_quizzes(
-        self,
-        nb,
-        model_for_generation,
-        models_for_validation,
-        min_ave_seq_logproba,
-        n_epoch,
-        result_dir,
-        logger,
-    ):
-        ###############################################################
-        # Generate quizzes with model
-
-        c_quizzes = torch.empty(
-            nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64
-        )
-
-        ar_mask = torch.full(c_quizzes.size(), 1, device=self.device)
-        seq_logproba = torch.empty(ar_mask.size(0), device=self.device)
-
-        temperature = 1
-        d_temperature = 1 / 3
-
-        while True:
-            seq_logproba[...] = 0
-
-            masked_inplace_autoregression(
-                model=model_for_generation,
-                batch_size=self.batch_size,
-                input=c_quizzes,
-                ar_mask=ar_mask,
-                seq_logproba=seq_logproba,
-                temperature=temperature,
-                deterministic_synthesis=False,
-                # progress_bar_desc="sampling c_quizzes",
-                device=self.device,
-            )
-
-            ave_seq_logproba = seq_logproba.mean()
-
-            if min_ave_seq_logproba is None:
-                break
-
-            # Oh man that's ugly
-            if ave_seq_logproba < min_ave_seq_logproba:
-                if d_temperature > 0:
-                    d_temperature *= -1 / 3
-                temperature += d_temperature
-            elif ave_seq_logproba > min_ave_seq_logproba * 0.99:
-                if d_temperature < 0:
-                    d_temperature *= -1 / 3
-                temperature += d_temperature
-            else:
-                break
-
-            logger(f"changing temperature to {temperature}")
-
-        ###############################################################
+    def comput_correctness(self, c_quizzes, models_for_validation):
         # Create the reverse quizzes
 
         token_forward, token_backward = self.problem.direction_tokens()
@@ -297,11 +327,9 @@ class QuizzMachine:
         ar_mask = self.make_ar_mask(c_quizzes)
         seq_logproba = torch.empty(ar_mask.size(0), device=self.device)
 
-        ###############################################################
-        # Check how many of the other models can solve them in both
-        # directions
+        # Check how many of models can solve the quizzes in both directions
 
-        nb_correct = []
+        nb_correct = 0
 
         for model in models_for_validation:
             result = c_quizzes.clone()
@@ -338,8 +366,103 @@ class QuizzMachine:
                 (reverse_c_quizzes == reverse_result).long().min(dim=-1).values
             )
 
-            nb_correct.append((correct * reverse_correct)[None, :])
+            nb_correct += correct * reverse_correct
+
+        return nb_correct
+
+    ###############################################################
+
+    def generate_quizzes(self, nb, model_for_generation, min_ave_seq_logproba):
+        c_quizzes = torch.empty(
+            nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64
+        )
+
+        ar_mask = torch.full(c_quizzes.size(), 1, device=self.device)
+        seq_logproba = torch.empty(ar_mask.size(0), device=self.device)
+
+        # bracketing of the temperature to get the target logproba
+
+        temperature = 1
+        d_temperature = 1 / 3
+
+        while True:
+            seq_logproba[...] = 0
+
+            masked_inplace_autoregression(
+                model=model_for_generation,
+                batch_size=self.batch_size,
+                input=c_quizzes,
+                ar_mask=ar_mask,
+                seq_logproba=seq_logproba,
+                temperature=temperature,
+                deterministic_synthesis=False,
+                # progress_bar_desc="sampling c_quizzes",
+                device=self.device,
+            )
+
+            ave_seq_logproba = seq_logproba.mean()
 
-        nb_correct = torch.cat(nb_correct, dim=0).sum(dim=0)
+            # If we do not have target logprobs, get out now
+            if min_ave_seq_logproba is None:
+                break
 
-        return c_quizzes, nb_correct, seq_logproba.mean()
+            # Oh man that's ugly
+            if ave_seq_logproba < min_ave_seq_logproba:
+                if d_temperature > 0:
+                    d_temperature *= -1 / 3
+                temperature += d_temperature
+            elif ave_seq_logproba > min_ave_seq_logproba * 0.99:
+                if d_temperature < 0:
+                    d_temperature *= -1 / 3
+                temperature += d_temperature
+            else:
+                break
+
+            logger(f"changing temperature to {temperature}")
+
+        return c_quizzes, seq_logproba.mean()
+
+    ######################################################################
+
+    def create_c_quizzes(
+        self,
+        nb,
+        model_for_generation,
+        models_for_validation,
+        min_ave_seq_logproba,
+        n_epoch,
+        result_dir,
+        logger,
+    ):
+        c_quizzes, ave_seq_logproba = self.generate_quizzes(
+            nb, model_for_generation, min_ave_seq_logproba
+        )
+
+        nb_correct = self.comput_correctness(c_quizzes, models_for_validation)
+
+        return c_quizzes, nb_correct, ave_seq_logproba
+
+    ######################################################################
+
+    def gang_create_c_quizzes(
+        self,
+        nb,
+        nb_models_for_generation,
+        models,
+        mode,
+        min_ave_seq_logproba,
+        n_epoch,
+        result_dir,
+        logger,
+    ):
+        model_for_generation = Gang(models, nb_models_for_generation, mode)
+        models_for_validation = models
+        return self.create_c_quizzes(
+            nb,
+            model_for_generation,
+            models_for_validation,
+            min_ave_seq_logproba,
+            n_epoch,
+            result_dir,
+            logger,
+        )