Refactoring TODOs

tfmiss/keras/layers/__init__.py

@@ -9,9 +9,9 @@
 from tfmiss.keras.layers.softmax import NoiseContrastiveEstimation
 from tfmiss.keras.layers.softmax import SampledSofmax
 from tfmiss.keras.layers.todense import ToDense
-from tfmiss.keras.layers.wordvec import BpeEmbedding
 from tfmiss.keras.layers.wordvec import CnnEmbedding
 from tfmiss.keras.layers.wordvec import NgramEmbedding
 from tfmiss.keras.layers.wordvec import WordEmbedding
+from tfmiss.keras.layers.wordvec import WPieceEmbedding
 from tfmiss.keras.layers.wrappers import MapFlat
 from tfmiss.keras.layers.wrappers import WithRagged

tfmiss/keras/layers/softmax.py

@@ -246,8 +246,7 @@ def _train_loss(self, inputs, targets, mask):
                 tail_mask, self._cutoff[0] + i, root_targets
             )
 
-            tail_inputs = inputs[tail_mask]
-            tail_logits = tail_inputs  # TODO: in one line?
+            tail_logits = inputs[tail_mask]
             for t in self.tails[i]:
                 tail_logits = t(tail_logits)
             tail_logits = tf.cast(tail_logits, "float32")
@@ -299,8 +298,7 @@ def _train_probs_loss(self, inputs, targets, mask):
             )
             parent_logprobs = root_logprobs[i + 1]
 
-            tail_inputs = inputs[tail_mask]
-            tail_logits = tail_inputs
+            tail_logits = inputs[tail_mask]
             for t in self.tails[i]:
                 tail_logits = t(tail_logits)
             tail_logits = tf.cast(tail_logits, "float32")
@@ -342,9 +340,7 @@ def _train_probs_loss(self, inputs, targets, mask):
         )
         full_loss.insert(0, root_loss)
 
-        full_loss = (
-            full_loss if mask is None else full_loss[mask]
-        )  # TODO: error?
+        full_loss = full_loss if mask is None else full_loss[mask]
         full_loss = tf.reduce_mean(full_loss)
 
         full_logprobs = tf.concat(full_logprobs, axis=-1)

tfmiss/keras/layers/softmax_test.py

@@ -153,7 +153,7 @@ def test_loss_and_output_2d_sum(self):
 
         self.assertGreater(eval_loss, train_loss)
 
-    # TODO: https://github.com/keras-team/keras/issues/19646
+    # TODO: https://github.com/keras-team/keras/issues/18414
     # def test_loss_mask_3d(self):
     #     inputs = tf.ragged.constant([
     #         [[1., 2.], [2., 3.], [2., 5.]],
@@ -264,7 +264,7 @@ def test_model(self):
         )
         self.assertAllClose(np.ones_like(predictsum), predictsum)
 
-    # TODO: https://github.com/keras-team/keras/issues/19646
+    # TODO: https://github.com/keras-team/keras/issues/18414
     # def test_ragged_input(self):
     #     layer = AdaptiveSoftmax(
     #     units=16, cutoff=[1], return_probs=True, factor=2)
@@ -429,7 +429,7 @@ def test_loss_and_output_2d(self):
 
         self.assertGreater(eval_loss, train_loss)
 
-    # TODO: https://github.com/keras-team/keras/issues/19646
+    # TODO: https://github.com/keras-team/keras/issues/18414
     # def test_loss_mask_3d(self):
     #     inputs = tf.ragged.constant([
     #         [[1., 2.], [2., 3.], [2., 5.]],
@@ -507,15 +507,13 @@ def test_model(self):
 
         self.assertGreater(history["loss"][0], history["loss"][-1])
         self.assertGreater(history["val_loss"][0], history["val_loss"][-1])
-        # self.assertGreater(history['val_loss'][0], history['loss'][0])
-        # self.assertGreater(history['val_loss'][-1], history['loss'][-1])
         self.assertEqual(
             [sample_size // 100, seq_length, num_classes],
             list(predictions.shape),
         )
         self.assertAllClose(np.ones_like(predictsum), predictsum)
 
-    # TODO: https://github.com/keras-team/keras/issues/19646
+    # TODO: https://github.com/keras-team/keras/issues/18414
     # def test_with_ragged_input(self):
     #     layer = SampledSofmax(units=16, negatives=8, return_probs=True)
     #     logits_data = tf.ragged.constant([
@@ -728,7 +726,7 @@ def test_model(self):
         )
         self.assertAllClose(np.ones_like(predictsum), predictsum)
 
-    # TODO: https://github.com/keras-team/keras/issues/19646
+    # TODO: https://github.com/keras-team/keras/issues/18414
     # def test_with_ragged_input(self):
     #     layer = NoiseContrastiveEstimation(
     #     units=16, negatives=8, return_probs=True)

