X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=problems.py;h=68a46b3cdb89dc79f892b5d3339927c039193e6a;hb=0f4c86c0e7730db4147f136df5aeb5528fc943a0;hp=1f4098a831ef120cb3f97ed22a0b6ba8dc090c1a;hpb=6c8bed86221baae24a7c2aaaa41c009444efb5c9;p=picoclvr.git diff --git a/problems.py b/problems.py index 1f4098a..68a46b3 100755 --- a/problems.py +++ b/problems.py @@ -17,43 +17,149 @@ 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 ProblemLenId(Problem): - def __init__(self, nb_sentences=100, len_max=5): - self.len_max = len_max +class ProblemTwoCuts(Problem): + def __init__(self, len_total=50, nb_values=100, global_constraint=True): + self.len_total = len_total + self.nb_values = nb_values + self.global_constraint = global_constraint + + def generate_sequences_internal(self, nb): + return u,v,a,b,c + + def generate_sequences(self,nb): + + u = torch.randint(self.len_total, (nb,)) + v = torch.randint(self.len_total, (nb,)) + + a = torch.randint(self.nb_values, (nb,)) + b = torch.randint(self.nb_values, (nb,)) + c = torch.randint(self.nb_values, (nb,)) + + while True: + to_compute = torch.logical_or(u>=v-self.len_total//10,u=v).long().sum() == 0 + assert (a==b).long().sum() == 0 + assert (a==c).long().sum() == 0 + assert (c==b).long().sum() == 0 + + t = torch.arange(self.len_total) + seq = (t[None,:] < u[:,None]).long() * a[:,None] + \ + (t[None,:] >= u[:,None]).long() * (t[None,:] < v[:,None]).long() * b[:,None] + \ + (t[None,:] >= v[:,None]).long() * c[:,None] + + 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 + i = torch.arange(result.size(1), device=result.device) + + for k in range(nb_total): + s = result[k] + a = s[0] + uu = (s != a).nonzero() + if uu.size(0) > 0: + u = uu.min() + b = s[u] + vv = torch.logical_and(s != b, i >= u).nonzero() + if vv.size(0) > 0: + v = vv.min() + c = s[v] + ww = torch.logical_and(s != c, i >= v).nonzero() + if ww.size(0) == 0: + if not self.global_constraint or (a*u+b*(v-u)+c*(self.len_total-v)) // self.len_total == self.nb_values//2: + nb_correct += 1 + + return nb_total, nb_correct + + def seq2str(self, seq): + return " ".join( [ f"{x:02d}" for x in seq ] ) + +#################### + + +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_max * 3 + 3)[None, :] - l = torch.randint(self.len_max, (2, nb))[:, :, None] + 1 - i = torch.randint(10, (2, nb))[:, :, None] - a = l[0] - b = l[0] + 1 + l[1] - c = l[0] + 1 + l[1] + 1 + l[0] - sequences = ( - (k < a) * i[0] - + (k == a) * 10 - + (k > a) * (k < b) * i[1] - + (k == b) * 11 - + (k > b) * (k < c) * i[1] - + (k == c) * 12 - + (k > c) * 13 + 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() ) - ar_mask = (sequences == 11).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) + return "".join("0123456789-+|"[x.item()] for x in seq) #################### -class ProblemLevel0(Problem): - def __init__(self, nb_sentences=100, len_prompt=5, len_result=5): +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 @@ -70,8 +176,8 @@ class ProblemLevel0(Problem): #################### -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 @@ -88,7 +194,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) @@ -104,7 +209,7 @@ class ProblemLevel1(Problem): #################### -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 @@ -181,18 +286,7 @@ 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 = ProblemTwoCuts(12) + s, m = p.generate_sequences(10000) + print(p.compute_nb_correct(None, None, s))