Repository: keras-team/keras-nlp

Issue:
Fix rank 1 outputs for WordPieceTokenizer
Right now, a scalar string inputs to the tokenizer will turn into rank 2 outputs, which we do not want. Scalar inputs should result in rank 1 outputs.

See this colab for a reproduction of the issue:
https://colab.research.google.com/gist/mattdangerw/8e1fdbebadf0ff6ff2dc065b2e00db65/rank1wordpiece.ipynb


Discussion:
This gets a little tricky in how this will work with ragged outputs.

When `sequence_length` is defined, the output for a scalar should be a dense rank 1 tensor with shape `[sequence_length]`.

When `sequence_length=None`, we can't output a ragged tensor for a scalar input, as there is no such thing as a ragged rank 1 tensor. Instead we should output a dense tensor, with static shape of `[None]`. That way the dense tensor could still be batched to a ragged tensor with [tf.data.experimental.dense_to_ragged_batch](https://www.tensorflow.org/api_docs/python/tf/data/experimental/dense_to_ragged_batch).
Hey
I'd like to try out this issue as well. I've raised a PR for the other one and would make changes there as advised
@mattdangerw 
I just wish to confirm a couple of things - 
- For the non ragged tensor case a simple reshape should suffice right?
- For the case with ragged tensors I should eliminate it all together and return a dense tensor so that would mean I add an else statement and handle it [here](https://github.com/keras-team/keras-nlp/blob/da3da0e8afd264eca095da003946d9c9271995e8/keras_nlp/tokenizers/word_piece_tokenizer.py#L270) right? as from what I see the other way would be to go and change `tf_text.FastWordpieceTokenizer` as that is what is essentially tokenizing the sentence

Am I missing something?
When `sequence_length` is set, should be fine to just squeeze the leading dim I think.

When `sequence_length` is not set, yeah don't change the underlying tf text tokenizer. Adding some logic around that if block you pointed to sounds good. You will need to covert to dense in the case where the output should be rank 1, squeeze the extra dimension, and then I think you could do `output.set_shape([None])` to make sure the static shape is None on the dim.

We should actually add a unit test where we make a dataset of scalar strings, map through our tokenizer without setting sequence length, and the use `tf.data.experimental.dense_to_ragged_batch` to batch the tokenized output to a ragged tensor.
Thanks for the review!
I'll get to work on this
Hey @mattdangerw 
I think it's working as intended now it would be great if you could have a look at it [here](https://colab.research.google.com/gist/aflah02/3bd4c0d002df308d9afbe32ce1d58890/fixingrank1wordpiece.ipynb) 
The changes done are [here](https://github.com/aflah02/keras-nlp/blob/5eef3434f0a266fb84ff58f6d8730a3ed99ee5ab/keras_nlp/tokenizers/word_piece_tokenizer.py#L270-L279)
If this seems to be the right approach I'll get on with making unit tests and fixing the docstrings
Looks close! But we can't always just flatten the inputs, we do want to handle batch inputs too. There will need to be some conditioning on the input, so the changes are special cased for. Desired behavior:

 - Rank 0 input (scalar), sequence_length set -> dense rank 1 output, known shape
 - Rank 0 input (scalar), sequence_length not set -> dense rank 1 output, unknown shape
 - Rank 1 input, sequence_length set -> dense rank 2 output
 - Rank 1 input, sequence_length not set -> ragged rank 2 output
 - Rank 2 input, sequence_length set -> dense rank 3 output
 - Rank 2 input, sequence_length not set -> ragged rank 3 output

So this should be about fixing the scalar input problem, not changing the shape for all inputs. Sorry if that was unclear!

You can think of this in terms of workflows too. Users should be able to tokenize before batching, or batch before tokenizing, with a simple mental model for how the shape will change (input rank + 1). 
Thanks!, it's more clear now I'll get to work on this
Hey @mattdangerw 
I'm facing some difficulty in tracking the rank of the input, for instance in this snippet  - 

 ```
tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
vocabulary=["[UNK]", "a", "b", "c"], sequence_length = 6)
dataset = tf.data.Dataset.from_tensor_slices([["a", "b c", "a", "b"], ["a", "b c", "a", "b"]])
dataset = dataset.map(tokenizer)
dataset.take(1).get_single_element()
```

The print lines add by me show - 

```
RANK  Tensor("word_piece_tokenizer_4/Rank:0", shape=(), dtype=int32)
INPUT  Tensor("args_0:0", shape=(4,), dtype=string)
SHAPE  (4,)
```

The rank of the input is 1 however when I try printing the rank of the input inside the tokenize function it's always 
a rank 0 tensor as map acts on individual elements of the tensor. I'm not really sure how do I handle the cases where the input is rank 0 given the function always will calculate rank as 0
dataset.map will not iterate over each individual element, it will iterate over the first dimension of the input. The print statements you are showing look like a rank 1 tensor to me--shape=(4,) rank=1.

I made another colab that shows the rank of the inputs. Try comparing this with your code?
https://colab.sandbox.google.com/gist/mattdangerw/580a52996e527e2d85d8599953f88be5/rank-explanation.ipynb

Resolving patch:
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -256,6 +256,12 @@ def get_config(self) -> Dict[str, Any]:
         return config
 
     def tokenize(self, inputs):
+        # Check if Input is Scalar or Not
+        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+            inputs = tf.convert_to_tensor(inputs)
+        scalar_input = tf.convert_to_tensor(inputs).shape.rank == 0
+        if scalar_input:
+            inputs = tf.expand_dims(inputs, 0)
         # Optionally normalize and split inputs.
         if self._lowercase:
             inputs = tf_text.case_fold_utf8(inputs)
@@ -282,6 +288,11 @@ def tokenize(self, inputs):
             output_shape = tokens.shape.as_list()
             output_shape[-1] = self._sequence_length
             tokens = tokens.to_tensor(shape=output_shape)
+        # Convert to a dense output if input in scalar
+        if scalar_input:
+            tokens = tf.squeeze(tokens, 0)
+            tf.ensure_shape(tokens, shape=[self._sequence_length])
+
         return tokens
 
     def detokenize(self, inputs):


Test patch:
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -163,6 +163,18 @@ def test_functional_model(self):
         model_output = model(input_data)
         self.assertAllEqual(model_output, ["the quick brown fox"])
 
+    def test_batching_ragged_tensors(self):
+        tokenizer = WordPieceTokenizer(
+            vocabulary=["[UNK]", "a", "b", "c", "d", "e", "f"]
+        )
+        dataset = tf.data.Dataset.from_tensor_slices(["a b c", "d e", "a f e"])
+        dataset = dataset.map(tokenizer)
+        dataset = dataset.apply(
+            tf.data.experimental.dense_to_ragged_batch(batch_size=1)
+        )
+        element = dataset.take(1).get_single_element().numpy()
+        self.assertAllEqual(element, [[1, 2, 3]])
+
     def test_from_file(self):
         vocab_path = os.path.join(self.get_temp_dir(), "vocab.txt")
         input_data = ["the quick brown fox."]