tfmiss/keras/layers/todense_test.py

@@ -32,7 +32,7 @@ def test_ragged_input_with_padding(self):
         output = model.predict(input_data)
         self.assertAllClose(output, expected_output)
 
-    # TODO https://github.com/keras-team/keras/issues/19646
+    # TODO https://github.com/keras-team/keras/issues/18414
     # @parameterized.named_parameters(*test_util.generate_combinations_with_testcase_name(
     #     layer=[layers.SimpleRNN, layers.GRU, layers.LSTM]))
     # def test_ragged_input_rnn_layer(self, layer):

tfmiss/keras/layers/wordvec.py

@@ -350,7 +350,7 @@ def get_config(self):
 
 
 @register_keras_serializable(package="Miss")
-class BpeEmbedding(WordEmbedding):  # TODO: rename to WordPiece
+class WPieceEmbedding(WordEmbedding):
     UNK_CHAR = "##[UNK]"
 
     def __init__(
@@ -610,7 +610,6 @@ def _kernel_init(kernel_size):
                 return initializers.RandomNormal(0.0, stddev)
 
             if "relu" == activations.serialize(self.activation):
-                # TODO: check for keras v3
                 return "random_uniform"
 
             return "glorot_uniform"

tfmiss/keras/layers/wordvec_test.py

@@ -6,11 +6,11 @@
 from keras.src import testing
 from keras.src.saving import register_keras_serializable
 
-from tfmiss.keras.layers.wordvec import BpeEmbedding
 from tfmiss.keras.layers.wordvec import CnnEmbedding
 from tfmiss.keras.layers.wordvec import Highway
 from tfmiss.keras.layers.wordvec import NgramEmbedding
 from tfmiss.keras.layers.wordvec import WordEmbedding
+from tfmiss.keras.layers.wordvec import WPieceEmbedding
 
 
 class WordEmbeddingTest(testing.TestCase):
@@ -992,7 +992,7 @@ def test_preprocess(self):
 
 
 @register_keras_serializable(package="Miss")
-class BpeEmbeddingWrap(BpeEmbedding):
+class WPieceEmbeddingWrap(WPieceEmbedding):
     def call(self, inputs, **kwargs):
         dense_shape = tf.unstack(tf.shape(inputs))
 
@@ -1016,14 +1016,14 @@ def call(self, inputs, **kwargs):
         return outputs
 
 
-class BpeEmbeddingTest(testing.TestCase):
+class WPieceEmbeddingTest(testing.TestCase):
     def test_reserved_words(self):
-        layer = BpeEmbedding()
+        layer = WPieceEmbedding()
         self.assertAllEqual(
             layer._reserved_words, [layer.UNK_MARK, layer.UNK_CHAR]
         )
 
-        layer = BpeEmbedding(reserved_words=["~TesT~"])
+        layer = WPieceEmbedding(reserved_words=["~TesT~"])
         self.assertAllEqual(
             layer._reserved_words, [layer.UNK_MARK, layer.UNK_CHAR, "~TesT~"]
         )
@@ -1045,18 +1045,18 @@ def test_merged_vocab(self):
             "d",
             "##g",
         ]
-        layer = BpeEmbedding(vocab)
+        layer = WPieceEmbedding(vocab)
         self.assertAllEqual(
             layer._vocabulary, [layer.UNK_MARK, layer.UNK_CHAR] + vocab
         )
 
-        layer = BpeEmbedding(vocab, reserved_words=["~TesT~"])
+        layer = WPieceEmbedding(vocab, reserved_words=["~TesT~"])
         self.assertAllEqual(
             layer._vocabulary,
             [layer.UNK_MARK, layer.UNK_CHAR, "~TesT~"] + vocab,
         )
 
-        layer = BpeEmbedding(vocab + ["~TesT~"], reserved_words=["~TesT~"])
+        layer = WPieceEmbedding(vocab + ["~TesT~"], reserved_words=["~TesT~"])
         self.assertAllEqual(
             layer._vocabulary,
             [layer.UNK_MARK, layer.UNK_CHAR, "~TesT~"] + vocab,
@@ -1081,7 +1081,7 @@ def test_build_vocab(self):
             ("##g", 1),
         ]
 
-        layer = BpeEmbedding(vocab_size=4)
+        layer = WPieceEmbedding(vocab_size=4)
         self.assertAllEqual(layer.vocab(counts).most_common(), expected)
 
     def test_adapt_1d(self):
@@ -1113,7 +1113,7 @@ def test_adapt_1d(self):
             [b"d", b"##o", b"##g"],
         ]
 
