Remove support for TF in whole word masking

Author: rjgleaton

src/transformers/data/data_collator.py

@@ -910,7 +910,6 @@ def numpy_mask_tokens(
             word_ids, no_mask_mask = self._calc_word_ids_and_prob_mask(
                 to_numpy(offset_mapping), to_numpy(special_tokens_mask)
             )
-            no_mask_mask = np.array(no_mask_mask, dtype=bool)
         else:
             no_mask_mask = (
                 special_tokens_mask.astype(bool)
@@ -998,7 +997,7 @@ def _calc_word_ids_and_prob_mask(
         word_ids = np.cumsum(is_new_word, axis=1)
         word_ids[special_tokens_mask] = -1
 
-        prob_mask = (~is_new_word).astype(int)
+        prob_mask = ~is_new_word
 
         return word_ids, prob_mask
 
@@ -1049,8 +1048,6 @@ def tolist(x) -> list[Any]:
 def to_numpy(x) -> np.ndarray[Any]:
     if isinstance(x, np.ndarray):
         return x
-    elif hasattr(x, "numpy"):
-        return x.numpy()
     elif hasattr(x, "detach"):
         return x.detach().cpu().numpy()
     else:

tests/trainer/test_data_collator.py

@@ -30,7 +30,6 @@
     DataCollatorWithFlattening,
     DataCollatorWithPadding,
     default_data_collator,
-    is_tf_available,
     is_torch_available,
     set_seed,
 )
@@ -558,19 +557,6 @@ def test_data_collator_for_whole_word_mask(self):
             self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
             self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
 
-        if is_tf_available():
-            import tensorflow as tf
-
-            # Features can already be tensors
-            features = [
-                tokenizer(" ".join(input_tokens), return_offsets_mapping=True).convert_to_tensors("tf")
-                for _ in range(2)
-            ]
-            data_collator = DataCollatorForWholeWordMask(tokenizer, return_tensors="tf")
-            batch = data_collator(features)
-            self.assertEqual(batch["input_ids"].shape, tf.TensorShape((2, 10)))
-            self.assertEqual(batch["labels"].shape, tf.TensorShape((2, 10)))
-
     def test_data_collator_for_whole_word_mask_with_seed(self):
         tokenizer = BertTokenizerFast(self.vocab_file)