Fine-Tune NER Models

Your off-the-shelf NER model can't tell a product name from a person's name in your domain, and every time you try to fix it with regex you end up with a brittle mess. The pain is real: generic models fail on specialized jargon, inconsistent labels, and your unique entity types. This lesson shows you how to fine-tune a transformer model for named entity recognition (NER) so your model finally understands your data and your entities — not just the ones it saw in Wikipedia.

The problem this lesson solves

Generic NER models are trained on open-domain text like news articles and web pages. They're great at spotting common names, places, and organizations, but they fall apart when your data looks different. Maybe you work with medical records and need to tag "MEDICATION", "DOSAGE", and "ADVERSE_EVENT" — entities a general model has never learned. Or your text is full of internal acronyms, product codes, and brand names that the model mangles or ignores. The result? You spend hours writing post-processing rules, but the rules don't scale and they still miss edge cases.

You need a custom NER model that:

• Recognizes your domain-specific entity types
• Handles your text's vocabulary and style
• Generalizes beyond your training examples

Fine-tuning is the answer. Instead of training from scratch, you take a pretrained transformer model and adapt it to your task with a relatively small amount of labeled data. This is the LLM Finetuning skill that directly addresses the pain of one-size-fits-all NER.

Core concept / mental model

Think of a pretrained transformer as a talented intern who has read millions of books and articles. They know grammar, context, and plenty of general facts, but they've never worked in your industry. Fine-tuning is a short, focused internship in your domain. You show them a few thousand annotated examples, and they quickly pick up your specific entity types and linguistic quirks.

Concretely, NER fine-tuning treats the problem as a token classification task. Each token in a sentence gets a label, like B-PER, I-PER, or O. The model learns to map a sequence of tokens to a sequence of labels. The architecture is a pretrained transformer (like BERT or RoBERTa) with a classification head on top that outputs a probability for each possible tag.

Here's the mental diagram:

Input tokens:  ["John", "lives", "in", "New", "York", "."]

Ground truth:  [B-PER,   O,     O,   B-LOC, I-LOC, O]

The model sees the entire sentence context and predicts a label for each token. It learns that "John" at the start is a person because of surrounding words, and that "New York" is a multi-word location because "New" starts it and "York" continues it.

Key terms to know:

• BIO/IOB tagging: B- (beginning), I- (inside), O (outside). Builds your label vocabulary.
• Token labels: You must align labels with subword tokens produced by the tokenizer.
• Classification head: A dense layer added on top of the transformer's hidden states.

Once you internalize this, fine-tuning becomes a straightforward data problem: prepare aligned labels, train the head and adapt the transformer layers, then evaluate.

How it works step by step

The process to fine-tune a model for NER always follows the same pipeline:

1. Prepare your dataset — Collect text and annotate named entities. You can use tools like Label Studio, Prodigy, or CoNLL-2003-style format.
2. Convert BIO tags to token-level labels — Use a tokenizer that splits words into subwords (e.g., WordPiece). Because a word like "New" might split into two tokens, you need to handle alignment carefully.
3. Tokenize and encode — Apply the tokenizer to each sentence with truncation=True and padding=True. Map word-level labels to token-level labels, assigning -100 to special tokens so they're ignored in loss.
4. Load a pretrained model — Use AutoModelForTokenClassification with your tag-to-id mapping. It automatically adds a classification head.
5. Train — Set up TrainingArguments and use the Hugging Face Trainer class. The model learns to map tokens to labeled sequences.
6. Evaluate — Compute metrics like Precision, Recall, and F1, typically using sequence-level accuracy (like seqeval).
7. Inference — Predict entity spans and convert token-level predictions back to words.

The cause-and-effect is important: proper label alignment directly determines whether the model can learn. If you get alignment wrong, your loss becomes noisy and the model learns nonsense. Most of the difficulty in NER fine-tuning is in this step, not in the training loop itself.

Hands-on walkthrough

Let's build a complete example using the Hugging Face transformers library. We'll use a small synthetic dataset so you can run this end-to-end. We'll fine-tune distilbert-base-uncased for a custom entity: PRODUCT.

First, install the required packages:

pip install transformers datasets seqeval torch

Step 1: Prepare data

We'll use a minimal dataset with two sentences. In practice you want hundreds or thousands, but this shows the pipeline.

# data.py
# Format: list of (tokens, tags) where tags are in BIO format.

dataset = [
    {"tokens": ["Python", "is", "a", "programming", "language"], "tags": ["B-PRODUCT", "O", "O", "O", "O"]},
    {"tokens": ["I", "love", "PyTorch", "for", "deep", "learning"], "tags": ["O", "O", "B-PRODUCT", "O", "O", "O"]},
]

For real projects, load your annotated data from a JSON, CSV, or CoNLL file.

Step 2: Tokenize and align labels

Now we'll tokenize with a tokenizer that splits into subwords. We need to handle the fact that "PyTorch" might split into ["Py", "##Torch"], so the label B-PRODUCT must be assigned to the first subword token and I-PRODUCT to subsequent subword tokens, or just repeat the same label for continuity.

from transformers import AutoTokenizer

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

