* Make padding mask optional.

* Perform preprocessing in parallel.

seq2seq/do_run.sh

@@ -3,7 +3,8 @@ python run_seq2seq_flax.py \
 	--train_file /data/CC12M/encoded-small-train.tsv \
 	--validation_file /data/CC12M/encoded-small-valid.tsv \
 	--output_dir output \
-	--per_device_train_batch_size 16 \
-	--per_device_eval_batch_size 16 \
+	--per_device_train_batch_size 24 \
+	--per_device_eval_batch_size 24 \
+	--preprocessing_num_workers 48 \
 	--do_train \
 	--do_eval \

seq2seq/run_seq2seq_flax.py

@@ -482,8 +482,6 @@ def main():
 
         # We need decoder_attention_mask so we can ignore pad tokens from loss
         # TODO: I don't believe we need "decoder_attention_mask" in this case because all labels have same length
-        # However, we need to provide a mask or modify the compute_loss function, which relies on having one
-        model_inputs["decoder_attention_mask"] = np.ones(labels.shape)
         #model_inputs["decoder_attention_mask"] = labels["attention_mask"]
 
         return model_inputs
@@ -647,6 +645,9 @@ def main():
         loss = optax.softmax_cross_entropy(logits, soft_labels)
         loss = loss - normalizing_constant
 
+        if padding_mask is None:
+            padding_mask = np.ones(loss.shape)
+
         # ignore padded tokens from loss
         loss = loss * padding_mask
         loss = loss.sum() / padding_mask.sum()
@@ -659,7 +660,8 @@ def main():
         def compute_loss(params):
             labels = batch.pop("labels")
             logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
-            loss = loss_fn(logits, labels, batch["decoder_attention_mask"], label_smoothing_factor)
+            padding_mask = batch.get("decoder_attention_mask", None)
+            loss = loss_fn(logits, labels, padding_mask, label_smoothing_factor)
             return loss
 
         grad_fn = jax.value_and_grad(compute_loss)
@@ -677,7 +679,8 @@ def main():
     def eval_step(params, batch, label_smoothing_factor=0.0):
         labels = batch.pop("labels")
         logits = model(**batch, params=params, train=False)[0]
-        loss = loss_fn(logits, labels, batch["decoder_attention_mask"], label_smoothing_factor)
+        padding_mask = batch.get("decoder_attention_mask", None)
+        loss = loss_fn(logits, labels, padding_mask, label_smoothing_factor)
 
         # summarize metrics
         metrics = {"loss": loss}