Add sample_weight conversion in case of boolean weights.
import tensorflow as tf
-def _apply_mask(y_true, sample_weight, masked_tokens):
+def _apply_mask(y_true, sample_weight, masked_tokens, dtype):
if sample_weight is None:
- sample_weight = tf.ones_like(y_true, tf.float32)
+ sample_weight = tf.ones_like(y_true, dtype)
+ else:
+ sample_weight = tf.cast(sample_weight, dtype)
for token in masked_tokens:
- mask = tf.cast(tf.not_equal(y_true, token), tf.float32)
+ mask = tf.cast(tf.not_equal(y_true, token), dtype)
sample_weight = sample_weight * mask
return sample_weight
@@ -53,7 +55,8 @@ class NumTokensCounter(tf.keras.metrics.Sum):
super().__init__(name, dtype)
def update_state(self, y_true, y_pred, sample_weight=None):
- sample_weight = _apply_mask(y_true, sample_weight, self._masked_tokens)
+ sample_weight = _apply_mask(y_true, sample_weight, self._masked_tokens,
+ self._dtype)
sample_weight = tf.reshape(sample_weight, [-1])
super().update_state(sample_weight)
@@ -66,7 +69,8 @@ class MaskedCategoricalAccuracy(tf.keras.metrics.SparseCategoricalAccuracy):
super().__init__(name, dtype=dtype)
def update_state(self, y_true, y_pred, sample_weight=None):
- sample_weight = _apply_mask(y_true, sample_weight, self._masked_tokens)
+ sample_weight = _apply_mask(y_true, sample_weight, self._masked_tokens,
+ self._dtype)
num_classes = tf.shape(y_pred)[-1]
y_true = tf.reshape(y_true, [-1])
y_pred = tf.reshape(y_pred, [-1, num_classes])