label_to_id = {"O": 0, "B-PRODUCT": 1, "I-PRODUCT": 2}
id_to_label = {v: k for k, v in label_to_id.items()}

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True,
        padding="max_length",
        max_length=128,
    )
    labels = []
    for i, label in enumerate(examples["tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)  # special tokens ignored
            elif word_idx != previous_word_idx:
                label_ids.append(label_to_id[label[word_idx]])
            else:
                # subword continuation: use I- tag for entity, else -100
                label_ids.append(label_to_id[label[word_idx]] if label[word_idx] != "O" else -100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_dataset = [tokenize_and_align_labels(item) for item in dataset]

Pro tip: Using -100 for tokens you don't want to contribute to loss (like special tokens) is a critical Hugging Face convention. The cross-entropy loss ignores labels with value -100.

Step 3: Train the model

Now we build the model and train it. We'll convert our dataset to a Dataset object for convenience.

from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer
from datasets import Dataset

# Convert to Hugging Face Dataset
hf_dataset = Dataset.from_list(tokenized_dataset)

# Load model with num_labels
model = AutoModelForTokenClassification.from_pretrained(
    "distilbert-base-uncased",
    num_labels=len(label_to_id),
    id2label=id_to_label,
    label2id=label_to_id,
)

# Training arguments
args = TrainingArguments(
    output_dir="./ner-finetuned",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    evaluation_strategy="no",
    save_strategy="epoch",
)

# Trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=hf_dataset,
    tokenizer=tokenizer,
)

trainer.train()

# Save the model and tokenizer
trainer.save_model("./ner-finetuned")
tokenizer.save_pretrained("./ner-finetuned")

When you run this, you'll see training progress. For our tiny dataset, the loss will drop quickly.

Step 4: Inference on new text

Now let's use the fine-tuned model to predict entities in new text.

from transformers import pipeline

# Load the fine-tuned model with pipeline
ner_pipeline = pipeline("ner", model="./ner-finetuned", tokenizer="./ner-finetuned")

text = "I use TensorFlow for computer vision projects."
predictions = ner_pipeline(text)

for pred in predictions:
    print(f"Entity: {pred['word']}, Label: {pred['entity']}, Score: {pred['score']:.2f}")

Expected output (may vary due to randomness):

Entity: tensorflow, Label: B-PRODUCT, Score: 0.98

Note: The pipeline groups subword tokens by default. If you need exact spans in the original text, use the grouped_entities=True parameter or process the offsets manually.

Compare options / when to choose what

When fine-tuning a model for NER, you have several choices. Here's a comparison to guide you:

When to choose fine-tuning over prompting:

• You need consistent, fast inference at scale (e.g., in your data pipeline).
• You have sensitive data you can't send to external APIs.
• You want full control over the model and it's your product's core feature.

If you have fewer than ~100 examples, consider a labeling sprint or use a general LLM with entity extraction. If you have thousands, fine-tuning is the right call.

Variations to consider:

• BIO vs BIOES schemes — BIOES adds E (end) and S (single) labels, which can improve span detection.
• Domain-specific base models — Instead of bert-base, try biobert for biomedical NER or legal-bert for legal text.
• Span-based models — Instead of token classification, use end-to-end span extraction models for nested or overlapping entities.

Troubleshooting & edge cases

Even with a solid pipeline, you'll hit common issues. Here are fixes:

1. Label misalignment after tokenization

Error or symptom: The model loss is NaN or predictions are garbage.

Fix: Double-check your tokenization function. Use tokenizer.word_ids() to map tokens to original words and assign labels correctly. Print a few tokenized examples with their labels to verify alignment.

2. Model only predicts O (outside) for everything

This often means the training data is too imbalanced (too many O tokens). Fix by: - Using class weights to penalize O more. - Ensuring the training set has a reasonable proportion of entity tokens. - Training for more epochs, but watch for overfitting.

3. GPU memory errors during training

Symptom: CUDA out of memory. Fix by: - Reducing per_device_train_batch_size (e.g., from 8 to 4 or 2). - Using gradient_accumulation_steps to keep the effective batch size. - If still stuck, try LoRA or use a smaller base model.

4. Tokenizer errors on special words

If your text contains characters the tokenizer doesn't know, it may split into many subword tokens. That's normal, but make sure your label alignment handles it gracefully. Use is_split_into_words=True when your input is already a list of tokens, and use word_ids to align.

5. Evaluation metrics are too low

Check if your seqeval metric is computing correctly. Make sure your predictions are converted to label strings before passing to seqeval, not integer IDs. Also use grouped_entities=True when extracting spans for evaluation.

Pro tip: Always evaluate on a holdout set you never train on. If you use the same data for evaluation, your metrics will mislead you.

What you learned & what's next

You've learned the core idea behind fine-tuning a model for NER: it's a token classification task where you adapt a pretrained transformer to your domain using annotated data. You completed a hands-on exercise where you tokenized and aligned labels, trained a model with the Hugging Face Trainer, and ran inference on new text. You also compared fine-tuning to other approaches and learned to troubleshoot common pitfalls.

Key takeaways:

• NER fine-tuning is token classification: each token gets a BIO label.
• Proper label alignment with subword tokens is the #1 success factor.
• AutoModelForTokenClassification adds a classification head automatically.
• Use -100 to mask special tokens in loss.

What's next:

Now that you can fine-tune for NER, the next lesson in the LLM Finetuning track will likely cover evaluating your system end-to-end with seqeval and handling more complex data. You'll build on the training loop you just wrote to include proper validation and metric tracking. You'll also learn to adapt your fine-tuning workflow to other sequence tasks classification, e.g., intent detection.