candle
candle copied to clipboard
feat: Add BertForTokenClassification for Named Entity Recognition in Rust
Summary
Implements BertForTokenClassification for pure-Rust Named Entity Recognition (NER) in Candle, following the
existing DeBERTa V2 pattern.
This enables Rust applications to perform token classification tasks (NER, POS tagging, etc.) without Python dependencies or heavyweight C++ libraries like libtorch.
Changes
- candle-transformers/src/models/bert.rs: Added
BertForTokenClassificationstruct with dropout and linear classifier layers (~35 lines) - candle-examples/examples/bert-ner/: Complete working example demonstrating NER inference with
dslim/bert-base-NERfrom HuggingFace Hub (~300 lines)
Implementation Details
- Follows the existing DeBERTa V2 NER architecture pattern
- Uses
VarBuilderfor loading model weights from safetensors - Full pipeline: tokenization → forward pass → IOB2 entity extraction
- Supports both HuggingFace Hub and local model loading
Testing
Tested with dslim/bert-base-NER model - successfully extracts named entities (PER, LOC, ORG, MISC) with
confidence scores and character offsets.
Example output:
NERItem { entity: "B-PER", word: "John", score: 0.998, start: 0, end: 4 }
NERItem { entity: "I-PER", word: "Smith", score: 0.997, start: 5, end: 10 }
NERItem { entity: "B-LOC", word: "New", score: 0.995, start: 14, end: 17 }
NERItem { entity: "I-LOC", word: "York", score: 0.994, start: 18, end: 22 }
Use Case
Enables pure-Rust NER for applications requiring lightweight ML inference without Python/libtorch runtime
dependencies. Perfect for:
- Embedded systems and edge devices
- WASM targets (browser/serverless)
- Production deployments where binary size and dependency management are critical
- Cross-platform applications without complex build toolchains
Comparison to Alternatives
- rust-bert: Requires libtorch (~2GB) with complex C++ build dependencies
- This implementation: Pure Rust, ~400MB model download, zero non-Rust dependencies
This fills a significant gap in the Rust ML ecosystem by providing production-ready NER without the complexity of
PyTorch bindings.