Please check the build logs for more information.
See Builds for ideas on how to fix a failed build, or Metadata for how to configure docs.rs builds.
If you believe this is docs.rs' fault, open an issue.
rust-bert
Rust native BERT implementation. Port of Huggingface's 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.
The following models are currently implemented:
DistilBERT | BERT | RoBERTa | GPT | GPT2 | |
---|---|---|---|---|---|
Masked LM | ✅ | ✅ | ✅ | ||
Sequence classification | ✅ | ✅ | ✅ | ||
Token classification | ✅ | ✅ | ✅ | ||
Question answering | ✅ | ✅ | ✅ | ||
Multiple choices | ✅ | ✅ | |||
Next token prediction | ✅ | ✅ |
Ready-to-use pipelines
Based on Huggingface's pipelines, ready to use end-to-end NLP pipelines are available as part of this crate. The following capabilities are currently available:
1. Question Answering
Extractive question answering from a given question and context. DistilBERT model finetuned on SQuAD (Stanford Question Answering Dataset)
let device = cuda_if_available;
let qa_model = new?;
let question = String from;
let context = String from;
let answers = qa_model.predict;
Output:
[Answer { score: 0.9976814985275269, start: 13, end: 21, answer: "Amsterdam" }]
2. Sentiment analysis
Predicts the binary sentiment for a sentence. DistilBERT model finetuned on SST-2.
let device = cuda_if_available;
let sentiment_classifier = new?;
let input = ;
let output = sentiment_classifier.predict;
(Example courtesy of IMDb (http://www.imdb.com))
Output:
[
Sentiment { polarity: Positive, score: 0.9981985493795946 },
Sentiment { polarity: Negative, score: 0.9927982091903687 },
Sentiment { polarity: Positive, score: 0.9997248985164333 }
]
3. Named Entity Recognition
Extracts entities (Person, Location, Organization, Miscellaneous) from text. BERT cased large model finetuned on CoNNL03, contributed by the MDZ Digital Library team at the Bavarian State Library
let device = cuda_if_available;
let ner_model = new?;
let input = ;
let output = ner_model.predict;
Output:
[
Entity { word: "Amy", score: 0.9986, label: "I-PER" }
Entity { word: "Paris", score: 0.9985, label: "I-LOC" }
Entity { word: "Paris", score: 0.9988, label: "I-LOC" }
Entity { word: "France", score: 0.9993, label: "I-LOC" }
]
Base models
The base model and task-specific heads are also available for users looking to expose their own transformer based models.
Examples on how to prepare the date using a native tokenizers Rust library are available in ./examples
for BERT, DistilBERT and RoBERTa.
Note that when importing models from Pytorch, the convention for parameters naming needs to be aligned with the Rust schema. Loading of the pre-trained weights will fail if any of the model parameters weights cannot be found in the weight files.
If this quality check is to be skipped, an alternative method load_partial
can be invoked from the variables store.
Setup
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.
- Compile the package:
cargo build --release
- 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/{}
- Run the example
cargo run --release