Repository: keras-team/keras-nlp

Issue:
Add compute_output_shape method to WordPieceTokenizer
When we run Pretraining Transformer from Scratch guide with PyTorch and JAX backend, it raises

```

RuntimeError: Exception encountered when calling WordPieceTokenizer.call().

Could not automatically infer the output shape / dtype of 'word_piece_tokenizer_1' (of type WordPieceTokenizer). Either the `WordPieceTokenizer.call()` method is incorrect, or you need to implement the `WordPieceTokenizer.compute_output_spec() / compute_output_shape()` method. Error encountered:

'string'

Arguments received by WordPieceTokenizer.call():
  • args=('<KerasTensor shape=(None,), dtype=string, name=keras_tensor_59>',)
  • kwargs=<class 'inspect._empty'>
```

cc: @mattdangerw 


Discussion:
I tried adding `compute_output_shape` method in wordpiece tokenizer

```
    def compute_output_shape(self, inputs_shape):
        if isinstance(inputs_shape[0], tuple):
            inputs_shape = inputs_shape[0]
        inputs_shape = list(inputs_shape)
        inputs_shape[-1] = self.sequence_length
        return tuple(inputs_shape)
```

It throws following error now:
<details> <summary> Error </summary>
--------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[3], line 477
    475 # Add our tokenization into our final model.
    476 inputs = keras.Input(shape=(), dtype=tf.string)
--> 477 tokens = tokenizer(inputs)
    478 outputs = finetuning_model(tokens)
    479 final_model = keras.Model(inputs, outputs)

File ~/Desktop/keras-nlp/keras_nlp/layers/preprocessing/preprocessing_layer.py:37, in PreprocessingLayer.__call__(self, *args, **kwargs)
     33 def __call__(self, *args, **kwargs):
     34     # Always place on CPU for preprocessing, to avoid expensive back and
     35     # forth copies to GPU before the trainable model.
     36     with tf.device("cpu"):
---> 37         outputs = super().__call__(*args, **kwargs)
     39         # Jax and Torch lack native string and ragged types.
     40         # If we are running on those backends and not running with tf.data
     41         # (we are outside a tf.function), we covert all ragged and string
     42         # tensor to pythonic types.
     43         is_tf_backend = config.backend() == "tensorflow"

File ~/anaconda3/envs/backend/lib/python3.10/site-packages/keras_core/src/utils/traceback_utils.py:123, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    120     filtered_tb = _process_traceback_frames(e.__traceback__)
    121     # To get the full stack trace, call:
    122     # `keras_core.config.disable_traceback_filtering()`
--> 123     raise e.with_traceback(filtered_tb) from None
    124 finally:
    125     del filtered_tb

File ~/anaconda3/envs/backend/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

IndexError: Exception encountered when calling WordPieceTokenizer.call().

tuple index out of range

Arguments received by WordPieceTokenizer.call():
  • args=('<KerasTensor shape=(None,), dtype=string, name=keras_tensor_29>',)
  • kwargs=<class 'inspect._empty'>

</details>

I don't think that `compute_output_shape` looks right. Tokenization will always be adding a new axis (the sequence axis), so what you probably want is something like this.

```
    def compute_output_shape(self, inputs_shape):
        return tuple(inputs_shape) + (self.sequence_length,)
```

Also, at a high level, there are some parts of the guide we should probably just ditch in the new multi-backend work. In particular this section -> https://keras.io/guides/keras_nlp/transformer_pretraining/#save-a-model-that-accepts-raw-text, we may want to just axe entirely. jax and torch have no concept of string tensor, so the section doesn't really apply.
@mattdangerw shall I open a PR adding compute_output_shape to tokenizer.py ?
Hey @mattdangerw in the pretraining guide, https://github.com/keras-team/keras-io/pull/1438

when we call fit on the model created with keras.Model, it throws following for **torch** backend

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-5-9f69c9ccacc7>](https://localhost:8080/#) in <cell line: 379>()
    377 
    378 # Pretrain the model on our wiki text dataset.
--> 379 pretraining_model.fit(
    380     pretrain_ds,
    381     validation_data=pretrain_val_ds,

1 frames
[/usr/local/lib/python3.10/dist-packages/keras_core/src/utils/traceback_utils.py](https://localhost:8080/#) in error_handler(*args, **kwargs)
    121             # To get the full stack trace, call:
    122             # `keras_core.config.disable_traceback_filtering()`
--> 123             raise e.with_traceback(filtered_tb) from None
    124         finally:
    125             del filtered_tb

[/usr/local/lib/python3.10/dist-packages/keras_core/src/backend/torch/trainer.py](https://localhost:8080/#) in _symbolic_build(self, data_batch)
    168                 y_pred = backend.compute_output_spec(self, x)
    169             except:
--> 170                 raise RuntimeError(
    171                     "Unable to automatically build the model. "
    172                     "Please build it yourself before calling "

RuntimeError: Unable to automatically build the model. Please build it yourself before calling fit/evaluate/predict. A model is 'built' when its variables have been created and its `self.built` attribute is True. Usually, calling the model on a batch of data is the right way to build it.
```
@shivance sure!

Resolving patch:
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
--- a/keras_nlp/tokenizers/tokenizer.py
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -123,3 +123,6 @@ def token_to_id(self, token: str) -> int:
 
     def call(self, inputs, *args, training=None, **kwargs):
         return self.tokenize(inputs, *args, **kwargs)
+
+    def compute_output_shape(self, inputs_shape):
+        return tuple(inputs_shape) + (self.sequence_length,)


Test patch: