Repository: keras-team/keras-nlp

Issue:
BertPreprocessor has confusing behavior on lists of strings
Now `BertPreprocessor` has confusing UX on batch input, mainly because we combine strings together (like NSP). We should be clear on our user journey on the following input types:
1) a string: "This is a string"
2) a list of string: ["holy", "geez"]
3) a list of list of strings [["holy"], ["geez"]]
4) a 1D tensor: tf.Tensor(b"This is a string")
5) a 2D tensor of shape batch * 1: tf.Tensor([["holy"], ["geez"]])
6) a 2D tensor of shape batch * 2: tf.Tensor([["holy", "geez"], ["holy", "geez"]])

One question we need to resolve is - do we default to string combination, or default to treat input as batch data? 

Discussion:

Patch:
diff --git a/keras_nlp/models/bert/bert_preprocessing.py b/keras_nlp/models/bert/bert_preprocessing.py
index b836dea1..0577c911 100644
--- a/keras_nlp/models/bert/bert_preprocessing.py
+++ b/keras_nlp/models/bert/bert_preprocessing.py
@@ -22,6 +22,9 @@ from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.models.bert.bert_presets import backbone_presets
 from keras_nlp.models.bert.bert_presets import classifier_presets
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
+from keras_nlp.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -182,16 +185,17 @@ class BertPreprocessor(keras.layers.Layer):
     `keras.Model.fit`.
 
     The call method of this layer accepts three arguments, `x`, `y`, and
-    `sample_weights`. `x` should be either a (possibly batched) string tensor,
-    or a tuple of (possibly batched) string tensors. `y` and `sample_weights`
-    are both optional, can have any format, and will be passed through
-    unaltered.
+    `sample_weight`. `x` can be a python string or tensor representing a single
+    segment, a list of python strings representing a batch of single segments,
+    or a list of tensors representing multiple segments to be packed together.
+    `y` and `sample_weight` are both optional, can have any format, and will be
+    passed through unaltered.
 
     Special care should be taken when using `tf.data` to map over an unlabeled
-    tuple of string segments. `tf.data` will unpack this tuple directly into the
-    call arguments of this layer, rather than forward all argument to `x`. To
-    handle this case, it is recommended to  explicitly call the layer, e.g.
-    `ds.map(lambda seg1, seg2: preprocessor(x=(seg1, seg2)))`.
+    tuple of string segments. `tf.data.Dataset.map` will unpack this tuple
+    directly into the call arguments of this layer, rather than forward all
+    argument to `x`. To handle this case, it is recommended to  explicitly call
+    the layer, e.g. `ds.map(lambda seg1, seg2: preprocessor(x=(seg1, seg2)))`.
 
     Args:
         tokenizer: A `keras_nlp.models.BertTokenizer` instance.
@@ -217,22 +221,43 @@ class BertPreprocessor(keras.layers.Layer):
     tokenizer = keras_nlp.models.BertTokenizer(vocabulary=vocab)
     preprocessor = keras_nlp.models.BertPreprocessor(tokenizer)
 
-    # Tokenize and pack a single sentence directly.
+    # Tokenize and pack a single sentence.
+    sentence = tf.constant("The quick brown fox jumped.")
+    preprocessor(sentence)
+    # Same output.
     preprocessor("The quick brown fox jumped.")
 
