spacy-accelerate

Accelerate spaCy transformer pipelines with TensorRT and ONNX Runtime. Drop-in replacement — one line of code. Speedup depends on your GPU: 1.2–3.5× faster inference on the tested setups, with small accuracy deltas.

Repository: GitHub • Package: PyPI

Requirements

• Python 3.11+
• CUDA 12.x
• NVIDIA GPU with TensorRT support (Ampere / Ada Lovelace recommended)
• spaCy 3.8+ with spacy-transformers

Installation

PyPI

pip install spacy-accelerate
pip install --force-reinstall \
    --extra-index-url https://pypi.nvidia.com \
    onnxruntime-gpu==1.23.2

The second command installs the TensorRT-enabled build of onnxruntime-gpu from NVIDIA's index. It is required because the default PyPI build does not include the TensorRT execution provider.

Source / editable install

pip install -r requirements.txt
pip install -e . --no-deps

Verify the installation

python -m spacy_accelerate

You should see TensorRT EP : OK and CUDA EP : OK in the output.

Quick Start

import spacy
import spacy_accelerate

# Load your spaCy transformer model
nlp = spacy.load("en_core_web_trf")

# Optimize with one line
nlp = spacy_accelerate.optimize(nlp, precision="fp16")

# Use as normal — same API, faster inference
doc = nlp("Apple Inc. was founded by Steve Jobs in Cupertino.")
print([(ent.text, ent.label_) for ent in doc.ents])
# [('Apple Inc.', 'ORG'), ('Steve Jobs', 'PERSON'), ('Cupertino', 'GPE')]

# Batch processing works too
texts = ["Text one.", "Text two.", "Text three."]
docs = list(nlp.pipe(texts, batch_size=32))

API Reference

optimize(nlp, **kwargs)

Optimize a spaCy transformer pipeline with ONNX Runtime / TensorRT.

Parameters:

TensorRT-specific parameters:

Advanced parameters:

Returns: The optimized spacy.Language object (modified in-place).

Cache Management

import spacy_accelerate

# List cached models
cached = spacy_accelerate.list_cached()
print(f"Cached models: {cached}")

# Get cache size
size_bytes = spacy_accelerate.get_cache_size()
print(f"Cache size: {size_bytes / 1024**2:.1f} MB")

# Clear cache
cleared = spacy_accelerate.clear_cache()
print(f"Cleared {cleared} cache entries")

Examples

Maximum performance with TensorRT

import spacy
import spacy_accelerate

nlp = spacy.load("en_core_web_trf")

nlp = spacy_accelerate.optimize(
    nlp,
    provider="tensorrt",
    precision="fp16",
    trt_max_workspace_size=8 * 1024**3,  # 8 GB
    trt_builder_optimization_level=5,     # Maximum optimization
)

# First inference builds the TensorRT engine (cached for subsequent runs)
doc = nlp("TensorRT provides maximum inference speed.")

Custom cache directory

import spacy
import spacy_accelerate

nlp = spacy.load("en_core_web_trf")

nlp = spacy_accelerate.optimize(
    nlp,
    cache_dir="/path/to/custom/cache",
    precision="fp16",
)

Verbose mode for debugging

import spacy
import spacy_accelerate

nlp = spacy.load("en_core_web_trf")

nlp = spacy_accelerate.optimize(
    nlp,
    verbose=True,
)

Performance

Conditions: en_core_web_trf, CoNLL-2003 test set, batch_size=128, 1 warmup pass + 3 measured passes averaged.

How much speedup to expect

The relative gain depends on your GPU. On a faster GPU, PyTorch is already faster — GPU compute stops being the bottleneck and the overhead shifts to the spaCy pipeline (tokenization, NER decoding, Python). ONNX Runtime does not accelerate that part.

NER-only mode disables tagger, parser, attribute_ruler, lemmatizer. Only the transformer + NER component run, which reduces non-GPU overhead and yields higher speedups.

