rust-bert

rust-bert is a Rust-native port of Hugging Face’s Transformers library, giving Rust applications the same task-specific NLP pipelines and pretrained transformer architectures without a Python runtime in the loop. It wraps tch-rs (libtorch bindings) or an ONNX Runtime backend for inference, and rust-tokenizers for multi-threaded, pre-processing-heavy tokenization, so a Rust binary can run BERT, GPT-2, BART, T5, and two dozen other architectures directly on CPU or GPU.

Beyond the base model implementations, the crate ships ready-to-use pipelines for translation, summarization, question answering, zero-shot classification, named entity recognition, part-of-speech tagging, sentence embeddings, masked language modeling, keyword extraction, and multi-turn conversational agents — each callable in a few lines of Rust and backed by downloadable pretrained weights converted from the original PyTorch checkpoints.

What You Get

• Ready-to-use pipelines: translation, summarization, question answering, zero-shot classification, NER, POS tagging, sentence embeddings, masked LM, keyword extraction, and conversational agents, each usable in a handful of lines
• Native Rust implementations of 25+ transformer architectures (BERT, RoBERTa, DistilBERT, GPT/GPT-2/GPT-Neo/GPT-J, BART, T5/LongT5, Marian, MBart, M2M100, NLLB, Electra, ALBERT, XLNet, Reformer, ProphetNet, Longformer, Pegasus, DeBERTa/DeBERTa-v2, FNet, MobileBERT) with task-specific heads (sequence classification, token classification, QA, multiple choice)
• Dual inference backends: tch (libtorch) for GPU-accelerated training-grade inference, or ort (ONNX Runtime) for a lighter, portable CPU-only deployment path
• RemoteResources abstraction that downloads and caches pretrained config, vocabulary, and converted .ot/ONNX weight files automatically, alongside a convert-tensor CLI utility for converting Hugging Face .bin checkpoints
• Multi-threaded tokenization via the companion rust-tokenizers crate, avoiding a Python preprocessing dependency

Common Use Cases

• Embedding sentiment analysis, NER, or zero-shot classification directly inside a Rust backend service without a Python microservice hop
• Running local, offline machine translation or summarization in a CLI tool or desktop app using pretrained MarianMT/T5/BART/NLLB weights
• Building conversational or text-generation features (GPT-2/GPT-Neo/GPT-J-based) inside a Rust application with deterministic, low-overhead deployment
• Deploying NLP inference to CPU-only or resource-constrained environments via the ONNX Runtime backend instead of full libtorch
• Generating sentence embeddings for semantic search or clustering pipelines written entirely in Rust

Under The Hood

Architecture Each transformer family lives in its own module under src/models/ (e.g. models::bert, models::gpt2, models::t5), exposing a base model plus per-task head structs (BertForSequenceClassification, BertForQuestionAnswering, etc.) built from shared attention/embeddings/encoder submodules following the tch nn::Module pattern. The src/pipelines/ layer sits above these models and owns the user-facing task API — QuestionAnsweringModel, SummarizationModel, TranslationModel, ZeroShotClassificationModel and others — handling tokenization, batching, and post-processing so callers never touch raw tensors. A generation_utils.rs module (2,300+ lines) centralizes beam search, sampling, and constrained decoding shared across every generative pipeline (GPT-2, BART, T5, Marian, M2M100, NLLB, Pegasus), which is the load-bearing abstraction: changing its decoding loop affects every text-generation-capable model at once. Resource loading is abstracted behind a ResourceProvider trait (LocalResource/RemoteResource) so weights, configs, and vocabularies can come from disk or be fetched and cached transparently.

Tech Stack The crate targets stable Rust (2018 edition) and depends on tch 0.17 (libtorch/PyTorch C++ bindings) as its primary tensor/autograd backend behind the default download-libtorch feature, with an optional onnx feature swapping in ort 1.16 (ONNX Runtime) plus ndarray for a CPU-portable path. Tokenization delegates to the sibling rust_tokenizers crate (with an optional Hugging Face tokenizers crate integration behind the hf-tokenizers feature), serialization uses serde/serde_json, error handling uses thiserror, and pretrained resource downloading/caching uses cached-path with a choice of rustls-tls or default-tls. CI builds run across Linux, Windows, and macOS runners via GitHub Actions with default, no-default-features, and libtorch-download build variants.

Code Quality The repository has an extensive integration test suite — one dedicated test file per architecture (30+ files under tests/, e.g. bert.rs, gpt2.rs, t5.rs, nllb.rs) exercising each pipeline against real pretrained weights, plus a matching set of 50+ runnable examples under examples/ covering every task and several ONNX variants. clippy.toml and rustfmt.toml are checked into the repo, thiserror is used consistently for typed pipeline errors rather than string errors, and public APIs are documented with runnable doctests (no_run blocks) throughout the model modules. No unit-test-per-function pattern is evident for internal model layers — correctness is validated primarily at the pipeline/integration level against known model outputs.

What Makes It Unique rust-bert’s distinguishing choice is being a faithful native-Rust port of Hugging Face Transformers’ pipeline abstraction rather than just Rust bindings to Python or a from-scratch model zoo — it reproduces the same task-level pipeline API (pipeline("question-answering")-equivalent structs) while running fully in-process with no Python interpreter, and it supports an unusually broad architecture matrix (25+ model families spanning encoder, decoder, and encoder-decoder designs) with both GPU (libtorch) and CPU-portable (ONNX Runtime) execution paths in a single crate.