better handling of GPU out-of-memory issues

Author: clovis

text_preprocessing/preprocessor.py

@@ -143,11 +143,11 @@ def process_texts(
         fetched_texts = self.text_fetcher(
             texts, do_nlp=self.do_nlp, keep_all=keep_all, progress=progress, post_func=self.post_func
         )
-        if self.text_fetcher.text_object_type in ("para", "sent") and self.do_nlp is True:
+        if self.text_fetcher.text_object_type == "sent" and self.do_nlp is True:
             fetched_texts = self.nlp.pipe(
                 ((make_spacy_doc(self.nlp, tokens), c) for tokens, c in fetched_texts),
                 as_tuples=True,
-                batch_size=500,
+                batch_size=250,
             )
         for tokens, doc_count in fetched_texts:
             count += 1

text_preprocessing/spacy_helpers.py

@@ -502,6 +502,14 @@ def __normalize_token(self, orig_token: Token | PreprocessorToken) -> str:
         return token
 
 
+@Language.component("clear_trf_data")
+def clear_trf_data(doc):
+    """Clear the cache of a doc to free GPU memory"""
+    if hasattr(doc._, "trf_data"):
+        doc._.trf_data = None
+    return doc
+
+
 def load_language_model(language, normalize_options: dict[str, Any]) -> Language:
     """Load language model based on name"""
     nlp = None
@@ -521,17 +529,17 @@ def load_language_model(language, normalize_options: dict[str, Any]) -> Language
             normalize_options["ents_to_keep"],
         )
     ):
-        diabled_pipelines = ["tokenizer", "textcat"]
+        disabled_pipelines = ["tokenizer", "textcat"]
         if not normalize_options["pos_to_keep"]:
-            diabled_pipelines.append("tagger")
+            disabled_pipelines.append("tagger")
         if not normalize_options["ents_to_keep"]:
-            diabled_pipelines.append("ner")
+            disabled_pipelines.append("ner")
         model_loaded = ""
         set_gpu_allocator("pytorch")
-        prefer_gpu()
+        use_gpu = prefer_gpu()
         for model in possible_models:
             try:
-                nlp = spacy.load(model, exclude=diabled_pipelines)
+                nlp = spacy.load(model, exclude=disabled_pipelines)
                 print("Using Spacy model", model)
             except OSError:
                 pass
@@ -541,6 +549,8 @@ def load_language_model(language, normalize_options: dict[str, Any]) -> Language
         if nlp is None:
             print(f"No Spacy model installed for the {language} language. Stopping...")
             exit(-1)
+        if use_gpu is True:
+            nlp.add_pipe("clear_trf_data", last=True)
         nlp.add_pipe("postprocessor", config=normalize_options, last=True)
         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...")