Playgrounds/nanosocrates-train-experiment-2.py

@@ -51,6 +51,7 @@ VALIDATION_STEPS = 50
 CHECKPOINT_STEPS = VALIDATION_STEPS * 4
 PATIENCE = 4
 CURRENT_EPOCH = 0
+VERBOSE = False
 
 SOS_TOKEN = TOKENANO.encode("<SOS>")[0]
 
@@ -125,6 +126,7 @@ while current_epoch < MAX_EPOCHS:
 
     batch_counter = 0
 
+    if VERBOSE:
         print(f"EPOCH {current_epoch} STARTING")
 
     for batch in TRAIN_BATCHER.batch(MINI_BATCH_SIZE):
@@ -146,6 +148,7 @@ while current_epoch < MAX_EPOCHS:
         dec_x[:, 1:] = tgt[:, :-1]
         dec_x_pad = dec_x.eq(PAD_TOKEN)
 
+        if VERBOSE:
             print(f"\tBATCH {batch_counter} Starting")
 
         # Task 1 and Task 2
@@ -173,6 +176,7 @@ while current_epoch < MAX_EPOCHS:
         # Task 3
         if tasktype == Batch.TaskType.MASKING:
 
+            if VERBOSE:
                 print(f"\tExecuting TASK 3 - BATCH {batch_counter}")
 
             encoder_only_optim.zero_grad()
@@ -192,6 +196,7 @@ while current_epoch < MAX_EPOCHS:
         # Task 4
         if tasktype == Batch.TaskType.COMPLETATION:
 
+            if VERBOSE:
                 print(f"\tExecuting TASK 4 - BATCH {batch_counter}")
 
             decoder_only_optim.zero_grad()