Merge pull request #91 from huridocs/error-handling

Handle GPU out of memory
huridocs · Sep 27, 2024 · de1d839 · de1d839
2 parents 11e6f10 + eaea29b
commit de1d839
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/extractors/pdf_to_multi_option_extractor/multi_labels_methods/SetFitEnglishMethod.py b/src/extractors/pdf_to_multi_option_extractor/multi_labels_methods/SetFitEnglishMethod.py
@@ -17,6 +17,8 @@
 from extractors.pdf_to_multi_option_extractor.MultiLabelMethod import MultiLabelMethod
 from send_logs import send_logs
 
+import gc
+
 
 class SetFitEnglishMethod(MultiLabelMethod):
     model_name = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
@@ -100,6 +102,10 @@ def train(self, extraction_data: ExtractionData):
 
         trainer.model.save_pretrained(self.get_model_path())
 
+        del model
+        del trainer
+        gc.collect()
+
     def predict(self, multi_option_data: ExtractionData) -> list[list[Option]]:
         model = SetFitModel.from_pretrained(self.get_model_path(), trust_remote_code=True)
         predict_texts = [sample.pdf_data.get_text() for sample in multi_option_data.samples]

diff --git a/...tors/pdf_to_multi_option_extractor/multi_labels_methods/SingleLabelSetFitEnglishMethod.py b/...tors/pdf_to_multi_option_extractor/multi_labels_methods/SingleLabelSetFitEnglishMethod.py
@@ -1,3 +1,4 @@
+import gc
 import os
 import shutil
 from os.path import join, exists
@@ -115,6 +116,10 @@ def train(self, extraction_data: ExtractionData):
 
         trainer.model.save_pretrained(self.get_model_path())
 
+        del model
+        del trainer
+        gc.collect()
+
     def predict(self, multi_option_data: ExtractionData) -> list[list[Option]]:
         model = SetFitModel.from_pretrained(self.get_model_path(), trust_remote_code=True)
         predict_texts = [sample.pdf_data.get_text() for sample in multi_option_data.samples]