Implement compute_output_spec() for tokenizers with vocabulary. (kera…

…s-team#1523)

* Implement compute_output_spec() for tokenizers with vocabulary. (restarted from new point in master branch)

* Remove type annotation from compute_output_spec() in tokenizers

keras_nlp/tokenizers/byte_pair_tokenizer.py

@@ -24,6 +24,7 @@
 from typing import Iterable
 from typing import List
 
+import keras
 import regex as re
 import tensorflow as tf
 
@@ -605,6 +606,11 @@ def detokenize(self, inputs):
             outputs = tf.squeeze(outputs, 0)
         return outputs
 
+    def compute_output_spec(self, input_spec):
+        return keras.KerasTensor(
+            input_spec.shape + (self.sequence_length,), dtype=self.compute_dtype
+        )
+
     def _transform_bytes(self, tokens):
         """Map token bytes to unicode using `byte2unicode`."""
         split_bytes = tf.strings.bytes_split(tokens)

keras_nlp/tokenizers/sentence_piece_tokenizer.py

@@ -17,6 +17,7 @@
 import os
 from typing import List
 
+import keras
 import tensorflow as tf
 
 from keras_nlp.api_export import keras_nlp_export
@@ -255,3 +256,8 @@ def detokenize(self, inputs):
         if unbatched:
             outputs = tf.squeeze(outputs, 0)
         return outputs
+
+    def compute_output_spec(self, input_spec):
+        return keras.KerasTensor(
+            input_spec.shape + (self.sequence_length,), dtype=self.compute_dtype
+        )

keras_nlp/tokenizers/word_piece_tokenizer.py

@@ -17,6 +17,7 @@
 from typing import Iterable
 from typing import List
 
+import keras
 import tensorflow as tf
 
 from keras_nlp.api_export import keras_nlp_export
@@ -528,3 +529,8 @@ def detokenize(self, inputs):
         if unbatched:
             outputs = tf.squeeze(outputs, 0)
         return outputs
+
+    def compute_output_spec(self, input_spec):
+        return keras.KerasTensor(
+            input_spec.shape + (self.sequence_length,), dtype=self.compute_dtype
+        )