[−][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 Sentiment Analysis, Named Entity Recognition, Question-Answering or Language Generation. More information on these can be found in the pipelines module.
use tch::Device;
use rust_bert::pipelines::question_answering::{QuestionAnsweringModel, QaInput};


let device = Device::cuda_if_available();
let qa_model = QuestionAnsweringModel::new(vocab_path,
                                           config_path,
                                           weights_path, device)?;

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
DistilBERTBERTRoBERTaGPTGPT2
Masked LM✅✅✅
Sequence classification✅✅✅
Token classification✅✅✅
Question answering✅✅✅
Multiple choices✅✅
Next token prediction✅✅
Natural Language Generation✅✅

Loading pre-trained models

The architectures defined in this crate are compatible with model trained in the Transformers library. The model configuration and vocabulary are downloaded directly from Huggingface's repository. The model weights 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). A Python script for downloading the required files & running the necessary steps is provided for all models classes in this library. Further models can be loaded by extending the python scripts to point to the desired model.

  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

Modules

bert

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

distilbert

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

gpt2

GPT2 (Radford et al.)

openai_gpt

GPT (Radford et al.)

pipelines

Ready-to-use NLP pipelines and models

roberta

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

Traits

Config

Utility to deserialize JSON config files