limit workers to 1 when using GPU to avoid running out of RAM when the buffer is full

Author: clovis

text_preprocessing/preprocessor.py

@@ -98,12 +98,14 @@ def __init__(
         if nlp_model is not None:
             self.nlp = nlp_model
         else:
-            self.nlp = load_language_model(self.language, self.normalize_options)
+            self.nlp, using_gpu = load_language_model(self.language, self.normalize_options)
         if workers is None:
             cpu_count = os.cpu_count() or 2
             self.workers = cpu_count - 1
         else:
             self.workers = workers
+        if using_gpu is True:
+            self.workers = 1
         ngrams = ngrams or 0
         if ngrams:
             self.ngram_config = {
@@ -124,7 +126,7 @@ def __init__(
             is_philo_db=is_philo_db,
             text_object_type=text_object_type,
             ngram_config=self.ngram_config,
-            workers=workers,
+            workers=self.workers,
         )
         if self.normalize_options["pos_to_keep"] or self.normalize_options["ents_to_keep"] or lemmatizer == "spacy":
             self.do_nlp = True

text_preprocessing/spacy_helpers.py

@@ -1,6 +1,7 @@
 """Helper functions for Spacy"""
 
 import os
+import pickle
 import re
 import sys
 import unicodedata
@@ -9,7 +10,6 @@
 from typing import Any, Dict, Iterable, List, Optional, Union
 from xml.sax.saxutils import unescape as unescape_xml
 
-import pickle
 import spacy
 from spacy.language import Language
 from spacy.tokens import Doc, Token
@@ -528,7 +528,7 @@ def clear_trf_data(doc):
     return doc
 
 
-def load_language_model(language, normalize_options: dict[str, Any]) -> Language:
+def load_language_model(language, normalize_options: dict[str, Any]) -> tuple[Language, bool]:
     """Load language model based on name"""
     nlp = None
     language = language.lower()
@@ -573,7 +573,7 @@ def load_language_model(language, normalize_options: dict[str, Any]) -> Language
         if normalize_options["ents_to_keep"] and "ner" not in nlp.pipe_names:
             print(f"There is no NER pipeline for model {model_loaded}. Exiting...")
             exit(-1)
-        return nlp
+        return nlp, use_gpu
     nlp = spacy.blank("en")
     nlp.add_pipe("postprocessor", config=normalize_options, last=True)
-    return nlp
+    return nlp, False