-    # Tokenize and pack a multiple sentence directly.
-    preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
+    # Tokenize and a batch of single sentences.
+    sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    preprocessor(sentences)
+    # Same output.
+    preprocessor(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+
+    # Tokenize and pack a sentence pair.
+    first_sentence = tf.constant("The quick brown fox jumped.")
+    second_sentence = tf.constant("The fox tripped.")
+    preprocessor((first_sentence, second_sentence))
 
     # Map a dataset to preprocess a single sentence.
-    features = ["The quick brown fox jumped.", "I forgot my homework."]
-    labels = [0, 1]
+    features = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    labels = tf.constant([0, 1])
     ds = tf.data.Dataset.from_tensor_slices((features, labels))
     ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
 
-    # Map a dataset to preprocess multiple sentence pairs.
-    first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
-    second_sentences = ["The fox tripped.", "Oh look, a whale."]
-    labels = [1, 1]
+    # Map a dataset to preprocess sentence pairs.
+    first_sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    second_sentences = tf.constant(
+        ["The fox tripped.", "Oh look, a whale."]
+    )
+    labels = tf.constant([1, 1])
     ds = tf.data.Dataset.from_tensor_slices(
         (
             (first_sentences, second_sentences), labels
@@ -241,8 +266,12 @@ class BertPreprocessor(keras.layers.Layer):
     ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
 
     # Map a dataset to preprocess unlabeled sentence pairs.
-    first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
-    second_sentences = ["The fox tripped.", "Oh look, a whale."]
+    first_sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    second_sentences = tf.constant(
+        ["The fox tripped.", "Oh look, a whale."]
+    )
     ds = tf.data.Dataset.from_tensor_slices((first_sentences, second_sentences))
     # Watch out for tf.data's default unpacking of tuples here!
     # Best to invoke the `preprocessor` directly in this case.
@@ -293,9 +322,7 @@ class BertPreprocessor(keras.layers.Layer):
         return cls(**config)
 
     def call(self, x, y=None, sample_weight=None):
-        if not isinstance(x, (list, tuple)):
-            x = [x]
-
+        x = convert_inputs_to_list_of_tensor_segments(x)
         x = [self.tokenizer(segment) for segment in x]
         token_ids, segment_ids = self.packer(x)
         x = {
diff --git a/keras_nlp/utils/keras_utils.py b/keras_nlp/utils/keras_utils.py
index d9a0b245..c1a98679 100644
--- a/keras_nlp/utils/keras_utils.py
+++ b/keras_nlp/utils/keras_utils.py
@@ -13,8 +13,11 @@
 # limitations under the License.
 
 
+import tensorflow as tf
 from tensorflow import keras
 
+from keras_nlp.utils.tf_utils import is_tensor_type
+
 
 def clone_initializer(initializer):
     """Clones an initializer to ensure a new seed.
@@ -45,3 +48,51 @@ def pack_x_y_sample_weight(x, y=None, sample_weight=None):
         return (x, y)
     else:
         return (x, y, sample_weight)
+
+
+def convert_inputs_to_list_of_tensor_segments(x):
+    """Converts user inputs to a list of a tensor segments.
+
+    For models and layers which accept lists of string tensors to pack together,
+    this method converts user inputs to a uniform format in a way that can be
+    considered canonical for the library.
+
+    We handle the following:
+
+    - A single string will be converted to a tensor and wrapped in a list.
+    - A list of strings will be converted to a tensor and wrapped in a list.
+    - A single tensor will be wrapped in a list.
+    - A list of tensors will be passed through unaltered.
+
+    All other inputs will result in an error. This effectively means that users
+    who would like to pack multiple segments together should convert those
+    segments to tensors before calling the layer. This removes any ambiguity
+    in the input for those cases.
+    """
+    # Check the input type.
+    is_string = isinstance(x, (str, bytes))
+    is_tensor = is_tensor_type(x)
+    is_string_list = (
+        isinstance(x, (list, tuple)) and x and isinstance(x[0], (str, bytes))
+    )
+    is_tensor_list = isinstance(x, (list, tuple)) and x and is_tensor_type(x[0])
+
+    if is_string or is_string_list:
+        # Automatically convert raw strings or string lists to tensors.
+        # Wrap this input as a single (possibly batched) segment.
+        x = [tf.convert_to_tensor(x)]
+    elif is_tensor:
+        # Automatically wrap a single tensor as a single segment.
+        x = [x]
+    elif is_tensor_list:
+        # Pass lists of tensors though unaltered.
+        x = x
+    else:
+        # Error for all other input.
+        raise ValueError(
+            f"Unsupported input for `x`. `x` should be a string, a list of "
+            "strings, or a list of tensors. If passing multiple segments "
+            "which should packed together, please convert your inputs to a "
+            f"list of tensors. Received `x={x}`"
+        )
+    return x
diff --git a/keras_nlp/utils/pipeline_model.py b/keras_nlp/utils/pipeline_model.py
index c3579f08..b5b5a0d1 100644
--- a/keras_nlp/utils/pipeline_model.py
+++ b/keras_nlp/utils/pipeline_model.py
@@ -17,23 +17,11 @@
 import functools
 import math
 
-import numpy as np
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
-
-try:
-    import pandas as pd
-except ImportError:
-    pd = None
-
-
-def _get_tensor_types():
-    if pd is None:
-        return (tf.Tensor, np.ndarray)
-    else:
-        return (tf.Tensor, np.ndarray, pd.Series, pd.DataFrame)
+from keras_nlp.utils.tf_utils import is_tensor_type
 
 
 def _convert_inputs_to_dataset(
@@ -77,8 +65,7 @@ def _train_validation_split(arrays, validation_split):
     """
 
     def _can_split(t):
-        tensor_types = _get_tensor_types()
-        return isinstance(t, tensor_types) or t is None
+        return is_tensor_type(t) or t is None
 
     flat_arrays = tf.nest.flatten(arrays)
     unsplitable = [type(t) for t in flat_arrays if not _can_split(t)]
diff --git a/keras_nlp/utils/tf_utils.py b/keras_nlp/utils/tf_utils.py
index 687b4261..1b6ceab7 100644
--- a/keras_nlp/utils/tf_utils.py
+++ b/keras_nlp/utils/tf_utils.py
@@ -12,8 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import tensorflow as tf
 
+try:
+    import pandas as pd
+except ImportError:
+    pd = None
+
 try:
     import tensorflow_text
 except ImportError:
@@ -67,3 +73,10 @@ def assert_tf_text_installed(symbol_name):
             f"{symbol_name} requires the `tensorflow-text` package. "
             "Please install with `pip install tensorflow-text`."
         )
+
+
+def is_tensor_type(x):
+    if pd is None:
+        return isinstance(x, (tf.Tensor, np.ndarray))
+    else:
+        return isinstance(x, (tf.Tensor, np.ndarray, pd.Series, pd.DataFrame))


Test patch:
diff --git a/keras_nlp/models/bert/bert_preprocessing_test.py b/keras_nlp/models/bert/bert_preprocessing_test.py
index 4f7ea449..20683383 100644
--- a/keras_nlp/models/bert/bert_preprocessing_test.py
+++ b/keras_nlp/models/bert/bert_preprocessing_test.py
@@ -80,25 +80,22 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
         self.vocab += ["the", "quick", "brown", "fox"]
-
-    def test_tokenize(self):
-        input_data = ["THE QUICK BROWN FOX."]
-        preprocessor = BertPreprocessor(
+        self.preprocessor = BertPreprocessor(
             BertTokenizer(vocabulary=self.vocab),
             sequence_length=8,
         )
-        output = preprocessor(input_data)
+
+    def test_tokenize_strings(self):
+        input_data = "THE QUICK BROWN FOX."
+        output = self.preprocessor(input_data)
         self.assertAllEqual(output["token_ids"], [2, 5, 6, 7, 8, 1, 3, 0])
         self.assertAllEqual(output["segment_ids"], [0, 0, 0, 0, 0, 0, 0, 0])
         self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
 
-    def test_tokenize_batch(self):
-        input_data = tf.constant(["THE QUICK BROWN FOX."] * 4)
-        preprocessor = BertPreprocessor(
-            BertTokenizer(vocabulary=self.vocab),
-            sequence_length=8,
-        )
-        output = preprocessor(input_data)
+    def test_tokenize_list_of_strings(self):
+        # We should handle a list of strings as as batch.
+        input_data = ["THE QUICK BROWN FOX."] * 4
+        output = self.preprocessor(input_data)
         self.assertAllEqual(output["token_ids"], [[2, 5, 6, 7, 8, 1, 3, 0]] * 4)
         self.assertAllEqual(
             output["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0]] * 4
@@ -111,11 +108,7 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         x = tf.constant(["THE QUICK BROWN FOX."] * 4)
         y = tf.constant([1] * 4)
         sw = tf.constant([1.0] * 4)
-        preprocessor = BertPreprocessor(
-            BertTokenizer(vocabulary=self.vocab),
-            sequence_length=8,
-        )
-        x_out, y_out, sw_out = preprocessor(x, y, sw)
+        x_out, y_out, sw_out = self.preprocessor(x, y, sw)
         self.assertAllEqual(x_out["token_ids"], [[2, 5, 6, 7, 8, 1, 3, 0]] * 4)
         self.assertAllEqual(
             x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0]] * 4
@@ -131,11 +124,7 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         y = tf.constant([1] * 4)
         sw = tf.constant([1.0] * 4)
         ds = tf.data.Dataset.from_tensor_slices((x, y, sw))
-        preprocessor = BertPreprocessor(
-            BertTokenizer(vocabulary=self.vocab),
-            sequence_length=8,
-        )
-        ds = ds.map(preprocessor)
+        ds = ds.map(self.preprocessor)
         x_out, y_out, sw_out = ds.batch(4).take(1).get_single_element()
         self.assertAllEqual(x_out["token_ids"], [[2, 5, 6, 7, 8, 1, 3, 0]] * 4)
         self.assertAllEqual(
@@ -148,15 +137,9 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(sw_out, sw)
 
     def test_tokenize_multiple_sentences(self):
-        sentence_one = "THE QUICK"
-        sentence_two = "BROWN FOX."
-        preprocessor = BertPreprocessor(
-            BertTokenizer(vocabulary=self.vocab),
-            sequence_length=8,
-        )
-        # The first tuple or list is always interpreted as an enumeration of
-        # separate sequences to concatenate.
-        output = preprocessor((sentence_one, sentence_two))
+        sentence_one = tf.constant("THE QUICK")
+        sentence_two = tf.constant("BROWN FOX.")
+        output = self.preprocessor((sentence_one, sentence_two))
         self.assertAllEqual(output["token_ids"], [2, 5, 6, 3, 7, 8, 1, 3])
         self.assertAllEqual(output["segment_ids"], [0, 0, 0, 0, 1, 1, 1, 1])
         self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1])
@@ -164,13 +147,9 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
     def test_tokenize_multiple_batched_sentences(self):
         sentence_one = tf.constant(["THE QUICK"] * 4)
         sentence_two = tf.constant(["BROWN FOX."] * 4)
-        preprocessor = BertPreprocessor(
-            BertTokenizer(vocabulary=self.vocab),
-            sequence_length=8,
-        )
         # The first tuple or list is always interpreted as an enumeration of
         # separate sequences to concatenate.
-        output = preprocessor((sentence_one, sentence_two))
+        output = self.preprocessor((sentence_one, sentence_two))
         self.assertAllEqual(output["token_ids"], [[2, 5, 6, 3, 7, 8, 1, 3]] * 4)
         self.assertAllEqual(
             output["segment_ids"], [[0, 0, 0, 0, 1, 1, 1, 1]] * 4
@@ -179,17 +158,18 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
             output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 1]] * 4
         )
 
+    def test_errors_for_2d_list_input(self):
+        ambiguous_input = [["one", "two"], ["three", "four"]]
+        with self.assertRaises(ValueError):
+            self.preprocessor(ambiguous_input)
+
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
     )
     def test_saving_model(self, save_format):
         input_data = tf.constant(["THE QUICK BROWN FOX."])
-        preprocessor = BertPreprocessor(
-            BertTokenizer(vocabulary=self.vocab),
-            sequence_length=8,
-        )
         inputs = keras.Input(dtype="string", shape=())
-        outputs = preprocessor(inputs)
+        outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
         path = os.path.join(self.get_temp_dir(), "model")
         model.save(path, save_format=save_format)