Introduction

Many TensorFlow 2 BERT problems are not caused by the encoder itself. They come from preprocessing: the model expects raw text in one place, token IDs in another, or a dictionary of tensors with exact keys and shapes. If the preprocessor, dataset, and encoder are not aligned, the result is usually a confusing shape or type error.

Understand What the Preprocessor Expects

A BERT preprocessor normally takes raw string tensors and returns a dictionary containing tokenized inputs such as input_word_ids, input_mask, and input_type_ids. That means your Keras model input should usually be tf.string, not already-tokenized integers.

1import tensorflow as tf
2import tensorflow_hub as hub
3
4preprocess = hub.KerasLayer(
5    "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3",
6    name="preprocess"
7)
8
9text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name="text")
10encoder_inputs = preprocess(text_input)
11
12print(type(encoder_inputs))

The key idea is that the preprocessor owns tokenization. If you pass the wrong dtype or shape into that layer, the rest of the model never gets a valid input structure.

Match the Preprocessor and Encoder

Another common issue is mixing components from different model families. The preprocessor and encoder need to agree on vocabulary, casing, tokenization rules, and expected sequence structure.

1import tensorflow as tf
2import tensorflow_hub as hub
3
4preprocess = hub.KerasLayer(
5    "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3",
6    name="preprocess"
7)
8encoder = hub.KerasLayer(
9    "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/2",
10    trainable=True,
11    name="encoder"
12)
13
14text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name="text")
15encoder_inputs = preprocess(text_input)
16outputs = encoder(encoder_inputs)
17pooled_output = outputs["pooled_output"]
18logits = tf.keras.layers.Dense(1, activation="sigmoid")(pooled_output)
19
20model = tf.keras.Model(text_input, logits)
21model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

When those layers are compatible, the model graph is much easier to reason about and the error messages get far less mysterious.

Keep the Dataset Simple

Your tf.data.Dataset should usually emit raw strings and labels. Let the model handle preprocessing internally unless you have a strong reason to move tokenization into the input pipeline.

1texts = tf.constant([
2    "this movie was great",
3    "the plot was boring",
4    "excellent acting",
5    "not worth the time"
6])
7labels = tf.constant([1, 0, 1, 0], dtype=tf.float32)
8
9train_ds = tf.data.Dataset.from_tensor_slices((texts, labels)).batch(2)
10model.fit(train_ds, epochs=1)

This pattern avoids many pipeline bugs. You are training on text, the model accepts text, and the preprocessor transforms it inside the graph.

Inspect the Preprocessor Output Directly

When the error message is unclear, run one batch through the preprocessor by itself and inspect the returned structure:

1sample = tf.constant(["bert preprocessing check"])
2encoded = preprocess(sample)
3
4for key, value in encoded.items():
5    print(key, value.shape, value.dtype)

This quickly tells you whether the layer is returning the keys and tensor ranks that your encoder expects.

Debug Shape and Type Errors Systematically

If the preprocessor fails, inspect three things first:

• Input dtype, which should usually be tf.string
• Input shape, which should usually be a scalar string per example
• Output keys from the preprocessor, which must match what the encoder expects

A fast debugging step is to run a small batch through the preprocessor by itself and inspect the returned dictionary. That tells you whether the problem is in the dataset, the preprocessing layer, or the encoder wiring.

Common Pitfalls

• Passing integer token IDs into a preprocessor layer that expects raw text.
• Mixing a cased encoder with an uncased preprocessor or otherwise combining incompatible components.
• Building a tf.data pipeline that emits shapes different from what the Keras input layer expects.
• Moving preprocessing into the dataset too early and making the training graph harder to debug.

Summary

• Most BERT preprocessing issues in TF2 come from mismatched dtypes, shapes, or model components.
• A preprocessor layer usually expects raw string input and returns a dictionary of tensors.
• Keep the preprocessor and encoder from the same family so tokenization assumptions stay aligned.
• Let the Keras model own preprocessing unless you have a strong reason to push it into tf.data.
• Debug by checking dtype, shape, and output keys before changing the encoder.