Wrapper around detoxify package for faster inference using ONNX runtime.
Project description
Speedtoxify :rocket:
Fast :speak_no_evil: Detoxify inference with ONNX runtime
:zap: Benchmarks | :gear: install | :star2: Quick Start | :page_with_curl: Docs
Speedtoxify is a wrapper around detoxify
that speeds up inference by 2-4x by using ONNX runtime.
Detoxify is a NLP library for detecting toxic / inappropriate / profane texts. Speedtoxify makes use of their pretrained models and runs them in ONNX runtime for much faster inference speeds, which makes it the better option for being used in production.
Speedtoxify provides the same Python API as Detoxify, so it can be used as a drop-in replacement.
However, if your focus is on fine-tuning / re-training the models with your own data, please refer to Detoxify.
:zap: Lightning fast
Model | Batch size | Device | Detoxify (ms/sample) | Speedtoxify (ms/sample) | Speedup |
---|---|---|---|---|---|
original-small | 8 | cpu | 13.34 | 5.43 | 2.46x |
original-small | 1 | cpu | 31.07 | 13.03 | 2.38x |
original-small | 8 | cuda | 1.55 | 0.79 | 1.98x |
original-small | 1 | cuda | 11.17 | 3.24 | 3.44x |
original | 8 | cpu | 22.99 | 5.39 | 4.26x |
original | 1 | cpu | 31.48 | 13.11 | 2.40x |
original | 8 | cuda | 1.60 | 0.75 | 2.12x |
original | 1 | cuda | 12.13 | 3.37 | 3.60x |
Evaluation script can be found in test_speed.py.
Evaluation is done on my laptop with AMD 4900HS and Nvidia 2060 Max-Q.
:gear: Installation
Pip
pip install speedtoxify
GPU Inference
Please additionally install onnxruntime-gpu
for inference on gpus.
Requires the machine to have CUDA installed.
pip install onnxruntime-gpu
:star2: Quick start
Speedtoxify provides the identical Python API as Detoxify.
from speedtoxify import Speedtoxify
model = Speedtoxify("original-small")
# Exporting to onnx format to ~/.cache/detoxify_onnx/original-small.onnx...
# Using framework PyTorch: 1.11.0+cu102
# Removing shared weights from ~/.cache/detoxify_onnx/original-small.onnx...
# Validating ONNX model...
# -[✓] ONNX model output names match reference model ({'logits'})
# - Validating ONNX Model output "logits":
# -[✓] (2, 6) matches (2, 6)
# -[✓] all values close (atol: 1e-05)
res = model.predict("I hate you!")
print(res)
# {'toxicity': 0.9393415, 'severe_toxicity': 0.015587699, 'obscene': 0.039672945, 'threat': 0.0733101, 'insult': 0.15676126, 'identity_attack': 0.019178415}
Please refer to detoxify for available model types.
The first time Speedtoxify("original-small")
is called, an onnx model is
exported and stored at ~/.cache/detoxify_onnx
.
This directory can be customized in the cache_dir
argument to
Speedtoxify()
.
:page_with_curl: Documentation
Please refer to docs.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for speedtoxify-0.0.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 1b80d06ab8296fa29aba6fa56bd163a8d29b25af02295629d5308957d61ed08a |
|
MD5 | 5acae71548664687c1c4a0e73c392b57 |
|
BLAKE2b-256 | d3ad63e7009b68924035592ab94702cab069f9b99cec1c4f4d3026e3156b820a |