candle icon indicating copy to clipboard operation
candle copied to clipboard

feat: Add BertForTokenClassification for Named Entity Recognition in Rust

Open MSkill1 opened this issue 1 week ago • 0 comments

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 BertForTokenClassification struct with dropout and linear classifier layers (~35 lines)
  • candle-examples/examples/bert-ner/: Complete working example demonstrating NER inference with dslim/bert-base-NER from HuggingFace Hub (~300 lines)

Implementation Details

  • Follows the existing DeBERTa V2 NER architecture pattern
  • Uses VarBuilder for 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.

MSkill1 avatar Nov 24 '25 19:11 MSkill1