X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=problems.py;h=7aa59bea856c205d232ec1bf49c80f18ef20bed1;hb=16e7952b7cc32ca21498fa3a12fb79f679ea8c21;hp=516158795d047002ef70163bea7e03b68b9b9181;hpb=16cb07f99cf770fb4e97824f874a68cbddd4c1cf;p=picoclvr.git diff --git a/problems.py b/problems.py index 5161587..7aa59be 100755 --- a/problems.py +++ b/problems.py @@ -17,12 +17,123 @@ class Problem: def seq2str(self, seq): return "[NOT IMPLEMENTED]" + def compute_nb_correct(self, input, ar_mask, result): + nb_total = ar_mask.sum().item() + nb_correct = ((result == input).long() * ar_mask).sum().item() + return nb_total, nb_correct #################### -class ProblemLevel0(Problem): - def __init__(self, nb_sentences=100, len_prompt=5, len_result=5): +class ProblemDegradation(Problem): + def __init__(self, nb_state_tokens=7, nb_time_steps=10, value_max=100, hard=False): + self.nb_state_tokens = nb_state_tokens + self.nb_time_steps = nb_time_steps + self.value_max = value_max + self.hard = hard + + def generate_sequences(self,nb): + + x = (torch.rand(nb,self.nb_state_tokens).sort(dim=-1).indices == 0).long() * self.value_max + seq = [x] + + for t in range(self.nb_time_steps-1): + v = torch.rand(x.size()) * (x > 0).float() + u = (v.max(dim=-1,keepdim=True).values == v).long() + n = (u*x*torch.rand(x.size())).long().sum(dim=-1,keepdim=True) // 2 + x = x + n * (u.roll(shifts=-1,dims=-1) - 2 * u + u.roll(shifts=1,dims=-1)) + seq.append(x) + + if self.hard: seq.reverse() + + seq = torch.cat(seq,dim=1) + return seq,seq.new_full(seq.size(), 1, dtype=torch.int64) + + def compute_nb_correct(self, input, ar_mask, result): + nb_total = result.size(0) + nb_correct = 0 + e=result.new_zeros(self.nb_state_tokens) + + for seq in result: + states = list(seq.split(self.nb_state_tokens)) + if self.hard: + states.reverse() + + d = states[0] + j=d.sort(descending=True).indices[0] + e.zero_() + e[j]=self.value_max + if (d-e).abs().sum() == 0: + nb_errors = 0 + for k in range(len(states)-1): + d=states[k]-states[k+1] + j=d.sort(descending=True).indices[0] + e.zero_() + e[j]=d[j] + e[(j+1)%e.size(0)]=-d[j]//2 + e[(j-1)%e.size(0)]=-d[j]//2 + if (d-e).abs().sum() > 0: + nb_errors += 1 + if nb_errors == 0: + nb_correct += 1 + + return nb_total, nb_correct + + def seq2str(self, seq): + return " | ".join( [ " ".join([f"{x:02d}" for x in s ]) for s in seq.split(self.nb_state_tokens) ] ) + +#################### + + +class ProblemTwoTargets(Problem): + def __init__(self, len_total=10, len_targets=3): + assert len_targets >= 3 + assert len_total >= 3 * len_targets - 1 + self.len_total = len_total + self.len_targets = len_targets + + def generate_sequences(self, nb): + k = torch.arange(self.len_total)[None, :] + s = torch.randint(10, (nb, self.len_total)) + l = torch.rand(nb, self.len_total) + l = l * (k <= self.len_total - self.len_targets).long() + k1 = l.argmax(dim=1, keepdim=True) + m = (k != k1).long() * (k != k1 + self.len_targets - 1).long() + s = s * m + 10 * (1 - m) + l = l * ( + 1 + - (k + self.len_targets - 1 >= k1).long() + * (k < k1 + self.len_targets).long() + ) + k2 = l.argmax(dim=1, keepdim=True) + m = (k != k2).long() * (k != k2 + self.len_targets - 1).long() + s = s * m + 11 * (1 - m) + a1 = s.gather(dim=1, index=k1 + 1 + torch.arange(self.len_targets - 2)[None, :]) + a2 = s.gather(dim=1, index=k2 + 1 + torch.arange(self.len_targets - 2)[None, :]) + sequences = torch.cat( + ( + s, + torch.full((nb, 1), 12), + a1, + torch.full((nb, 1), 12), + a2, + torch.full((nb, 1), 12), + ), + 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) + + +#################### + + +class ProblemByHeart(Problem): + def __init__(self, nb_sentences=100, len_prompt=8, len_result=8): self.seq = torch.randint(10, (nb_sentences, len_prompt + 1 + len_result)) self.seq[:, len_prompt] = 10 @@ -32,9 +143,15 @@ class ProblemLevel0(Problem): 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) + -class ProblemLevel1(Problem): - def __init__(self, nb_operators=100, len_source=5, len_result=8): +#################### + + +class ProblemLearnOperator(Problem): + def __init__(self, nb_operators=100, len_source=6, len_result=9): self.len_source = len_source self.len_result = len_result self.len_nb_operator = int(math.log(nb_operators) / math.log(10)) + 1 @@ -51,7 +168,6 @@ class ProblemLevel1(Problem): // 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) @@ -64,7 +180,10 @@ class ProblemLevel1(Problem): return "".join("0123456789|>"[x.item()] for x in seq) -class ProblemLevel2(Problem): +#################### + + +class ProblemGuessOperator(Problem): def __init__(self, len_source=5, len_result=8): self.len_source = len_source self.len_result = len_result @@ -141,18 +260,9 @@ class ProblemAddition(Problem): return "".join(self.id2char[x.item()] for x in seq) -# class ProblemUnion(Problem): -# problems = [ProblemByheart()] -# nb_common_codes = 100 - -# def generate_sequences(nb_samples): -# problem_indexes = torch.randint(len(problems), (nb_samples,)) -# nb_samples_per_problem = torch.one_hot(problem_indexes).sum(0) -# print(f"{nb_samples_per_problem}") -# all_seq = [] -# for nb, p in zip(nb_samples_per_problem, problems): -# all_seq.append(p.generate_sequences(nb_samples_per_problem[nb])) -# return all_seq - -# for strain, stest in zip(train_seq, test_seq): -# s = torch.cat((strain, stest), 0) +if __name__ == "__main__": + p = ProblemDegradation(hard=False) + s, m = p.generate_sequences(10000) + for x in s[:100]: + print(p.seq2str(x)) + print(p.compute_nb_correct(None, None, s))