RTX 4000 SFF Ada Generation — full results ¹

Full pipeline:

NER only:

A100 80 GB — full results ²

Full pipeline:

NER only:

On A100, CUDA EP FP16 outperforms TensorRT FP16. This is expected: A100 was optimized for BF16 and large-batch datacenter workloads; TensorRT gains are less pronounced for the NLP batch sizes typical in spaCy pipelines.

RTX 4090 — full results ³

Full pipeline:

NER only:

On RTX 4090, the PyTorch baseline is already fast (~10k WPS full pipeline). The remaining gains come mostly from FP16 precision, not from the runtime switch. Switching to NER-only mode shows the clearest improvement (1.31×).

¹ Cloud instance (Hetzner). Ada Lovelace architecture benefits most from TensorRT FP16. ² Virtual partition (GRID A100D-80C). On this GPU CUDA EP FP16 is the recommended provider — TensorRT does not outperform it for typical spaCy batch sizes. ³ Cloud instance (RunPod). Local RTX 4090 results may differ due to power limits or virtualization overhead.

Supported Models

Currently tested, confirmed, and supported:

• en_core_web_trf (RoBERTa-based)

Other spaCy transformer packages should be treated as unsupported for now, even if related architecture-detection code exists internally.

How It Works

1. Weight Mapping — extracts transformer weights from spaCy's internal format and maps them to HuggingFace format.
2. ONNX Export — exports the model to ONNX with dynamic batch and sequence dimensions.
3. FP16 Optimization (optional) — applies BERT-style graph optimizations and converts weights to FP16.
4. Runtime Patching — replaces the PyTorch transformer with an ONNX Runtime proxy that provides the same spaCy interface.
5. Caching — converted models are cached to disk to avoid re-conversion on subsequent runs.

Troubleshooting

TensorRT provider not available

Run the diagnostic tool first:

python -m spacy_accelerate

If you see TensorRT EP : MISSING, the NVIDIA build of onnxruntime-gpu is not installed. Fix it with:

pip install --force-reinstall \
    --extra-index-url https://pypi.nvidia.com \
    onnxruntime-gpu==1.23.2

libnvinfer.so / libcublas.so / libcublasLt.so not found

If you see errors like libnvinfer.so.10: cannot open shared object file:

Automatic fix: spacy-accelerate configures the native library paths on import. Make sure to import spacy_accelerate before calling spacy.require_gpu() or creating any ONNX Runtime sessions.

Manual fix: Set LD_LIBRARY_PATH explicitly:

SITE=$(python -c "import site; print(site.getsitepackages()[0])")
export LD_LIBRARY_PATH="$SITE/tensorrt_libs:$SITE/nvidia/cublas/lib:$SITE/nvidia/cuda_runtime/lib:$SITE/nvidia/cudnn/lib:$LD_LIBRARY_PATH"

CUDA out of memory

Reduce workspace size or batch size:

# For TensorRT provider — reduce workspace
nlp = spacy_accelerate.optimize(
    nlp,
    provider="tensorrt",
    trt_max_workspace_size=2 * 1024**3,  # 2 GB instead of 4 GB
)

# For any provider — reduce batch size
nlp = spacy_accelerate.optimize(
    nlp,
    max_batch_size=16,
)

First inference is slow

This only applies to the TensorRT provider. TensorRT compiles an optimized engine on the first run — this can take tens of seconds. Subsequent runs reuse the cached engine and are fast.

The timing cache is enabled by default and carries over build history between runs. If build time matters, prefer a lower optimization level:

nlp = spacy_accelerate.optimize(
    nlp,
    provider="tensorrt",
    trt_timing_cache=True,             # on by default
    trt_builder_optimization_level=3,  # lower = faster build, 5 = best runtime perf
)

Contributing

Contributions are welcome! Please open an issue or submit a pull request on GitHub.

License

MIT License