X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=problems.py;h=68a46b3cdb89dc79f892b5d3339927c039193e6a;hb=0f4c86c0e7730db4147f136df5aeb5528fc943a0;hp=aa3acf038559195f4e68d7fd7e1625ffb41e332d;hpb=687d5b2d9f465577665991b84faec7c789685271;p=picoclvr.git diff --git a/problems.py b/problems.py index aa3acf0..68a46b3 100755 --- a/problems.py +++ b/problems.py @@ -17,6 +17,96 @@ 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 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 ] ) #################### @@ -47,51 +137,29 @@ class ProblemTwoTargets(Problem): 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), 1 + ( + 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 ProblemLenId(Problem): - def __init__(self, len_max=10): - self.len_max = len_max - - 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 - ) - ar_mask = (sequences == 11).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 @@ -108,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 @@ -126,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) @@ -142,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 @@ -220,7 +287,6 @@ class ProblemAddition(Problem): if __name__ == "__main__": - p = ProblemTwoTargets(12, 4) - s, m = p.generate_sequences(10) - for x in s: - print(p.seq2str(x)) + p = ProblemTwoCuts(12) + s, m = p.generate_sequences(10000) + print(p.compute_nb_correct(None, None, s))