index_update

Author: zheyuye

scripts/pretraining/pretraining_utils.py

@@ -493,13 +493,14 @@ def dynamic_masking(self, F, input_ids, valid_lengths):
         valid_candidates = valid_candidates.astype(np.float32)
         num_masked_position = F.np.maximum(
             1, F.np.minimum(N, round(valid_lengths * self._mask_prob)))
+
         # The categorical distribution takes normalized probabilities as input
         # softmax is used here instead of log_softmax
         sample_probs = F.npx.softmax(
-            self._proposal_distribution * valid_candidates, axis=-1)  # (B, L)
-        # Top-k Sampling is an alternative solution to avoid duplicates positions
+           self._proposal_distribution * valid_candidates, axis=-1)  # (B, L)
         masked_positions = F.npx.random.categorical(
             sample_probs, shape=N, dtype=np.int32)
+
         masked_weights = F.npx.sequence_mask(
             F.np.ones_like(masked_positions),
             sequence_length=num_masked_position,
@@ -508,7 +509,7 @@ def dynamic_masking(self, F, input_ids, valid_lengths):
         length_masks = F.npx.sequence_mask(
             F.np.ones_like(input_ids, dtype=np.float32),
             sequence_length=valid_lengths,
-            use_sequence_length=True, axis=1, value=0).astype(np.float32)
+            use_sequence_length=True, axis=1, value=0)
         unmasked_tokens = select_vectors_by_position(
             F, input_ids, masked_positions) * masked_weights
         masked_weights = masked_weights.astype(np.float32)
@@ -518,11 +519,8 @@ def dynamic_masking(self, F, input_ids, valid_lengths):
                 F.np.zeros_like(masked_positions),
                 F.np.ones_like(masked_positions)) > self._mask_prob) * masked_positions
         #  deal with multiple zeros
-        filled = F.np.where(
-            replaced_positions,
-            self.vocab.mask_id,
-            masked_positions).astype(np.int32)
-        masked_input_ids, _ = updated_vectors_by_position(F, input_ids, filled, replaced_positions)
+        filled = F.np.where(replaced_positions, self.vocab.mask_id, masked_positions)
+        masked_input_ids = updated_vectors_by_position(F, input_ids, filled, replaced_positions)
         masked_input = self.MaskedInput(input_ids=masked_input_ids,
                                         masks=length_masks,
                                         unmasked_tokens=unmasked_tokens,

src/gluonnlp/models/electra.py

@@ -36,7 +36,7 @@
 from mxnet import use_np
 from mxnet.gluon import HybridBlock, nn
 from ..registry import BACKBONE_REGISTRY
-from ..op import gumbel_softmax, select_vectors_by_position, updated_vectors_by_position
+from ..op import gumbel_softmax, select_vectors_by_position, add_vectors_by_position, updated_vectors_by_position
 from ..base import get_model_zoo_home_dir, get_repo_model_zoo_url, get_model_zoo_checksum_dir
 from ..layers import PositionalEmbedding, get_activation
 from .transformer import TransformerEncoderLayer
@@ -833,13 +833,14 @@ def get_corrupted_tokens(self, F, inputs, unmasked_tokens, masked_positions, log
             use_np_gumbel=False)
         corrupted_tokens = F.np.argmax(prob, axis=-1).astype(np.int32)
 
-        # Following the Official electra to deal with duplicate positions as
-        # https://github.com/google-research/electra/issues/41
-        original_data, updates_mask = updated_vectors_by_position(F,
+        original_data = updated_vectors_by_position(F,
             inputs, unmasked_tokens, masked_positions)
-        fake_data, _ = updated_vectors_by_position(F,
+        fake_data = updated_vectors_by_position(F,
             inputs, corrupted_tokens, masked_positions)
-
+        updates_mask = add_vectors_by_position(F, F.np.zeros_like(inputs),
+                F.np.ones_like(masked_positions), masked_positions)
+        # Dealing with duplicate positions
+        updates_mask = F.np.minimum(updates_mask, 1)
         labels = updates_mask * F.np.not_equal(fake_data, original_data)
         return corrupted_tokens, fake_data, labels
 

src/gluonnlp/op.py

@@ -100,7 +100,7 @@ def updated_vectors_by_position(F, base, data, positions):
     """
     Update each batch with the given positions. Considered as a reversed process of
     "select_vectors_by_position", this is an advanced operator of add_vectors_by_position
-    that updates the results instead of add and avoids duplicate positions.
+    that updates the results instead of adding.
     Once advanced indexing can be hybridized, we can revise the implementation.
 
     updates[i, positions[i, j], :] = data[i, j, :]
@@ -127,22 +127,16 @@ def updated_vectors_by_position(F, base, data, positions):
     out
         The updated result.
         Shape (batch_size, seq_length)
-    updates_mask
-        The state of the updated  for the whole sequence
-        1 -> updated, 0 -> not updated.
-        Shape (batch_size, seq_length)
     """
-    # TODO(zheyuye), update when npx.index_update implemented
-    updates = add_vectors_by_position(F, F.np.zeros_like(base), data, positions)
-    updates_mask = add_vectors_by_position(F, F.np.zeros_like(base),
-            F.np.ones_like(positions), positions)
-    updates = (updates / F.np.maximum(1, updates_mask)).astype(np.int32)
-
-    out = F.np.where(updates, updates, base)
-    updates_mask = F.np.minimum(updates_mask, 1)
-
-    return out, updates_mask
+    positions = positions.astype(np.int32)
+    # batch_idx.shape = (batch_size, 1) as [[0], [1], [2], ...]
+    batch_idx = F.np.expand_dims(F.npx.arange_like(positions, axis=0),
+                                 axis=1).astype(np.int32)
+    batch_idx = batch_idx + F.np.zeros_like(positions)
+    indices = F.np.stack([batch_idx.reshape(-1), positions.reshape(-1)])
 
+    out = F.npx.index_update(base, indices, data.reshape(-1))
+    return out
 
 @use_np
 def gumbel_softmax(F, logits, temperature: float = 1.0, eps: float = 1E-10,