+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):