[][src]Crate rust_bert

Ready-to-use NLP pipelines and Transformer-based models

Rust native Transformer-based models implementation. Port of the Transformers library, using the tch-rs crate and pre-processing from rust-tokenizers. Supports multithreaded tokenization and GPU inference. This repository exposes the model base architecture, task-specific heads (see below) and ready-to-use pipelines.

Quick Start

This crate can be used in two different ways:

  • Ready-to-use NLP pipelines for:
    • Translation
    • Summarization
    • Multi-turn dialogue
    • Sentiment Analysis
    • Named Entity Recognition
    • Question-Answering
    • Language Generation.

More information on these can be found in the pipelines module

use rust_bert::pipelines::question_answering::{QaInput, QuestionAnsweringModel};

let qa_model = QuestionAnsweringModel::new(Default::default())?;

let question = String::from("Where does Amy live ?");
let context = String::from("Amy lives in Amsterdam");
let answers = qa_model.predict(&vec![QaInput { question, context }], 1, 32);
  • Transformer models base architectures with customized heads. These allow to load pre-trained models for customized inference in Rust
Masked LM
Sequence classification
Token classification
Question answering
Multiple choices
Next token prediction
Natural Language Generation

Loading pre-trained models

A number of pretrained model configuration, weights and vocabulary are downloaded directly from Huggingface's model repository. The list of models available with Rust-compatible weights is available in the example ./examples/download_all_dependencies.rs. Additional models can be added if of interest, please raise an issue.

In order to load custom weights to the library, these need to be converter to a binary format that can be read by Libtorch (the original .bin files are pickles and cannot be used directly). Several Python scripts to load Pytorch weights and convert them to the appropriate format are provided and can be adapted based on the model needs.

The procedure for building custom weights or re-building pretrained weights is as follows:

  1. Compile the package: cargo build --release
  2. Download the model files & perform necessary conversions
    • Set-up a virtual environment and install dependencies
    • run the conversion script python /utils/download-dependencies_{MODEL_TO_DOWNLOAD}.py. The dependencies will be downloaded to the user's home directory, under ~/rustbert/{}
  3. Run the example cargo run --release



ALBERT: A Lite BERT for Self-supervised Learning of Language Representations (Lan et al.)


BART (Lewis et al.)


BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin et al.)


DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter (Sanh et al.)


Electra: Pre-training Text Encoders as Discriminators Rather Than Generators (Clark et al.)


GPT2 (Radford et al.)




GPT (Radford et al.)


Ready-to-use NLP pipelines and models


Resource definitions for model weights, vocabularies and configuration files


RoBERTa: A Robustly Optimized BERT Pretraining Approach (Liu et al.)


T5 (Text-To-Text Transfer Transformer)



Utility to deserialize JSON config files