In [9]: from trankit import Pipeline
In [10]: p = Pipeline(lang='english')
Loading pretrained XLM-Roberta, this may take a while...
Loading tokenizer for english
Loading tagger for english
Loading lemmatizer for english
Loading NER tagger for english
==================================================
Active language: english
==================================================
In [11]: p._use_gpu
Out[11]: False
In [12]: import torch
In [13]: torch.has_mps
Out[13]: True
A look into pipeline.py shows that it only works on CUDA:
def _setup_config(self, lang):
torch.cuda.empty_cache()
# decide whether to run on GPU or CPU
if self._gpu and torch.cuda.is_available():
self._use_gpu = True
This is a feature request to add support for the Apple M1 chip, which is supported by PyTorch since v1.12.
Currently, Trankit only seems to use Cuda:
Confirming that MPS is available through PyTorch:
A look into
pipeline.py
shows that it only works on CUDA: