Add data util to mask data

LLNL · Mar 5, 2024 · 4058954 · 4058954
1 parent e904f87
commit 4058954
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/applications/FLASK/Transformer/datasets/pretokenize/data_utils.py b/applications/FLASK/Transformer/datasets/pretokenize/data_utils.py
@@ -0,0 +1,9 @@
+import numpy as np
+
+
+def random_zero_array(arr, p, mask_token):
+    """
+    Randomly zero out elements of an array with probability p
+    """
+    mask = np.random.choice([0, 1], size=arr.shape, p=[p, 1 - p])
+    return arr * mask + mask_token * (1 - mask)