fallback to cpu on cuda oom error

VERSION

-1.11.2
+1.11.3

texta_mlp/mlp.py

@@ -79,7 +79,6 @@ class MLP:
         self.default_lang = default_language_code
         self.use_default_lang = use_default_language_code
         self.resource_dir = resource_dir
-        self.use_gpu = use_gpu
 
         self.resource_dir_pathlib = pathlib.Path(resource_dir)
         self.not_entities_path = self.resource_dir_pathlib / "concatenator" / "not_entities.txt"
@@ -88,7 +87,13 @@ class MLP:
 
         self.prepare_resources(refresh_data)
 
-        self.stanza_pipelines = self._load_stanza_pipelines(logging_level)
+        # This is for CUDA OOM exceptions. Fall back to CPU if needed.
+        try:
+            self.stanza_pipelines = self._load_stanza_pipelines(logging_level, use_gpu)
+        except RuntimeError:
+            # Try loading using CPU
+            self.stanza_pipelines = self._load_stanza_pipelines(logging_level, False)
+
         self.entity_mapper = self._load_entity_mapper()
         self.loaded_entity_files = []
 
@@ -258,7 +263,7 @@ class MLP:
         return [entity for sentence in pipeline.sentences for entity in sentence.entities]
 
 
-    def _load_stanza_pipelines(self, logging_level):
+    def _load_stanza_pipelines(self, logging_level, use_gpu):
         """
         Initializes Stanza Pipeline objects all at once to save time later.
         """