RustBERTa-SNLI
A Rust implementation of a RoBERTa classification model for the SNLI dataset, with support for fine-tuning, predicting, and serving. This is built on top of tch-rs and rust-bert.
Background
This was the result of the AI2 2020 Employee Hackathon. The motivation for this project was to demonstrate that Rust is already a viable alternative to Python for deep learning projects in the context of both research and production.
Setup
RustBERTa-SNLI is packaged as a Rust binary, and will work on any operating system that PyTorch supports.
The only prerequisite is that you have the Rust toolchain installed. So if you're already a Rustacean, skip ahead to the "Additional setup for CUDA" section (optional).
Now, luckily, installing Rust is nothing like installing a proper Python environment, i.e. it doesn't require a PhD in system administration or the courage to blindly run every sudo command you can find on Stack Overflow until something works or completely breaks your computer.
All you have to do is run this:
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
Then just make sure ~/.cargo/bin is in your $PATH, and you're good to go. You can test for success by running rustup --version and cargo --version.
rustupcan be used to update your toolchain when a new version of Rust is released (which happens monthly).cargois used to compile, run, and test your code, as well as to build documentation, publish your crate (the Rust term for a module/library) to crates.io, and install binaries from other crates on crates.io.
Additional setup for CUDA
If you have CUDA-enabled GPUs available on your machine, you'll probably want to compile this library with CUDA support.
To do that, you just need to download the right version of LibTorch from the PyTorch website: https://pytorch.org/get-started/locally/.
Then unzip the downloaded file to someplace safe like ~/torch/libtorch and set the environment variables:
LIBTORCH="$HOME/torch/libtorch" # or wherever you unzipped it
LD_LIBRARY_PATH="$HOME/torch/libtorch/lib:$LD_LIBRARY_PATH"
cargo. But if you accidentally start the build for this step, just delete the target/ directory and start over.
Compiling and running
To build the release binary, just run make.
To see all of the available commands, run
./roberta-snli --help
For example, to fine-tune a pretrained RoBERTa model, run
./roberta-snli train --out weights.ot
To interactively get predictions with a fine-tuned model, run
./roberta-snli predict --weigths weights.ot
To evaluate a fine-tuned model on the test set, run
./roberta-snli evaluate
And to serve a fine-tuned model as a production-grade webservice with batched prediction, run
./roberta-snli serve
This will serve on port 3030 by default. You can then test it out by running:
curl \
-d '{"premise":"A soccer game with multiple people playing.","hypothesis":"Some people are playing a sport."}' \
-H "Content-Type: application/json" \
http://localhost:3030/predict
You can also test the batching functionality by sending a bunch of requests at once with:
./scripts/test_server.sh
26 Dec 16, 2022
327 Dec 26, 2022
662 Dec 31, 2022
84 Dec 15, 2022
19 Dec 24, 2022
4 Sep 4, 2023
5 Apr 4, 2022
2 Jul 5, 2018
3 Jul 24, 2023
4 Nov 21, 2023
34 Dec 20, 2022
322 Dec 26, 2022
81 Dec 6, 2022
805 Dec 28, 2022
72 Dec 16, 2022
75 Jan 8, 2023
1.3k Jan 8, 2023