pyt-crf

A modern linear-chain CRF (Conditional Random Field) layer for PyTorch.

• forward() returns the negative log-likelihood directly (ready for .backward())
• decode() runs batched Viterbi decoding
• forward_and_decode() computes both in one fused sweep (for eval during training), with separate loss/decode masks
• Batch-first (batch, seq_len, num_tags) by default

See docs/explanation.md for how the CRF math maps to the code.

Requirements

• Python >= 3.12
• PyTorch >= 2.10

Usage

import torch
from pyt_crf import CRF

crf = CRF(num_tags=5)
emissions = torch.randn(2, 7, 5)          # (batch, seq_len, num_tags)
tags = torch.randint(5, (2, 7))           # (batch, seq_len)
mask = torch.ones(2, 7, dtype=torch.bool) # True = real token, False = padding

loss = crf(emissions, tags, mask)  # scalar NLL loss
loss.backward()

best_paths = crf.decode(emissions, mask)  # List[List[int]]

# eval during training: loss + decode in one fused sweep,
# with separate masks for the loss (-100 labels) and decoding (attention mask)
loss, best_paths = crf.forward_and_decode(
    emissions, tags, loss_mask=mask, decode_mask=mask
)

Warning

Run the CRF in float32. The forward/Viterbi recursions accumulate logsumexp/max scores over the whole sequence, which is numerically unstable in bf16/fp16. Keep the CRF parameters in fp32 and cast emissions before the layer — crf(emissions.float(), ...) — even when the backbone runs in bf16/autocast.

Examples

• examples/conll2003 — NER fine-tuning on tner/conll2003 with a Gemma3-based backbone + CRF head.

Build (wheel)

uv build

# or, without uv:
pip install build
python -m build

Both commands produce dist/pyt_crf-<version>-py3-none-any.whl (plus the sdist dist/pyt_crf-<version>.tar.gz). The version comes from pyproject.toml.

Install

# from the built wheel
pip install dist/pyt_crf-<version>-py3-none-any.whl

# or editable, for development
pip install -e .

Test

Tests live in tests/ and run with pytest. pyproject.toml already puts src/ on the import path for pytest, so no install is needed beforehand.

conda activate py312
pip install pytest

pytest            # run the whole suite
pytest -v         # per-test names
pytest -k fused   # only tests matching a keyword

# or with uv (pytest comes from the dev dependency group):
uv run pytest

The suite checks forward/decode against brute-force enumeration of all tag paths (including masks with mid-sequence holes), verifies forward_and_decode returns exactly the same loss/paths/gradients as the separate calls, and covers input validation.

Changelog

• 0.1.2 — forward_and_decode (fused loss + Viterbi), pytest suite, docs, publish workflow.
• 0.1.1 — initial release; full list of differences from pytorch-crf.

Acknowledgements

• This package mostly follows implementation of kmkurn/pytorch-crf.