This project provides an example implementation for inference on the BERT family of models. The following compatible
bert-variants: roberta-base(default)/roberta-large, bert-base-uncased/bert-large-uncased/bert-base-cased/bert-large-cased
can be loaded as following. The pre-trained weights and config files are automatically downloaded
from: HuggingFace Model hub
Add this to your Cargo.toml:
[dependencies]
bert-burn = { git = "https://github.com/tracel-ai/models", package = "bert-burn", default-features = false }Example usage for getting sentence embedding from given input text. The model supports multiple backends from burn
(e.g. ndarray, wgpu, tch-gpu, tch-cpu, cuda) which can be selected using the --features flag. An example with wgpu
backend is shown below. The fusion flag is used to enable kernel fusion for the wgpu backend. It is not required
with other backends. The safetensors flag is used to support loading weights in safetensors format via candle-core
crate.
cd bert-burn/
# Get sentence embeddings from the RobBERTa encoder (default)
cargo run --example infer-embedding --release --features wgpu,fusion,safetensors
# Using bert-base-uncased model
cargo run --example infer-embedding --release --features wgpu,fusion,safetensors bert-base-uncased
# Using roberta-large model
cargo run --example infer-embedding --release --features wgpu,fusion,safetensors roberta-large