X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=d21e2648466a606b3067fc680feb7305a6b95781;hb=8ea809c43242d3a2e063692105919a86c3f6fe6b;hp=1ea3b5d588d37e4c89ca4d06f74d20356b7a2fcf;hpb=4f489998d6e73680c3a031e8932a7678c16268e3;p=picoclvr.git

diff --git a/tasks.py b/tasks.py
index 1ea3b5d..d21e264 100755
--- a/tasks.py
+++ b/tasks.py
@@ -71,7 +71,7 @@ class Task:
 
 
 class TaskFromFile(Task):
-    def tensorize(self, pairs):
+    def tensorize(self, pairs, shuffle):
         len_max = max([len(x[0]) for x in pairs])
 
         input = torch.cat(
@@ -98,6 +98,12 @@ class TaskFromFile(Task):
             0,
         ).to("cpu")
 
+        if shuffle:
+            print("SHUFFLING!")
+            i = torch.randperm(input.size(0))
+            input = input[i].contiguous()
+            pred_mask = pred_mask[i].contiguous()
+
         return input, pred_mask
 
     # trim all the tensors in the tuple z to remove as much token from
@@ -117,35 +123,52 @@ class TaskFromFile(Task):
 
     def __init__(
         self,
-        filename,
+        train_filename,
+        test_filename,
         nb_train_samples,
         nb_test_samples,
         batch_size,
+        shuffle=False,
         device=torch.device("cpu"),
     ):
         self.batch_size = batch_size
         self.device = device
 
-        pairs = []
-        with open(filename, "r") as f:
-            for _ in range(nb_train_samples + nb_test_samples):
-                sequence = f.readline().strip()
-                pred_mask = f.readline().strip()
-                assert len(sequence) == len(pred_mask)
-                assert set(pred_mask).issubset({"0", "1", "2"}), f"{set(pred_mask)}"
-                pairs.append((sequence, pred_mask))
-
-        symbols = ["#"] + list(set("".join([x[0] for x in pairs])) - set(["#"]))
+        def read_file(filename, nb=-1):
+            pairs = []
+            with open(filename, "r") as f:
+                while True:
+                    sequence = f.readline().strip()
+                    if not sequence:
+                        break
+                    pred_mask = f.readline().strip()
+                    assert len(sequence) == len(pred_mask)
+                    assert set(pred_mask).issubset({"0", "1", "2"}), f"{set(pred_mask)}"
+                    pairs.append((sequence, pred_mask))
+                    if len(pairs) == nb:
+                        break
+
+            if nb > 0:
+                pairs = pairs[:nb]
+                assert len(pairs) == nb
+
+            return pairs
+
+        train_pairs = read_file(train_filename, nb_train_samples)
+        test_pairs = read_file(test_filename, nb_test_samples)
+
+        symbols = ["#"] + list(
+            set("".join([x[0] for x in train_pairs + test_pairs])) - set(["#"])
+        )
         self.char2id = dict([(c, n) for n, c in enumerate(symbols)])
         self.id2char = dict([(n, c) for c, n in self.char2id.items()])
 
         self.train_input, self.train_pred_masks = self.tensorize(
-            pairs[:nb_train_samples]
+            train_pairs, shuffle=shuffle
+        )
+        self.test_input, self.test_pred_masks = self.tensorize(
+            test_pairs, shuffle=shuffle
         )
-        self.test_input, self.test_pred_masks = self.tensorize(pairs[nb_train_samples:])
-
-        assert self.train_input.size(0) == nb_train_samples
-        assert self.test_input.size(0) == nb_test_samples
 
     def batches(self, split="train", nb_to_use=-1, desc=None):
         assert split in {"train", "test"}
@@ -176,7 +199,7 @@ class TaskFromFile(Task):
 
         logger(f"----------------------------------------------------------")
 
-        for e in self.tensor2str(result[:10]):
+        for e in self.tensor2str(result[:50]):
             logger(f"test_before {e}")
 
         masked_inplace_autoregression(
@@ -190,7 +213,7 @@ class TaskFromFile(Task):
 
         logger(f"----------------------------------------------------------")
 
-        for e, c in zip(self.tensor2str(result[:10]), self.tensor2str(correct[:10])):
+        for e, c in zip(self.tensor2str(result[:50]), self.tensor2str(correct[:50])):
             logger(f"test_after  {e}")
             logger(f"correct     {c}")