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

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.