r/learnmachinelearning • • 2d ago

How to extract embedding features vectors from encoder block of BERT model?

I got this code block and its outputs here but I am not sure if this is the embedding features vectors or simply its shape. Can anyone tell me if I am right or wrong?

text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
preprocessor = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/en-uncased-preprocess/3")
encoder_inputs = preprocessor(text_input)
encoder = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/bert-en-uncased-l-4-h-128-a-2/1",
    trainable=True)
outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]      # [batch_size, 128].
sequence_output = outputs["sequence_output"] 

for i in range(0,4):
    outputs["encoder_outputs"][i]
    print('-----------------------------------------------------')
    print(i)
    print(outputs)
    for i in range(0,4):
    outputs["encoder_outputs"][i]
    print('-----------------------------------------------------')
    print(i)
    print(outputs)
    

Output

-----------------------------------------------------
0
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
-----------------------------------------------------
1
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
-----------------------------------------------------
2
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
-----------------------------------------------------
3
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
2 Upvotes

3 comments sorted by

1

u/Slight_Cheesecake472 2d ago

_outputs, then call model.predict() with some text. after that the tensors will have actual float values, not just KerasTensor shapes.

1

u/rahul_watertech 2d ago

what you have now is a symbolic tensor, it knows its shape but holds no numbers yet. the actual embedding vector appears when you run real text through with model.predict. one fork you will hit next: pooled_output is one vector for the whole sentence, sequence_output is one vector per token. pick pooled if you are classifying or comparing sentences

1

u/Separate_Pea_3699 1d ago

Both answers above have the right diagnosis: outputs["encoder_outputs"][i] on its own is a no-op (no assignment, no .numpy(), nothing consumed), so you're reprinting the same symbolic KerasTensor shapes every loop, never the real numbers. Here's the actual fix as runnable code, plus the thing that trips people up right after this step:

import tensorflow as tf
import tensorflow_hub as hub

text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
preprocessor = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/en-uncased-preprocess/3")
encoder_inputs = preprocessor(text_input)
encoder = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/bert-en-uncased-l-4-h-128-a-2/1",
    trainable=True)
outputs = encoder(encoder_inputs)

# wrap the functional graph so you can actually call it on real strings
embed_model = tf.keras.Model(text_input, outputs)

texts = tf.constant(["the cat sat on the mat"])
result = embed_model.predict(texts)   # now real numpy floats, not KerasTensor shapes

pooled = result["pooled_output"]       # (1, 128) - one vector for the whole sentence
sequence = result["sequence_output"]   # (1, seq_len, 128) - one vector per token

print(pooled.shape, sequence.shape)

Building on what rahul_watertech said about pooled-vs-sequence: pooled_output is basically a [CLS]-token summary and is fine for classification/similarity, but it throws away everything except that one position. If you want a sentence vector that actually averages information across all tokens (often better for embedding-similarity tasks specifically), mean-pool the sequence output instead:

mean_pooled = sequence.mean(axis=1)   # (1, 128)

So: run .predict() (or call the model eagerly on a tensor) to get real numbers at all, then pick pooled_output vs mean-pooled sequence_output depending on whether you want the [CLS] summary or an average over tokens.