How to Implement Named Entity Recognition (NER) with TensorFlow Using NER-TF.ipynb

The NER-TF.ipynb notebook provides a complete, beginner-friendly pipeline for training a bidirectional LSTM model to identify named entities using TensorFlow's Keras API, achieving over 98% token-level accuracy even with minimal training.

The microsoft/AI-For-Beginners repository contains a practical introduction to Natural Language Processing through the lessons/5-NLP/19-NER/NER-TF.ipynb notebook. This tutorial demonstrates how to build a Named Entity Recognition (NER) system from scratch using TensorFlow, covering everything from dataset preprocessing to model inference without requiring pre-trained transformer weights.

Dataset Acquisition and Preparation

The tutorial uses the Annotated Corpus for Named Entity Recognition from Kaggle. The notebook expects ner_dataset.csv in the working directory, loading it with standard pandas operations.

Loading the NER Dataset

In lessons/5-NLP/19-NER/NER-TF.ipynb, the data ingestion starts with pandas.read_csv using unicode escape encoding to handle special characters:

import pandas as pd
df = pd.read_csv('ner_dataset.csv', encoding='unicode-escape')

Building Tag and Word Mappings

The notebook creates bidirectional dictionaries to convert between textual labels and integer IDs. First, it extracts unique entity tags using df.Tag.unique() and builds tag2id and id2tag mappings. Then it constructs a word vocabulary from the "Word" column, lower-casing each token and adding a special <UNK> token for out-of-vocabulary handling:

# Tag dictionaries

tags = df.Tag.unique()
id2tag = dict(enumerate(tags))
tag2id = {v: k for k, v in id2tag.items()}

# Word vocabulary with UNK handling

vocab = set(df['Word'].apply(lambda x: x.lower()))
id2word = {i+1: w for i, w in enumerate(vocab)}
id2word[0] = '<UNK>'
vocab.add('<UNK>')
word2id = {v: k for k, v in id2word.items()}

Data Preprocessing and Vectorization

Before feeding text into the TensorFlow model, the flat CSV structure must be converted into sentence sequences and padded to uniform length.

Sentence Segmentation

The notebook processes the "Sentence #" column to group tokens into lists. It handles NaN separators in the DataFrame to aggregate words (X) and their corresponding tags (Y) into parallel sentence structures:

X, Y = [], []
s, t = [], []
for _, row in df[['Sentence #', 'Word', 'Tag']].iterrows():
    if pd.isna(row['Sentence #']):
        s.append(row['Word'])
        t.append(row['Tag'])
    else:
        if s:
            X.append(s)
            Y.append(t)
        s, t = [row['Word']], [row['Tag']]
X.append(s); Y.append(t)

Padding Sequences

The vectorize() and tagify() helper functions convert word and tag sequences to integer IDs. The notebook then uses keras.preprocessing.sequence.pad_sequences with padding='post' to create uniform-length tensors:

def vectorize(seq): 
    return [word2id[w.lower()] for w in seq]

def tagify(seq): 
    return [tag2id[t] for t in seq]

Xv = list(map(vectorize, X))
Yv = list(map(tagify, Y))

from tensorflow import keras
X_data = keras.preprocessing.sequence.pad_sequences(Xv, padding='post')
Y_data = keras.preprocessing.sequence.pad_sequences(Yv, padding='post')

Model Architecture

The NER implementation uses a stacked bidirectional LSTM architecture for token-level classification, as defined in the model construction cell of NER-TF.ipynb.

Embedding and Bidirectional LSTM Layers

The model combines three key components:

• Embedding layer: Maps word IDs to 300-dimensional dense vectors (Embedding(vocab_size, 300))
• Two Bidirectional LSTM layers: Each with 100 units and return_sequences=True to capture context from both directions
• TimeDistributed Dense layer: Applies softmax classification to each timestep independently across the 17 entity tag classes

maxlen = X_data.shape[1]
vocab_size = len(vocab)
num_tags = len(tags)

model = keras.models.Sequential([
    keras.layers.Embedding(vocab_size, 300, input_length=maxlen),
    keras.layers.Bidirectional(keras.layers.LSTM(100, activation='tanh',
                                                return_sequences=True)),
    keras.layers.Bidirectional(keras.layers.LSTM(100, activation='tanh',
                                                return_sequences=True)),
    keras.layers.TimeDistributed(keras.layers.Dense(num_tags,
                                                    activation='softmax'))
])

The model compiles with sparse_categorical_crossentropy loss and the Adam optimizer:

model.compile(loss='sparse_categorical_crossentropy',
              optimizer='adam', metrics=['acc'])
model.summary()

Training and Inference

The notebook demonstrates both quick training cycles and practical prediction workflows.

Compiling and Fitting the Model

For demonstration purposes, the notebook runs a single epoch on the full dataset using model.fit(X_data, Y_data). Despite minimal training time, this achieves greater than 98% token-level accuracy, illustrating the effectiveness of bidirectional LSTMs for sequence labeling tasks.

Making Predictions on New Text

For inference, the notebook tokenizes input sentences, maps words to IDs (handling unknown tokens via the <UNK> mapping), and pads sequences to match the training length. Predictions use np.argmax to extract the highest-probability tag for each token:

import numpy as np

# Example inference

sent = 'John Smith went to Paris to attend a conference in cancer development institute'
words = sent.lower().split()
v = keras.preprocessing.sequence.pad_sequences(
        [[word2id.get(w, 0) for w in words]], padding='post', maxlen=maxlen)
res = model(v)[0]
pred_ids = np.argmax(res.numpy(), axis=1)

for token, pid in zip(words, pred_ids):
    print(f"{token} -> {id2tag[pid]}")

Summary

• The NER-TF.ipynb notebook in microsoft/AI-For-Beginners provides a complete TensorFlow implementation for Named Entity Recognition without requiring external pre-trained models.
• The pipeline processes the Kaggle NER dataset through pandas, creates vocabulary mappings with <UNK> handling, and pads sequences for uniform input dimensions.
• The model architecture uses 300-dimensional embeddings followed by two stacked bidirectional LSTM layers (100 units each) and a TimeDistributed Dense classifier.
• Training with sparse_categorical_crossentropy and the Adam optimizer achieves >98% accuracy even in a single epoch.
• For production-grade results, the notebook recommends migrating to transformer-based models like BERT, referencing the companion TransformersTF.ipynb tutorial.

Frequently Asked Questions

What dataset does the NER-TF.ipynb notebook use?

The notebook uses the Annotated Corpus for Named Entity Recognition from Kaggle, specifically the ner_dataset.csv file. This dataset contains over 1,000 sentences annotated with BIO (Begin-Inside-Outside) tags for entities like persons, organizations, and locations.

Why does the model use bidirectional LSTM layers?

Bidirectional LSTMs process sequences in both forward and backward directions, allowing each token's representation to incorporate context from both preceding and following words. This is essential for NER because entity boundaries often depend on future context (e.g., determining whether "Paris" is a person or city requires seeing subsequent words).

How does the notebook handle out-of-vocabulary words?

The vocabulary construction explicitly reserves index 0 for a special <UNK> token. When the vectorize() function encounters words not present in the training vocabulary, it maps them to 0, ensuring the model can process unseen tokens during inference without throwing errors.

What accuracy can I expect from this TensorFlow NER model?

According to the notebook's training logs, the bidirectional LSTM model achieves over 98% token-level accuracy after just one training epoch on the full dataset. However, the notebook notes that transformer-based approaches (like BERT) typically deliver superior generalization and state-of-the-art F1 scores for NER tasks.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}