Update

[beaver.git] / beaver.py

diff --git a/beaver.py b/beaver.py
index 6a6343d..5407859 100755 (executable)
--- a/beaver.py
+++ b/beaver.py
@@ -66,6 +66,8 @@ parser.add_argument("--deterministic_synthesis", action="store_true", default=Fa
 
 parser.add_argument("--random_regression_order", action="store_true", default=False)
 
+parser.add_argument("--noncausal_prompt", action="store_true", default=False)
+
 parser.add_argument("--no_checkpoint", action="store_true", default=False)
 
 parser.add_argument("--overwrite_results", action="store_true", default=False)
@@ -131,10 +133,10 @@ for n in vars(args):
 ######################################################################
 
 
-def generation_order(x, fixed_len=0):
+def generation_order(x, prompt_len=0):
     if args.random_regression_order:
         order = torch.rand(x.size(), device=x.device)
-        order[:, :fixed_len] = torch.arange(-fixed_len, 0, device=x.device)
+        order[:, :prompt_len] = torch.arange(-prompt_len, 0, device=x.device)
         order = order.sort(1).indices
     else:
         order = (
@@ -154,13 +156,13 @@ def reorder(x, order, reverse=False):  # x is NxTxD1x...xDk, order is NxT'
     return v
 
 
-def shuffle(x, fixed_len):
-    order = generation_order(x, fixed_len)
+def shuffle(x, prompt_len):
+    order = generation_order(x, prompt_len)
     return reorder(x, order), order
 
 
-def eval_mygpt(model, input, mode="standard", fixed_len=0):
-    x, order = shuffle(input, fixed_len)
+def eval_mygpt(model, input, mode="standard", prompt_len=0):
+    x, order = shuffle(input, prompt_len)
     x = model(mygpt.BracketedSequence(x), mode=mode, order=order).x
     return reorder(x, order, reverse=True)
 
@@ -193,7 +195,7 @@ def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None)
 ######################################################################
 
 
-def compute_perplexity(model, task, fixed_len, split="train"):
+def compute_perplexity(model, task, prompt_len, split="train"):
     with torch.autograd.no_grad():
         t = model.training
         model.eval()
@@ -202,8 +204,12 @@ def compute_perplexity(model, task, fixed_len, split="train"):
 
         for input in task.batches(split=split):
             input = input.to(device)
-            output = eval_mygpt(model, input, fixed_len=fixed_len)
-            loss = F.cross_entropy(output.transpose(1, 2), input)
+            output = eval_mygpt(model, input, prompt_len=prompt_len)
+            if args.noncausal_prompt:
+                d = input.size(1) // 2
+                loss = F.cross_entropy(output[:, d:].transpose(1, 2), input[:, d:])
+            else:
+                loss = F.cross_entropy(output.transpose(1, 2), input)
             acc_loss += loss.item() * input.size(0)
             nb_samples += input.size(0)
 
@@ -267,7 +273,7 @@ def oneshot(gpt, task):
         acc_train_loss, nb_train_samples = 0, 0
         for mazes, policies in task.policy_batches(split="train"):
             output_gpt = eval_mygpt(
-                gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width
+                gpt, mazes, mode=args.oneshot_input, prompt_len=task.height * task.width
             )
             output = model(output_gpt)
 
@@ -282,7 +288,7 @@ def oneshot(gpt, task):
         acc_test_loss, nb_test_samples = 0, 0
         for mazes, policies in task.policy_batches(split="test"):
             output_gpt = eval_mygpt(
-                gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width
+                gpt, mazes, mode=args.oneshot_input, prompt_len=task.height * task.width
             )
             output = model(output_gpt)
             loss = compute_loss(mazes, output, policies, task.height, task.width)
@@ -297,7 +303,7 @@ def oneshot(gpt, task):
         mazes = task.test_input[:32, : task.height * task.width]
         policies = task.test_policies[:32]
         output_gpt = eval_mygpt(
-            gpt, mazes, mode=args.oneshot_input, fixed_len=task.height * task.width
+            gpt, mazes, mode=args.oneshot_input, prompt_len=task.height * task.width
         )
         output = model(output_gpt)
         if args.oneshot_output == "policy":
@@ -517,6 +523,20 @@ log_string(f"vocabulary_size {vocabulary_size}")
 
 ##############################
 
+
+def noncausal_prompt_amm_generator(d):
+    q = torch.arange(d)[:, None]
+    k = torch.arange(d)[None, :]
+    s = args.maze_height * args.maze_width
+    #    return torch.logical_and(q < k, torch.logical_or(q >= s, k >= s))
+    return q < k
+
+
+amm_generator = None
+
+if args.noncausal_prompt:
+    amm_generator = noncausal_prompt_amm_generator
+
 model = mygpt.MyGPT(
     vocabulary_size=vocabulary_size,
     dim_model=args.dim_model,
@@ -526,6 +546,7 @@ model = mygpt.MyGPT(
     nb_blocks=args.nb_blocks,
     causal=True,
     dropout=args.dropout,
+    amm_generator=amm_generator,
 )
 
 model.to(device)
@@ -597,10 +618,10 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}")
 if nb_epochs_finished >= args.nb_epochs:
     n_epoch = nb_epochs_finished
     train_perplexity = compute_perplexity(
-        model, task, fixed_len=task.height * task.width, split="train"
+        model, task, prompt_len=task.height * task.width, split="train"
     )
     test_perplexity = compute_perplexity(
-        model, task, fixed_len=task.height * task.width, split="test"
+        model, task, prompt_len=task.height * task.width, split="test"
     )
 
     log_string(
@@ -631,10 +652,12 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
 
     for input in task.batches(split="train"):
         input = input.to(device)
-        output = eval_mygpt(
-            model, input, mode=args.oneshot_input, fixed_len=task.height * task.width
-        )
-        loss = F.cross_entropy(output.transpose(1, 2), input)
+        output = eval_mygpt(model, input, prompt_len=task.height * task.width)
+        if args.noncausal_prompt:
+            d = input.size(1) // 2
+            loss = F.cross_entropy(output[:, d:].transpose(1, 2), input[:, d:])
+        else:
+            loss = F.cross_entropy(output.transpose(1, 2), input)
         acc_train_loss += loss.item() * input.size(0)
         nb_train_samples += input.size(0)
 
@@ -644,7 +667,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
 
     train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
     test_perplexity = compute_perplexity(
-        model, task, fixed_len=task.height * task.width, split="test"
+        model, task, prompt_len=task.height * task.width, split="test"
     )
 
     log_string(