nlp-uoregon / trankit

Trankit is a Light-Weight Transformer-based Python Toolkit for Multilingual Natural Language Processing
Apache License 2.0
724 stars 99 forks source link

GPU on Apple M1 chip support #62

Open carschno opened 1 year ago

carschno commented 1 year ago

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:

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

Confirming that MPS is available through PyTorch:

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