Fix sentinel token IDs in data collator for Flax T5 pretraining script (#14477)

2021-11-29 08:30:17 -08:00
parent 2bd950ca47
commit 8332327dca
1 changed files with 1 additions and 1 deletions
--- a/examples/flax/language-modeling/run_t5_mlm_flax.py
+++ b/examples/flax/language-modeling/run_t5_mlm_flax.py
@@ -291,7 +291,7 @@ class FlaxDataCollatorForT5MLM:
        start_indices[:, 0] = mask_indices[:, 0]

        sentinel_ids = np.where(start_indices != 0, np.cumsum(start_indices, axis=-1), start_indices)
-        sentinel_ids = np.where(sentinel_ids != 0, (sentinel_ids + self.tokenizer.vocab_size - 1), 0)
+        sentinel_ids = np.where(sentinel_ids != 0, (len(self.tokenizer) - sentinel_ids), 0)
        sentinel_ids -= mask_indices - start_indices

        return sentinel_ids