TensorFlow
data embedding
machine learning
AI tools
data processing

Feeding data through an embedding wrapper in TensorFlow

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In TensorFlow, an embedding layer does not take raw words or arbitrary floating-point features directly. It expects integer token ids. Feeding data through an embedding wrapper therefore means converting tokens into integer indices, padding sequences to a consistent length, and then passing the resulting integer tensor into tf.keras.layers.Embedding.

What the Embedding Layer Expects

An embedding layer is basically a learnable lookup table. Each integer id points to one row in that table, and the row becomes the dense vector for that token.

That means the input should usually have shape batch by sequence length and dtype integer. For text, a sequence like "red car fast" might become [14, 209, 87]. The embedding layer turns those ids into vectors such as shape sequence length by embedding dimension.

A Minimal Working Example

The following example uses integer-tokenized sequences directly:

python
1import tensorflow as tf
2
3vocab_size = 1000
4sequence_length = 6
5embedding_dim = 16
6
7x_train = tf.constant([
8    [12, 45, 87, 0, 0, 0],
9    [5, 19, 333, 44, 91, 0],
10    [71, 8, 2, 19, 55, 101],
11], dtype=tf.int32)
12
13y_train = tf.constant([0, 1, 1], dtype=tf.int32)
14
15model = tf.keras.Sequential([
16    tf.keras.layers.Embedding(
17        input_dim=vocab_size,
18        output_dim=embedding_dim,
19        input_length=sequence_length,
20        mask_zero=True,
21    ),
22    tf.keras.layers.GlobalAveragePooling1D(),
23    tf.keras.layers.Dense(16, activation="relu"),
24    tf.keras.layers.Dense(2, activation="softmax"),
25])
26
27model.compile(
28    optimizer="adam",
29    loss="sparse_categorical_crossentropy",
30    metrics=["accuracy"],
31)
32
33model.fit(x_train, y_train, epochs=3)

Here, each row in x_train is a padded sequence of integer ids. The embedding layer converts each integer into a learnable dense vector, and the later layers operate on those vectors.

Tokenization Comes Before the Embedding

If your source data is raw text, tokenize it first. In Keras, TextVectorization is a common way to do that inside the model pipeline:

python
1import tensorflow as tf
2
3texts = tf.constant([
4    "red car fast",
5    "blue bike slow",
6    "fast train red",
7])
8labels = tf.constant([1, 0, 1])
9
10vectorizer = tf.keras.layers.TextVectorization(
11    max_tokens=1000,
12    output_mode="int",
13    output_sequence_length=5,
14)
15vectorizer.adapt(texts)
16
17model = tf.keras.Sequential([
18    tf.keras.Input(shape=(1,), dtype=tf.string),
19    vectorizer,
20    tf.keras.layers.Embedding(input_dim=1000, output_dim=16, mask_zero=True),
21    tf.keras.layers.GlobalAveragePooling1D(),
22    tf.keras.layers.Dense(1, activation="sigmoid"),
23])
24
25model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
26model.fit(texts, labels, epochs=3)

The embedding layer never sees raw strings. It only sees the integer ids produced by the vectorizer.

Important Shape and Range Rules

Two rules matter more than anything else:

  • every token id must be in the range from 0 to input_dim - 1
  • the input tensor must contain integers, not one-hot vectors or floating-point features

If an id is outside the valid range, TensorFlow raises an index error. If you pass one-hot encoded vectors into an embedding layer, the data shape is wrong for the lookup operation.

mask_zero=True is useful when 0 is your padding token. It tells compatible downstream layers to ignore padding positions, which prevents short sequences from being dominated by filler values.

Pretrained Embeddings Still Use Integer Ids

The same input rules apply when you initialize the embedding table with pretrained weights. Whether the vectors start randomly or come from a pretrained matrix, the layer still expects integer token ids as input.

The only thing that changes is the initial value of the lookup table, not the format of the incoming data.

Common Pitfalls

The biggest mistake is feeding raw text directly into Embedding without tokenization. Embeddings work on integer ids, not strings.

Another issue is mismatching input_dim with the vocabulary size. If your tokenizer produces ids larger than the embedding table size, training fails.

Developers also forget padding. Batches need consistent sequence length unless you use ragged-tensor aware handling throughout the model.

Summary

  • An embedding layer expects integer token ids, not raw strings or one-hot vectors.
  • Tokenize first, then pad sequences to a consistent length.
  • Set input_dim large enough to cover the full token-id range.
  • Use mask_zero=True when 0 represents padding.
  • Think carefully about shapes, because most embedding errors are input-format mistakes rather than model mistakes.

Course illustration
Course illustration

All Rights Reserved.