TensorFlow Transformers and BERT

Real-World Scenario

A news aggregator needs to classify articles into 6 topics (politics, technology, sports, business, health, entertainment). Fine-tuning DistilBERT on 3,000 labeled articles for 3 epochs achieves 93% accuracy — far beyond what TF-IDF + logistic regression (78%) delivers — and the TF SavedModel integrates seamlessly with the existing TF Serving infrastructure.

Fine-Tuning DistilBERT with TF Backend

import tensorflow as tf
from transformers import (
    TFAutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
)
import numpy as np

# Sample data
TEXTS = [
    "President signs new climate legislation",
    "Tech company releases AI model breakthrough",
    "Home team wins championship in overtime",
    "Central bank raises interest rates again",
    "New vaccine shows promising trial results",
    "Actor wins award at film ceremony",
] * 50   # 300 examples

LABELS    = [0, 1, 2, 3, 4, 5] * 50
N_CLASSES = 6
LABEL_NAMES = ["politics", "technology", "sports", "business", "health", "entertainment"]

MODEL_NAME = "distilbert-base-uncased"
tokenizer  = AutoTokenizer.from_pretrained(MODEL_NAME)

# Tokenize the full dataset
encodings = tokenizer(
    TEXTS,
    truncation=True,
    padding="max_length",
    max_length=64,
    return_tensors="tf",
)

# Create tf.data.Dataset
dataset = tf.data.Dataset.from_tensor_slices((
    dict(encodings),
    tf.constant(LABELS, dtype=tf.int32),
))
dataset = dataset.shuffle(1000, seed=42)

train_size = int(0.8 * len(TEXTS))
train_ds = dataset.take(train_size).batch(16).prefetch(tf.data.AUTOTUNE)
val_ds   = dataset.skip(train_size).batch(32).prefetch(tf.data.AUTOTUNE)

# Load TF variant of DistilBERT with classification head
model = TFAutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=N_CLASSES,
)

# Linear warmup + linear decay schedule
total_steps  = len(train_ds) * 3   # 3 epochs
warmup_steps = total_steps // 10

lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
    initial_learning_rate=3e-5,
    decay_steps=total_steps - warmup_steps,
    end_learning_rate=0.0,
)
# Wrap with warmup
class WarmupSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, post_warmup_schedule, warmup_steps):
        super().__init__()
        self.schedule      = post_warmup_schedule
        self.warmup_steps  = tf.cast(warmup_steps, tf.float32)

    def __call__(self, step):
        step   = tf.cast(step, tf.float32)
        warmup = step / self.warmup_steps
        post   = self.schedule(step - self.warmup_steps)
        return tf.cond(step < self.warmup_steps, lambda: warmup * 3e-5, lambda: post)

optimizer = tf.keras.optimizers.AdamW(
    learning_rate=WarmupSchedule(lr_schedule, warmup_steps),
    weight_decay=0.01,
)

model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=3,
    callbacks=[
        tf.keras.callbacks.EarlyStopping(patience=2, restore_best_weights=True),
    ],
)

# Save as SavedModel for TF Serving
model.save_pretrained("./tf_distilbert_news")
tokenizer.save_pretrained("./tf_distilbert_news")
print("Model saved")

Zero-Shot Classification with Pre-Trained Models

from transformers import pipeline

# No fine-tuning needed — works out of the box
classifier = pipeline(
    "zero-shot-classification",
    model="facebook/bart-large-mnli",   # or distilbart-mnli for speed
)

texts = [
    "Apple announces record quarterly revenue despite supply chain issues",
    "Scientists discover potential treatment for Alzheimer's disease",
    "Municipal elections show record voter turnout in major cities",
]

candidate_labels = ["technology", "health", "politics", "business", "sports"]

for text in texts:
    result = classifier(text, candidate_labels)
    top_label = result["labels"][0]
    top_score = result["scores"][0]
    print(f"Text:  '{text[:60]}...'")
    print(f"Label: {top_label}  ({top_score:.1%})\n")

Sentence Embeddings with TF

import tensorflow as tf
from transformers import TFAutoModel, AutoTokenizer
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model     = TFAutoModel.from_pretrained("distilbert-base-uncased")

def mean_pooling(model_output, attention_mask):
    """Average token embeddings weighted by attention mask."""
    token_embeddings = model_output.last_hidden_state
    mask_expanded    = tf.cast(
        tf.expand_dims(attention_mask, -1), tf.float32
    )
    token_sum = tf.reduce_sum(token_embeddings * mask_expanded, axis=1)
    mask_sum  = tf.clip_by_value(tf.reduce_sum(mask_expanded, axis=1), 1e-9, 1e9)
    return token_sum / mask_sum


def embed_texts(texts: list[str]) -> np.ndarray:
    inputs = tokenizer(
        texts, padding=True, truncation=True,
        max_length=128, return_tensors="tf"
    )
    outputs   = model(**inputs)
    embeddings = mean_pooling(outputs, inputs["attention_mask"])
    # L2 normalize for cosine similarity via dot product
    norms = tf.norm(embeddings, axis=1, keepdims=True)
    return (embeddings / norms).numpy()


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.dot(a, b))


# Semantic similarity
pairs = [
    ("The stock market crashed today", "Equities fell sharply on Wall Street"),
    ("The stock market crashed today", "My cat likes to sleep in the sun"),
]
for s1, s2 in pairs:
    embs  = embed_texts([s1, s2])
    score = cosine_similarity(embs[0], embs[1])
    print(f"Similarity: {score:.3f}")
    print(f"  '{s1}'")
    print(f"  '{s2}'\n")

Custom Token Classification (NER)

from transformers import (
    TFAutoModelForTokenClassification,
    AutoTokenizer,
    DataCollatorForTokenClassification,
)
import tensorflow as tf
import numpy as np

# Named Entity Recognition: tag each token as B-PER, I-PER, B-ORG, O, etc.
LABEL2ID = {"O": 0, "B-PER": 1, "I-PER": 2, "B-ORG": 3, "I-ORG": 4, "B-LOC": 5, "I-LOC": 6}
ID2LABEL = {v: k for k, v in LABEL2ID.items()}
N_LABELS = len(LABEL2ID)

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model     = TFAutoModelForTokenClassification.from_pretrained(
    "distilbert-base-uncased",
    num_labels=N_LABELS,
    id2label=ID2LABEL,
    label2id=LABEL2ID,
)
model.compile(
    optimizer=tf.keras.optimizers.AdamW(learning_rate=5e-5),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)


def predict_entities(text: str) -> list[dict]:
    """Run NER on a single sentence."""
    inputs  = tokenizer(text, return_tensors="tf", return_offsets_mapping=True)
    offsets = inputs.pop("offset_mapping").numpy()[0]

    logits  = model(**inputs).logits[0]
    labels  = tf.argmax(logits, axis=-1).numpy()

    tokens  = tokenizer.convert_ids_to_tokens(inputs["input_ids"].numpy()[0])
    entities = []
    for token, label_id, (start, end) in zip(tokens, labels, offsets):
        if token in ("[CLS]", "[SEP]", "[PAD]"):
            continue
        label = ID2LABEL[label_id]
        if label != "O":
            entities.append({"token": token, "label": label, "span": (start, end)})

    return entities


# Example (labels will be random without fine-tuning — shows the pipeline)
sample_text = "Apple CEO Tim Cook announced new products in Cupertino."
entities = predict_entities(sample_text)
print(f"Text: {sample_text}")
print(f"Entities (random labels, model not fine-tuned): {entities[:5]}")