Embeddings

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers
import numpy as np

# 1) The Embedding layer — a learnable lookup table
VOCAB = 10_000
EMBED_DIM = 64
SEQ_LEN = 50

emb = layers.Embedding(input_dim=VOCAB, output_dim=EMBED_DIM, mask_zero=True)
# input_dim:  vocab size (max token id + 1)
# output_dim: dimension of each vector
# mask_zero:  token 0 is PAD; downstream layers ignore it

sample = tf.constant([[1, 5, 0, 0], [42, 7, 9, 0]])
out = emb(sample)
print(out.shape)                                # (2, 4, 64)

# 2) Embedding as part of a text classifier
model = keras.Sequential([
    layers.TextVectorization(max_tokens=VOCAB, output_sequence_length=SEQ_LEN),
    layers.Embedding(VOCAB, EMBED_DIM, mask_zero=True),
    layers.GlobalAveragePooling1D(),
    layers.Dense(64, activation='relu'),
    layers.Dense(2),
])

model.layers[0].adapt(train_texts)
model.compile(
    optimizer='adam',
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['sparse_categorical_accuracy'],
)
model.fit(train_texts, train_labels, validation_split=0.1, epochs=5)

# 3) Extract trained embeddings
weights = model.layers[1].get_weights()[0]      # (VOCAB, EMBED_DIM)
vocab = model.layers[0].get_vocabulary()
print(vocab[:5], weights[:5].shape)

# 4) Two-tower model — user + item embeddings for recommendations
NUM_USERS = 5000
NUM_ITEMS = 1000
EMBED = 32

user_in = keras.Input(shape=(1,), dtype=tf.int32, name='user')
item_in = keras.Input(shape=(1,), dtype=tf.int32, name='item')

user_emb = layers.Embedding(NUM_USERS, EMBED, embeddings_regularizer=keras.regularizers.l2(1e-6))(user_in)
item_emb = layers.Embedding(NUM_ITEMS, EMBED, embeddings_regularizer=keras.regularizers.l2(1e-6))(item_in)

user_vec = layers.Flatten()(user_emb)
item_vec = layers.Flatten()(item_emb)

score = layers.Dot(axes=1, normalize=True)([user_vec, item_vec])    # cosine similarity

two_tower = keras.Model(inputs=[user_in, item_in], outputs=score)
two_tower.compile(optimizer='adam', loss='mse')

# Train with (user_id, item_id, rating) triples — predicts rating from embeddings.
# At serve time, ANN-search nearest items for a user vector.

# 5) Pretrained word embeddings — initialise from GloVe / fastText
import pickle
glove = {}                                       # word -> np.ndarray
with open('glove.6B.100d.txt', 'r', encoding='utf-8') as f:
    for line in f:
        parts = line.split(); glove[parts[0]] = np.array(parts[1:], dtype=np.float32)

embedding_matrix = np.zeros((VOCAB, 100), dtype=np.float32)
for i, w in enumerate(vocab):
    if w in glove:
        embedding_matrix[i] = glove[w]

emb_init = layers.Embedding(VOCAB, 100, embeddings_initializer=keras.initializers.Constant(embedding_matrix),
                                                          trainable=False)
# trainable=False to KEEP pretrained vectors fixed; True to fine-tune them.

# 6) Visualise embeddings
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

# Sample 500 words for 2-D projection
top_idx = np.arange(500)
vecs = weights[top_idx]
proj = TSNE(n_components=2, perplexity=30, random_state=42).fit_transform(vecs)

plt.figure(figsize=(12, 12))
plt.scatter(proj[:, 0], proj[:, 1], s=5)
for i in range(0, 100, 5):
    plt.annotate(vocab[i], proj[i])

# TensorBoard Projector — built-in visualisation
import io, os
from tensorboard.plugins import projector

log_dir = 'logs/projector'
os.makedirs(log_dir, exist_ok=True)
with open(os.path.join(log_dir, 'metadata.tsv'), 'w', encoding='utf-8') as f:
    for word in vocab:
        f.write(word + '\n')

cp = tf.train.Checkpoint(embedding=tf.Variable(weights))
cp.save(os.path.join(log_dir, 'embedding.ckpt'))

config = projector.ProjectorConfig()
e = config.embeddings.add()
e.tensor_name = 'embedding/.ATTRIBUTES/VARIABLE_VALUE'
e.metadata_path = 'metadata.tsv'
projector.visualize_embeddings(log_dir, config)
# Launch: tensorboard --logdir logs/projector → Projector tab

# 7) Nearest neighbours for recommendations / search
# Normalise + use FAISS for fast ANN
import faiss
weights_norm = weights / np.linalg.norm(weights, axis=1, keepdims=True)
index = faiss.IndexFlatIP(EMBED_DIM)
index.add(weights_norm.astype('float32'))
query = weights_norm[42:43]                       # vector for token id 42
d, idx = index.search(query, k=10)
print([vocab[i] for i in idx[0]])

# 8) Categorical embeddings vs one-hot
# For categorical features with HIGH cardinality (user id, country, product), an Embedding layer
# trained jointly with the model wins over one-hot:
#   • Smaller model: dim ≈ min(50, sqrt(cardinality))
#   • Captures similarity: 'cat' and 'dog' end up nearby; 'cat' and 'volkswagen' don't
#   • Initialised randomly; trained end-to-end

# Rule of thumb: dim = round(cardinality ** 0.25) for many features; tune by validation loss

# 9) Multi-feature embedding model
# In tabular ML with many categoricals, build an Embedding per feature, concat, then dense layers.
from tensorflow.keras.layers import Concatenate
features = [layers.Input(shape=(1,), dtype=tf.int32, name=f) for f in ['user', 'item', 'category']]
embs = [
    layers.Flatten()(layers.Embedding(NUM_USERS, 16)(features[0])),
    layers.Flatten()(layers.Embedding(NUM_ITEMS, 16)(features[1])),
    layers.Flatten()(layers.Embedding(50, 8)(features[2])),
]
x = Concatenate()(embs)
x = layers.Dense(64, activation='relu')(x)
x = layers.Dense(32, activation='relu')(x)
out = layers.Dense(1, activation='sigmoid')(x)
model = keras.Model(features, out)

# 10) Sentence / document embeddings — average pooling baseline
mean_pool = keras.Sequential([
    layers.TextVectorization(max_tokens=VOCAB, output_sequence_length=SEQ_LEN),
    layers.Embedding(VOCAB, EMBED_DIM, mask_zero=True),
    layers.GlobalAveragePooling1D(),
])

# Or use pre-trained sentence transformers (much better):
# tensorflow_hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')

# 11) Production serving
# • TFLite for mobile inference of small embedding tables
# • Vector DB (FAISS, Annoy, Milvus, Pinecone) for fast similarity search at scale
# • Embed at write-time; precompute when possible; cache by entity id

# 12) Common bugs
# • Embedding input_dim too small → KeyError when token id exceeds it; +1 for PAD
# • mask_zero=True downstream layers don't honour mask → check layer-level mask propagation
# • Trainable=True for pretrained but tiny dataset → catastrophic forgetting; freeze first epochs
# • Cosine similarity without L2 normalisation → distances reflect magnitude, not direction
# • Dim too high for small data → overfits; smaller dim + L2 regularisation
# • Saving model but losing TextVectorization vocab → save the whole Sequential, not just the trained dense head
# • Stop word filtering then expecting good embeddings for them — they're gone from the vocab
# • Hashing categorical features without enough hash buckets → collisions hurt quality

from tensorflow.keras import layers
embed = layers.Embedding(input_dim=10_000, output_dim=64, mask_zero=True)