Add GPU option for spaCy preprocessor (#1523)

Author: henryre

snorkel/labeling/lf/nlp.py

@@ -18,6 +18,7 @@ class SpacyPreprocessorParameters(NamedTuple):
     disable: Optional[List[str]]
     pre: List[BasePreprocessor]
     memoize: bool
+    gpu: bool
 
 
 class SpacyPreprocessorConfig(NamedTuple):
@@ -47,6 +48,7 @@ def _create_or_check_preprocessor(
         disable: Optional[List[str]],
         pre: List[BasePreprocessor],
         memoize: bool,
+        gpu: bool,
     ) -> None:
         # Create a SpacyPreprocessor if one has not yet been instantiated.
         # Otherwise, check that configuration matches already instantiated one.
@@ -57,6 +59,7 @@ def _create_or_check_preprocessor(
             disable=disable,
             pre=pre,
             memoize=memoize,
+            gpu=gpu,
         )
         if not hasattr(cls, "_nlp_config"):
             nlp = cls._create_preprocessor(parameters)
@@ -78,9 +81,10 @@ def __init__(
         language: str = EN_CORE_WEB_SM,
         disable: Optional[List[str]] = None,
         memoize: bool = True,
+        gpu: bool = False,
     ) -> None:
         self._create_or_check_preprocessor(
-            text_field, doc_field, language, disable, pre or [], memoize
+            text_field, doc_field, language, disable, pre or [], memoize, gpu
         )
         super().__init__(name, f, resources=resources, pre=[self._nlp_config.nlp])
 
@@ -128,6 +132,8 @@ class NLPLabelingFunction(BaseNLPLabelingFunction):
         See https://spacy.io/usage/processing-pipelines#disabling
     memoize
         Memoize preprocessor outputs?
+    gpu
+        Prefer Spacy GPU processing?
 
     Raises
     ------
@@ -176,13 +182,15 @@ def __init__(
         language: str = EN_CORE_WEB_SM,
         disable: Optional[List[str]] = None,
         memoize: bool = True,
+        gpu: bool = False,
     ) -> None:
         super().__init__(name, resources, pre)
         self.text_field = text_field
         self.doc_field = doc_field
         self.language = language
         self.disable = disable
         self.memoize = memoize
+        self.gpu = gpu
 
     def __call__(self, f: Callable[..., int]) -> BaseNLPLabelingFunction:
         """Wrap a function to create an ``BaseNLPLabelingFunction``.
@@ -210,6 +218,7 @@ def __call__(self, f: Callable[..., int]) -> BaseNLPLabelingFunction:
             language=self.language,
             disable=self.disable,
             memoize=self.memoize,
+            gpu=self.gpu,
         )
 
 

snorkel/labeling/lf/nlp_spark.py

@@ -36,6 +36,8 @@ class SparkNLPLabelingFunction(BaseNLPLabelingFunction):
         See https://spacy.io/usage/processing-pipelines#disabling
     memoize
         Memoize preprocessor outputs?
+    gpu
+        Prefer Spacy GPU processing?
 
     Raises
     ------

snorkel/preprocess/nlp.py

@@ -40,6 +40,8 @@ class SpacyPreprocessor(Preprocessor):
         Preprocessors to run before this preprocessor is executed
     memoize
         Memoize preprocessor outputs?
+    gpu
+        Prefer Spacy GPU processing?
     """
 
     def __init__(
@@ -50,6 +52,7 @@ def __init__(
         disable: Optional[List[str]] = None,
         pre: Optional[List[BasePreprocessor]] = None,
         memoize: bool = False,
+        gpu: bool = False,
     ) -> None:
         name = type(self).__name__
         super().__init__(
@@ -59,6 +62,9 @@ def __init__(
             pre=pre,
             memoize=memoize,
         )
+        self.gpu = gpu
+        if self.gpu:
+            spacy.prefer_gpu()
         self._nlp = spacy.load(language, disable=disable or [])
 
     def run(self, text: str) -> FieldMap:  # type: ignore