-        layer = BpeEmbedding(vocab)
+        layer = WPieceEmbedding(vocab)
 
         result = layer.adapt(data).numpy()
 
@@ -1155,7 +1155,7 @@ def test_adapt_2d(self):
             ],
         ]
 
-        layer = BpeEmbedding(vocab)
+        layer = WPieceEmbedding(vocab)
 
         result = layer.adapt(data).numpy()
 
@@ -1197,7 +1197,7 @@ def test_adapt_ragged(self):
             ],
         ]
 
-        layer = BpeEmbedding(vocab)
+        layer = WPieceEmbedding(vocab)
 
         result = layer.adapt(tf.ragged.constant(data)).numpy()
 
@@ -1232,7 +1232,7 @@ def test_preprocess(self):
             [14, 3, 15],
         ]
 
-        layer = BpeEmbedding(vocab)
+        layer = WPieceEmbedding(vocab)
 
         result = layer.preprocess(data).numpy()
 
@@ -1260,11 +1260,11 @@ def test_preprocess(self):
     #         "##g",
     #     ]
     #
-    #     inputs = BpeEmbedding(vocab).preprocess(data)
+    #     inputs = WPieceEmbedding(vocab).preprocess(data)
     #     inputs = inputs.to_tensor(-1)
     #
     #     self.run_layer_test(
-    #         BpeEmbeddingWrap,
+    #         WPieceEmbeddingWrap,
     #         init_kwargs={
     #             "vocabulary": vocab,
     #             "output_dim": 12,
@@ -1279,7 +1279,7 @@ def test_preprocess(self):
     #         expected_output_shape=(8, 12),
     #     )
     #     self.run_layer_test(
-    #         BpeEmbedding,
+    #         WPieceEmbedding,
     #         init_kwargs={
     #             "vocabulary": vocab,
     #             "output_dim": 12,
@@ -1317,11 +1317,11 @@ def test_preprocess(self):
     #         "##g",
     #     ]
     #
-    #     inputs = BpeEmbedding(vocab).preprocess(data)
+    #     inputs = WPieceEmbedding(vocab).preprocess(data)
     #     inputs = inputs.to_tensor(-1)
     #
     #     self.run_layer_test(
-    #         BpeEmbeddingWrap,
+    #         WPieceEmbeddingWrap,
     #         init_kwargs={
     #             "vocabulary": vocab,
     #             "output_dim": 12,
@@ -1363,7 +1363,7 @@ def test_preprocess(self):
     #         "##g",
     #     ]
     #
-    #     outputs = BpeEmbedding(vocab, 5, with_prep=True)(data)
+    #     outputs = WPieceEmbedding(vocab, 5, with_prep=True)(data)
     #     self.assertLen(outputs, 2)
     #     self.assertLen(outputs[0], 1)
     #     self.assertLen(outputs[0][0], 3)
@@ -1372,8 +1372,8 @@ def test_preprocess(self):
     #     self.assertLen(outputs[1][1], 3)
     #     self.assertLen(outputs[1][1][2], 5)
     #
-    #     inputs = BpeEmbedding(vocab).preprocess(data)
-    #     outputs = BpeEmbedding(vocab, 5)(inputs)
+    #     inputs = WPieceEmbedding(vocab).preprocess(data)
+    #     outputs = WPieceEmbedding(vocab, 5)(inputs)
     #     self.assertLen(outputs, 2)
     #     self.assertLen(outputs[0], 1)
     #     self.assertLen(outputs[0][0], 3)

tfmiss/keras/layers/wrappers_test.py

@@ -43,7 +43,7 @@ def test_layer(self):
         self.assertLen(outputs.shape, 3)
         self.assertEqual(outputs.shape[-1], 4)
 
-    # TODO https://github.com/keras-team/keras/issues/19646
+    # TODO https://github.com/keras-team/keras/issues/18414
     # def test_model(self):
     #     logits = tf.ragged.constant([
     #         [[1., 2.], [2., 3.], [2., 5.]],

tfmiss/keras/metrics/f1_test.py

@@ -366,16 +366,10 @@ class F1MacroTest(testing.TestCase):
     def test_config(self):
         r_obj = F1Macro(name="my_f1Macro")
         self.assertEqual(r_obj.name, "my_f1Macro")
-        # self.assertEqual(len(r_obj.variables), 4)
-        # self.assertEqual(
-        #     [v.name for v in r_obj.variables],
-        #     ['true_positives', 'false_positives',
-        #     'true_positives', 'false_negatives'])
 
         # Check save and restore config
         r_obj2 = F1Macro.from_config(r_obj.get_config())
         self.assertEqual(r_obj2.name, "my_f1Macro")
-        # self.assertEqual(len(r_obj2.variables), 4)
 
     def test_value_is_idempotent(self):
         r_obj = F1Macro()