cdfba85e6a2abb8d6cd1b14bb3b0f76fe2afad30
[culture.git] / quiz_machine.py
1 #!/usr/bin/env python
2
3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
5
6 # Written by Francois Fleuret <francois@fleuret.org>
7
8 import math, os, tqdm, warnings
9
10 import torch, torchvision
11
12 from torch import nn
13 from torch.nn import functional as F
14
15 import mygpt
16 from mygpt import BracketedSequence
17
18 ######################################################################
19
20 # ar_mask is a tensor with 0s and 1s, of same shape as input, with
21 # 1s where tokens should be generated. The others are kept
22 # unchanged.
23
24
25 def one_batch_masked_inplace_autoregression(
26     model,
27     input,
28     ar_mask,
29     seq_logproba,
30     temperature,
31     deterministic_synthesis,
32 ):
33     to_generate = (ar_mask.sum(0) > 0).nonzero()
34
35     if to_generate.min() > 0:
36         model(
37             BracketedSequence(input, 0, to_generate.min())
38         )  # Needed to initialize the model's cache
39     for s in range(to_generate.min(), to_generate.max() + 1):
40         output = model(BracketedSequence(input, s, 1)).x
41
42         logits = output[:, s]
43
44         logits = (logits / temperature).log_softmax(dim=-1)
45
46         if deterministic_synthesis:
47             t_next = logits.argmax(-1)
48         else:
49             dist = torch.distributions.categorical.Categorical(logits=logits)
50             t_next = dist.sample()
51
52         all_n = torch.arange(t_next.size(0))
53
54         seq_logproba += logits[all_n, t_next]
55
56         input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s]
57
58
59 def masked_inplace_autoregression(
60     model,
61     batch_size,
62     input,
63     ar_mask,
64     seq_logproba,
65     temperature,
66     deterministic_synthesis,
67     forbidden_tokens=None,
68     logit_biases=None,
69     progress_bar_desc=None,
70     device=torch.device("cpu"),
71 ):
72     assert input.size() == ar_mask.size()
73
74     batches = zip(
75         input.split(batch_size),
76         ar_mask.split(batch_size),
77         seq_logproba.split(batch_size),
78     )
79
80     if progress_bar_desc is not None:
81         batches = tqdm.tqdm(
82             batches,
83             dynamic_ncols=True,
84             desc=progress_bar_desc,
85             total=(input.size(0) + batch_size - 1) // batch_size,
86         )
87
88     with torch.autograd.no_grad():
89         t = model.training
90         model.eval()
91
92         for input, ar_mask, seq_logproba in batches:
93             one_batch_masked_inplace_autoregression(
94                 model=model,
95                 input=input,
96                 ar_mask=ar_mask,
97                 seq_logproba=seq_logproba,
98                 temperature=temperature,
99                 deterministic_synthesis=deterministic_synthesis,
100             )
101
102         model.train(t)
103
104
105 ######################################################################
106
107
108 class QuizMachine:
109     def indices_forward_and_backward(self, quizzes):
110         i_forward = quizzes[:, 0] == self.token_forward
111         j_forward = quizzes[:, 1 + self.prompt_len] == self.token_forward
112         i_backward = quizzes[:, 0] == self.token_backward
113         j_backward = quizzes[:, 1 + self.answer_len] == self.token_backward
114         assert torch.logical_or(
115             torch.logical_and(i_forward, j_forward),
116             torch.logical_and(i_backward, j_backward),
117         ).all()
118         return i_forward, i_backward
119
120     def non_trivial(self, quizzes):
121         quizzes = quizzes.clone()
122         n_forward = quizzes[quizzes[:, 0] == self.token_forward]
123         n_backward = quizzes[:, 0] == self.token_backward
124         backward = quizzes[n_backward]
125         quizzes[n_backward] = self.reverse_time(quizzes[n_backward])
126         return torch.logical_not(
127             self.problem.trivial_prompts_and_answers(
128                 quizzes[:, 1 : 1 + self.prompt_len],
129                 quizzes[:, 2 + self.prompt_len :],
130             )
131         )
132
133     def reverse_time(self, quizzes):
134         i_forward, i_backward = self.indices_forward_and_backward(quizzes)
135
136         forward_to_backward = torch.cat(
137             [
138                 quizzes[:, 0:1],
139                 quizzes[:, 2 + self.prompt_len : 2 + self.prompt_len + self.answer_len],
140                 quizzes[:, 1 + self.prompt_len : 1 + self.prompt_len + 1],
141                 quizzes[:, 1 : 1 + self.prompt_len],
142             ],
143             dim=1,
144         )
145
146         forward_to_backward[:, 0] = self.token_backward
147         forward_to_backward[:, 1 + self.answer_len] = self.token_backward
148
149         backward_to_forward = torch.cat(
150             [
151                 quizzes[:, 0:1],
152                 quizzes[:, 2 + self.answer_len :],
153                 quizzes[:, 1 + self.answer_len : 2 + self.answer_len],
154                 quizzes[:, 1 : 1 + self.answer_len],
155             ],
156             dim=1,
157         )
158
159         backward_to_forward[:, 0] = self.token_forward
160         backward_to_forward[:, 1 + self.prompt_len] = self.token_forward
161
162         m = i_forward.long()[:, None]
163
164         return m * forward_to_backward + (1 - m) * backward_to_forward
165
166     def reverse_random_half_in_place(self, quizzes):
167         i = torch.rand(quizzes.size(0)) < 0.5
168         if i.any():
169             quizzes[i] = self.reverse_time(quizzes[i])
170
171     def make_ar_mask(self, quizzes, first=False):
172         i_forward, i_backward = self.indices_forward_and_backward(quizzes)
173
174         t = torch.arange(quizzes.size(1), device=quizzes.device)
175
176         if first:
177             m_forward = (t >= 1).long() * (t < 1 + self.prompt_len).long()
178             m_backward = (t >= 1).long() * (t < 1 + self.answer_len).long()
179         else:
180             m_forward = (t >= 2 + self.prompt_len).long()
181             m_backward = (t >= 2 + self.answer_len).long()
182
183         m = i_forward.long()[:, None]
184
185         return m * m_forward + (1 - m) * m_backward
186
187     def generate_token_sequences(self, nb):
188         prompts, answers = self.problem.generate_prompts_and_answers(nb)
189
190         if self.prompt_len is None:
191             self.prompt_len = prompts.size(1)
192
193         if self.answer_len is None:
194             self.answer_len = answers.size(1)
195
196         assert prompts.size(1) == self.prompt_len and answers.size(1) == self.answer_len
197
198         result = []
199
200         for prompt, answer in zip(prompts, answers):
201             a = [
202                 torch.tensor([self.token_forward]),
203                 prompt,
204                 torch.tensor([self.token_forward]),
205                 answer,
206             ]
207
208             result.append(torch.cat(a, dim=0)[None, :])
209
210         return torch.cat(result, dim=0)
211
212     def __init__(
213         self,
214         problem,
215         nb_train_samples,
216         nb_test_samples,
217         back_accuracy,
218         batch_size,
219         result_dir,
220         logger,
221         device=torch.device("cpu"),
222     ):
223         super().__init__()
224
225         v = problem.nb_token_values()
226         self.token_forward = v
227         self.token_backward = v + 1
228         self.nb_token_values = v + 2
229
230         self.problem = problem
231         self.back_accuracy = back_accuracy
232         self.batch_size = batch_size
233         self.device = device
234         self.logger = logger
235         self.prompt_len = None
236         self.answer_len = None
237
238         # self.train_w_quizzes = self.generate_token_sequences(nb_train_samples)
239         # self.reverse_random_half_in_place(self.train_w_quizzes)
240
241         # self.test_w_quizzes = self.generate_token_sequences(nb_test_samples).to(device)
242         # self.reverse_random_half_in_place(self.test_w_quizzes)
243
244         self.train_c_quizzes = []
245         self.test_c_quizzes = []
246
247         # if result_dir is not None:
248         # self.save_quizzes(
249         # result_dir,
250         # "culture_w_quizzes",
251         # self.train_w_quizzes[:72],
252         # )
253
254     def save_quizzes(
255         self,
256         result_dir,
257         filename_prefix,
258         quizzes,
259         mistakes=None,
260     ):
261         quizzes = quizzes.clone().to("cpu")
262         n_forward = quizzes[quizzes[:, 0] == self.token_forward]
263         n_backward = quizzes[:, 0] == self.token_backward
264         backward = quizzes[n_backward]
265         assert n_forward.size(0) + backward.size(0) == quizzes.size(0)
266         quizzes[n_backward] = self.reverse_time(quizzes[n_backward])
267
268         predicted_prompts = n_backward.long()
269         predicted_answers = 1 - predicted_prompts
270         if mistakes is not None:
271             # 0/-1/+1 ~ not-to-predict / predicted wrong / predicted correct
272             predicted_prompts *= mistakes.to("cpu")
273             predicted_answers *= mistakes.to("cpu")
274         else:
275             # 0/2 ~ not-to-predict / to predict
276             predicted_prompts *= 2
277             predicted_answers *= 2
278
279         self.problem.save_quizzes(
280             result_dir,
281             filename_prefix,
282             quizzes[:, 1 : 1 + self.prompt_len],
283             quizzes[:, 2 + self.prompt_len :],
284             predicted_prompts,
285             predicted_answers,
286         )
287
288     def batches(self, model, split="train", desc=None):
289         assert split in {"train", "test"}
290         if split == "train":
291             w_quizzes = model.train_w_quizzes
292             c_quizzes = self.train_c_quizzes
293         else:
294             w_quizzes = model.test_w_quizzes
295             c_quizzes = self.test_c_quizzes
296
297         if len(c_quizzes) > 0:
298             c_quizzes = torch.cat(c_quizzes, dim=0)
299             if c_quizzes.size(0) > w_quizzes.size(0) // 2:
300                 i = torch.randperm(c_quizzes.size(0))[: w_quizzes.size(0) // 2]
301                 c_quizzes = c_quizzes[i]
302
303             i = torch.randperm(w_quizzes.size(0))[
304                 : w_quizzes.size(0) - c_quizzes.size(0)
305             ]
306             w_quizzes = w_quizzes[i]
307
308             self.nb_batch_w_quizzes = w_quizzes.size(0)
309             self.nb_batch_c_quizzes = c_quizzes.size(0)
310
311             input = torch.cat([w_quizzes, c_quizzes], dim=0)
312         else:
313             input = w_quizzes
314             self.nb_batch_w_quizzes = w_quizzes.size(0)
315             self.nb_batch_c_quizzes = 0
316
317         # Shuffle
318         input = input[torch.randperm(input.size(0))]
319
320         if desc is None:
321             desc = f"epoch-{split}"
322         for batch in tqdm.tqdm(
323             input.split(self.batch_size), dynamic_ncols=True, desc=desc
324         ):
325             yield batch
326
327     def vocabulary_size(self):
328         return self.nb_token_values
329
330     def produce_results(
331         self, n_epoch, model, result_dir, deterministic_synthesis, nmax=1000
332     ):
333         def compute_accuracy(input, log_prefix=None):
334             ar_mask = self.make_ar_mask(input)
335             result = input.clone() * (1 - ar_mask)
336             seq_logproba = torch.empty(input.size(0), device=self.device)
337
338             masked_inplace_autoregression(
339                 model=model,
340                 batch_size=self.batch_size,
341                 input=result,
342                 ar_mask=ar_mask,
343                 seq_logproba=seq_logproba,
344                 temperature=1.0,
345                 deterministic_synthesis=deterministic_synthesis,
346                 progress_bar_desc=None,
347                 device=self.device,
348             )
349
350             correct = torch.empty(input.size(0), dtype=torch.int64, device=input.device)
351
352             n_forward = input[:, 0] == self.token_forward
353             n_backward = input[:, 0] == self.token_backward
354
355             correct[n_forward] = (
356                 (input[n_forward] == result[n_forward]).long().min(dim=1).values
357             )
358
359             if self.back_accuracy and n_backward.any():
360                 # accuracy of B->A*->B*=B instead of B->A*=A
361                 back_input = self.reverse_time(result[n_backward])
362                 back_input[:, 2 + self.prompt_len :] = input[
363                     n_backward, 1 : 1 + self.answer_len
364                 ]
365                 _, correct[n_backward] = compute_accuracy(back_input)
366
367             if log_prefix is not None:
368                 forward_nb_correct = correct[n_forward].sum()
369                 forward_nb_total = correct[n_forward].size(0)
370                 backward_nb_correct = correct[n_backward].sum()
371                 backward_nb_total = correct[n_backward].size(0)
372
373                 self.logger(
374                     f"{log_prefix}_forward_accuracy {n_epoch} model {model.id} nb_correct {forward_nb_correct} / {forward_nb_total} ({forward_nb_correct*100/forward_nb_total} %)"
375                 )
376
377                 self.logger(
378                     f"{log_prefix}_backward_accuracy {n_epoch} model {model.id} nb_correct {backward_nb_correct} / {backward_nb_total} ({backward_nb_correct*100/backward_nb_total} %)"
379                 )
380
381             return result, correct
382
383         compute_accuracy(model.train_w_quizzes[:nmax], log_prefix="train")
384
385         test_result, test_correct = compute_accuracy(
386             model.test_w_quizzes[:nmax], log_prefix="test"
387         )
388
389         main_test_accuracy = test_correct.sum() / test_correct.size(0)
390         self.logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}")
391
392         ##############################
393
394         self.save_quizzes(
395             result_dir,
396             f"culture_prediction_{n_epoch:04d}_{model.id:02d}",
397             quizzes=test_result[:72],
398             mistakes=test_correct[:72] * 2 - 1,
399         )
400
401         return main_test_accuracy
402
403     def renew_w_quizzes(self, model, nb, for_train=True):
404         input = model.train_w_quizzes if for_train else model.test_w_quizzes
405         nb = min(nb, input.size(0))
406         input[:-nb] = input[nb:].clone()
407         fresh_w_quizzes = self.generate_token_sequences(nb)
408         self.reverse_random_half_in_place(fresh_w_quizzes)
409         input[-nb:] = fresh_w_quizzes.to(self.device)
410
411     def store_c_quizzes(self, new_c_quizzes, for_train=True):
412         if for_train:
413             self.train_c_quizzes.append(new_c_quizzes)
414         else:
415             self.test_c_quizzes.append(new_c_quizzes)
416
417     def logproba_solution(self, models, c_quizzes):
418         logproba = c_quizzes.new_zeros(c_quizzes.size(0), len(models))
419
420         for model in models:
421             for input, l in zip(
422                 c_quizzes.split(self.batch_size), logproba.split(self.batch_size)
423             ):
424                 ar_mask = self.make_ar_mask(input)
425                 output = model(mygpt.BracketedSequence(input)).x
426                 ce = (
427                     F.cross_entropy(output.transpose(1, 2), input, reduction="none")
428                     * ar_mask
429                 )
430                 l[:, model.id] = -ce.sum(dim=-1)
431
432         return logproba
433
434     ###############################################################
435
436     def compute_correctness(
437         self,
438         c_quizzes,
439         models_for_validation,
440         bidirectional_validation=False,
441         deterministic_validation=True,
442     ):
443         if bidirectional_validation:
444             backward_c_quizzes = self.forward_to_backward(c_quizzes)
445
446         seq_logproba = torch.zeros(
447             c_quizzes.size(0),
448             max([m.id for m in models_for_validation]) + 1,
449             device=self.device,
450         )
451
452         nb_correct = 0
453
454         seq_logproba[...] = 0.0
455
456         for model in models_for_validation:
457             result = c_quizzes.clone()
458
459             ar_mask = self.make_ar_mask(result)
460
461             masked_inplace_autoregression(
462                 model=model,
463                 batch_size=self.batch_size,
464                 input=result,
465                 ar_mask=ar_mask,
466                 seq_logproba=seq_logproba[:, model.id],
467                 temperature=1.0,
468                 deterministic_synthesis=deterministic_validation,
469                 # progress_bar_desc="solving c_quizzes",
470                 device=self.device,
471             )
472
473             correct = (c_quizzes == result).long().min(dim=-1).values
474
475             if bidirectional_validation:
476                 backward_result = backward_c_quizzes.clone()
477
478                 ar_mask = self.make_ar_mask(backward_result)
479
480                 masked_inplace_autoregression(
481                     model=model,
482                     batch_size=self.batch_size,
483                     input=backward_result,
484                     ar_mask=ar_mask,
485                     seq_logproba=seq_logproba[:, model.id],
486                     temperature=1.0,
487                     deterministic_synthesis=deterministic_validation,
488                     # progress_bar_desc="solving backward c_quizzes",
489                     device=self.device,
490                 )
491
492                 backward_correct = (
493                     (backward_c_quizzes == backward_result).long().min(dim=-1).values
494                 )
495
496                 correct *= backward_correct
497
498             # endif
499
500             nb_correct += correct
501
502         return nb_correct, seq_logproba
503
504     ###############################################################
505
506     def generate_quizzes(self, nb, model_for_generation, temperature=1.0):
507         c_quizzes = torch.empty(
508             nb,
509             self.prompt_len + self.answer_len + 2,
510             device=self.device,
511             dtype=torch.int64,
512         )
513
514         seq_logproba = torch.zeros(nb, device=self.device)
515
516         # First, we generate the answer at high temperature
517
518         c_quizzes[:, 0] = self.token_backward
519         c_quizzes[:, 1 + self.answer_len] = self.token_backward
520
521         masked_inplace_autoregression(
522             model=model_for_generation,
523             batch_size=self.batch_size,
524             input=c_quizzes,
525             ar_mask=self.make_ar_mask(c_quizzes, first=True),
526             seq_logproba=seq_logproba,
527             temperature=temperature,
528             deterministic_synthesis=False,
529             device=self.device,
530         )
531
532         # Then, we generate the prompt at low temperature
533
534         masked_inplace_autoregression(
535             model=model_for_generation,
536             batch_size=self.batch_size,
537             input=c_quizzes,
538             ar_mask=self.make_ar_mask(c_quizzes),
539             seq_logproba=seq_logproba,
540             temperature=1 / temperature,
541             deterministic_synthesis=False,
542             device=self.device,
543         )
544
545         # Then we return the quizz, and re-generate the response, now
546         # at low temperature
547
548         c_quizzes = self.reverse_time(c_quizzes)
549
550         masked_inplace_autoregression(
551             model=model_for_generation,
552             batch_size=self.batch_size,
553             input=c_quizzes,
554             ar_mask=self.make_ar_mask(c_quizzes),
555             seq_logproba=seq_logproba,
556             temperature=1 / temperature,
557             deterministic_synthesis=False,
558             device=self.device,
559         )
560
561         return c_quizzes