Fix TensorFlow BERT Text Classification on Colab/Kaggle: Resolve KerasLayer and tensorflow-text Conflicts

Reproduce TensorFlow's BERT text classification without errors: fix KerasLayer symbolic tensor issues by pinning TensorFlow, tf_keras, tensorflow-text

Reproducing TensorFlow’s BERT text classification tutorial sounds straightforward until the environment bites back. Running the official code on recent Colab/Kaggle runtimes led to a symbolic tensor error inside a KerasLayer and a wall of pip resolver conflicts around TensorFlow, tf-keras, and tensorflow-text. Below is a minimal, reproducible walkthrough of the issue and the setup that made it work reliably.

What goes wrong

The environment was prepared with TensorFlow 2.14.* and matching ecosystem packages, following the tutorial’s spirit of using a “similar” version. Then the BERT encoder and preprocessing layers from TensorFlow Hub were wired together as shown in the tutorial. On model build, execution failed with a ValueError pointing at the preprocessing KerasLayer call and highlighting that a KerasTensor is symbolic and cannot be converted to a NumPy array. At the same time, pip reported dependency conflicts across multiple installed packages and warned that previously imported packages required a runtime restart.

Problem example

The following snippet mirrors the setup that triggered the error. It installs the packages the usual way and then constructs the classifier model with a Hub-based BERT encoder and preprocessing.

!pip install --upgrade pip
!pip install -U "tensorflow-text==2.14.*"
!pip install -U "tf-models-official==2.14.*"
import tensorflow as tf
from tensorflow.keras.datasets import imdb
import tensorflow_hub as thub
import tensorflow_text as ttext
from tensorflow.keras.callbacks import EarlyStopping
hub_encoder_url = "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/1"
hub_prep_url = "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"
def create_text_clf():
    txt_in = tf.keras.layers.Input(shape=(), dtype=tf.string, name="text_in")
    prep_layer = thub.KerasLayer(hub_prep_url, name="bert_prep", trainable=False)
    enc_layer = thub.KerasLayer(hub_encoder_url, name="bert_enc", trainable=True)
    packed = prep_layer(txt_in)
    enc_outputs = enc_layer(packed)
    pooled_vec = enc_outputs["pooled_output"]
    x = tf.keras.layers.Dropout(0.1)(pooled_vec)
    x = tf.keras.layers.Dense(1, activation=None, name="logit_head")(x)
    return tf.keras.Model(txt_in, x)
model_cls = create_text_clf()

In this setup, the call to the preprocessing layer raises the ValueError, while pip outputs incompatibility notices involving tensorflow, tf-keras, tensorflow-text and other packages. The tutorial also suggests using an earlier line of packages, for example:

pip install -U "tensorflow-text==2.13.*"
pip install "tf-models-official==2.13.*"

However, attempting to install tensorflow-text==2.13.* resulted in an error that no matching version could be found.

Why it happens here

The observed failure coincides with dependency conflicts reported by pip and with packages already loaded in the runtime. The error surfaces when the KerasLayer for preprocessing is invoked, and the environment simultaneously shows mismatched versions across TensorFlow, tf-keras, and tensorflow-text. There is also a note that simply picking a “similar version” is insufficient; exact versions requested by dependencies matter. Finally, the runtime explicitly warns that previously imported packages require a restart to take effect.

The setup that works

What unblocked the model build was installing pinned versions of TensorFlow, tf_keras, and tensorflow-text, and doing so without pulling additional dependencies. With those exact versions, the model constructed and ran. The working command is as follows:

!pip install --no-deps tensorflow==2.14.1 tf_keras==2.14.1 tensorflow-text==2.14.0

After installation, restarting the runtime ensures the newly installed packages are actually used, as prompted by the environment warning.

Running the model after fixing the environment

With the versions pinned as above and the runtime restarted, the same model definition builds cleanly. Here is the full classifier model again, unchanged in behavior:

import tensorflow as tf
import tensorflow_hub as thub
import tensorflow_text as ttext
hub_encoder_url = "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/1"
hub_prep_url = "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"
def create_text_clf():
    txt_in = tf.keras.layers.Input(shape=(), dtype=tf.string, name="text_in")
    prep_layer = thub.KerasLayer(hub_prep_url, name="bert_prep", trainable=False)
    enc_layer = thub.KerasLayer(hub_encoder_url, name="bert_enc", trainable=True)
    packed = prep_layer(txt_in)
    enc_outputs = enc_layer(packed)
    pooled_vec = enc_outputs["pooled_output"]
    x = tf.keras.layers.Dropout(0.1)(pooled_vec)
    x = tf.keras.layers.Dense(1, activation=None, name="logit_head")(x)
    return tf.keras.Model(txt_in, x)
model_cls = create_text_clf()

In practice, it was also noted that incompatibilities seemed to originate from libraries pulled alongside tensorflow-text, which is consistent with the decision to install these packages with --no-deps. Additionally, the code was reported to run successfully with TensorFlow==2.14.1.

Why this matters

Training and serving NLP models across Colab or Kaggle hinges on predictable environments. Even minor drifts between tensorflow, tf-keras, and tensorflow-text can break symbolic graph execution or trigger resolver conflicts that are not automatically corrected. Relying on approximate versions often leads to mismatches, and skipping the runtime restart leaves the old stack in memory, hiding the actual installed state.

Takeaways

Pin exact versions for TensorFlow, tf_keras, and tensorflow-text instead of using loose “similar” specs. When installing in an existing notebook session, follow the environment’s prompt and restart the runtime so that the new packages are loaded. If the resolver keeps pulling in conflicting transitive dependencies, install the core trio with --no-deps as shown above and keep them aligned. When earlier package lines are recommended but unavailable, expect installation to fail rather than silently degrade.

With the above environment discipline, the BERT preprocessing and encoder KerasLayers from TensorFlow Hub work as intended, and the classifier builds and runs without tripping over symbolic tensor errors.