X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=af94979e937d852912dc79c01ada589436de461f;hb=d2844d7a2d09ef38dc6f62d5e131059cccc872c5;hp=1b0d39a004436466724e13144599102e4e96b3a3;hpb=ef3bef5253ff719953dfffff28d4122c19acdd77;p=picoclvr.git

diff --git a/main.py b/main.py
index 1b0d39a..af94979 100755
--- a/main.py
+++ b/main.py
@@ -12,7 +12,7 @@ from torch import nn
 from torch.nn import functional as F
 
 import ffutils
-import mygpt, tasks
+import mygpt, tasks, problems
 
 ######################################################################
 
@@ -335,19 +335,19 @@ picoclvr_pruner_eval = (
 
 if args.task == "sandbox":
     if args.sandbox_level == 0:
-        problem = tasks.ProblemLevel0(
+        problem = problems.ProblemLevel0(
             nb_sentences=args.sandbox_levels_nb_items,
             len_prompt=args.sandbox_levels_len_source,
             len_result=args.sandbox_levels_len_result,
         )
     elif args.sandbox_level == 1:
-        problem = tasks.ProblemLevel1(
+        problem = problems.ProblemLevel1(
             nb_operators=args.sandbox_levels_nb_items,
             len_source=args.sandbox_levels_len_source,
             len_result=args.sandbox_levels_len_result,
         )
     elif args.sandbox_level == 2:
-        problem = tasks.ProblemLevel2(
+        problem = problems.ProblemLevel2(
             len_source=args.sandbox_levels_len_source,
             len_result=args.sandbox_levels_len_result,
         )
@@ -355,8 +355,9 @@ if args.task == "sandbox":
         raise ValueError(f"Unknown sandbox level {args.sandbox_level}")
 
     task = tasks.SandBox(
-        problem,
-        # tasks.ProblemAddition(zero_padded=False, inverted_result=False),
+        # problem,
+        # problems.ProblemAddition(zero_padded=False, inverted_result=False),
+        problems.ProblemLenId(len_max=args.sandbox_levels_len_source),
         nb_train_samples=args.nb_train_samples,
         nb_test_samples=args.nb_test_samples,
         batch_size=args.batch_size,
@@ -541,34 +542,37 @@ token_probas = token_count / token_count.sum()
 entropy = -torch.xlogy(token_probas, token_probas).sum()
 train_set_perplexity = math.exp(entropy)
 
-##############################
-
+######################################################################
 # A bit of paranoia never hurts
 
-train_examples = {}
 
+def subsets_as_tuples(batches, cs):
+    s = set()
+    for batch in batches:
+        for x in batch:
+            s.add(tuple([v.item() for v in x]))
+            if len(s) == cs:
+                yield s
+                s = set()
+    yield s
 
-for input in task.batches(split="train"):
-    assert input.dim() == 2 and input.dtype == torch.int64
-    for x in input:
-        train_examples[x.sum().item()] = x
-
-nb_total, nb_collisions = 0, 0
-for input in task.batches(split="test"):
-    assert input.dim() == 2 and input.dtype == torch.int64
-    for x in input:
-        nb_total += 1
-        y = train_examples.get(x.sum().item())
-        if y is not None:
-            if x.size() == y.size() and (x - y).abs().sum() == 0:
-                nb_collisions += 1
-
-del train_examples
+
+nb_test, nb_in_train = 0, 0
+for test_subset in subsets_as_tuples(task.batches(split="test"), 25000):
+    in_train = set()
+    for train_subset in subsets_as_tuples(task.batches(split="train"), 25000):
+        in_train.update(test_subset.intersection(train_subset))
+    nb_in_train += len(in_train)
+    nb_test += len(test_subset)
 
 log_string(
-    f"data_check {nb_collisions*100/nb_total:.02f}% ({nb_collisions}/{nb_total}) of test samples are in the train set"
+    f"data_check {nb_in_train*100/nb_test:.02f}% ({nb_in_train}/{nb_test}) of test samples are in the train set"
 )
 
+assert (
+    nb_in_train <= nb_test // 100
+), "More than 1% of test samples are in the train set"
+
 ##############################
 
 if args.learning_rate_schedule == "cos":