Merge branch 'release/0.3.2'

.pylintrc

@@ -9,3 +9,7 @@ disable =
 [FORMAT]
 
 max-line-length = 128
+
+[TYPECHECK]
+
+generated-members=torch.device

README.rst

@@ -113,8 +113,8 @@ Models
 Model Usage
 ^^^^^^^^^^^
 
-| You may use our model directly from the HuggingFace's transformers library
-| 您可直接透過 HuggingFace's transformers 套件使用我們的模型
+| You may use our model directly from the HuggingFace's transformers library.
+| 您可直接透過 HuggingFace's transformers 套件使用我們的模型。
 
 .. code-block:: bash
 

ckip_transformers/__init__.py

@@ -10,7 +10,7 @@
 __copyright__ = "2021 CKIP Lab"
 
 __title__ = "CKIP Transformers"
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 __description__ = "CKIP Transformers"
 __license__ = "GPL-3.0"
 

ckip_transformers/nlp/driver.py

@@ -32,7 +32,7 @@ class CkipWordSegmenter(CkipTokenClassification):
             The pretrained model name provided by CKIP Transformers.
         model_name : ``str`` *optional*, overwrites **model**
             The custom pretrained model name (e.g. ``'ckiplab/bert-base-chinese-ws'``).
-        device : ``int``, *optional*, defaults to -1
+        device : ``int`` or ``torch.device``, *optional*, defaults to -1
             Device ordinal for CPU/GPU supports.
             Setting this to -1 will leverage CPU, a positive will run the model on the associated CUDA device id.
     """
@@ -132,7 +132,7 @@ class CkipPosTagger(CkipTokenClassification):
             The pretrained model name provided by CKIP Transformers.
         model_name : ``str`` *optional*, overwrites **model**
             The custom pretrained model name (e.g. ``'ckiplab/bert-base-chinese-pos'``).
-        device : ``int``, *optional*, defaults to -1
+        device : ``int`` or ``torch.device``, *optional*, defaults to -1
             Device ordinal for CPU/GPU supports.
             Setting this to -1 will leverage CPU, a positive will run the model on the associated CUDA device id.
     """
@@ -222,7 +222,7 @@ class CkipNerChunker(CkipTokenClassification):
             The pretrained model name provided by CKIP Transformers.
         model_name : ``str`` *optional*, overwrites **model**
             The custom pretrained model name (e.g. ``'ckiplab/bert-base-chinese-ner'``).
-        device : ``int``, *optional*, defaults to -1
+        device : ``int`` or ``torch.device``, *optional*, defaults to -1
             Device ordinal for CPU/GPU supports.
             Setting this to -1 will leverage CPU, a positive will run the model on the associated CUDA device id.
     """

ckip_transformers/nlp/util.py

@@ -51,7 +51,7 @@ class CkipTokenClassification(metaclass=ABCMeta):
             The pretrained model name (e.g. ``'ckiplab/bert-base-chinese-ws'``).
         tokenizer_name : ``str``, *optional*, defaults to **model_name**
             The pretrained tokenizer name (e.g. ``'bert-base-chinese'``).
-        device : ``int``, *optional*, defaults to -1
+        device : ``int`` or ``torch.device``, *optional*, defaults to -1
             Device ordinal for CPU/GPU supports.
             Setting this to -1 will leverage CPU, a positive will run the model on the associated CUDA device id.
     """
@@ -61,12 +61,17 @@ def __init__(
         model_name: str,
         tokenizer_name: Optional[str] = None,
         *,
-        device: int = -1,
+        device: Union[int, torch.device] = -1,
     ):
         self.model = AutoModelForTokenClassification.from_pretrained(model_name)
         self.tokenizer = BertTokenizerFast.from_pretrained(tokenizer_name or model_name)
 
-        self.device = torch.device("cpu" if device < 0 else f"cuda:{device}")  # pylint: disable=no-member
+        # Allow passing a customized torch.device.
+        if isinstance(device, torch.device):
+            self.device = device
+        else:
+            self.device = torch.device("cpu" if device < 0 else f"cuda:{device}")
+
         self.model.to(self.device)
 
     ########################################################################################################################