Implement compute_output_spec() for tokenizers with vocabulary. (#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
 
@@ -617,6 +618,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
+ )