Update BERTopic/my_topic_modeling.py

BERTopic/my_topic_modeling.py

@@ -52,9 +52,13 @@ from cuml.manifold import UMAP
 import gensim.corpora as corpora
 from gensim.models.coherencemodel import CoherenceModel
 
+import torch
+from GPUtil import showUtilization as gpu_usage
+from numba import cuda
+
 
 # Get working directory
-working_dir =  os.path.abspath(os.path.join("/notebooks", "TopicModelingRepo"))
+working_dir =  os.path.abspath(os.path.join("/workspace", "TopicModelingRepo"))
 data_dir = os.path.join(working_dir, 'data')
 lib_dir = os.path.join(working_dir, 'libs')
 outer_output_dir = os.path.join(working_dir, 'outputs')
@@ -74,6 +78,18 @@ doc_type = 'reviews'
 doc_level = 'sentence'
 target_col = 'normalized_content'
 
+def free_gpu_cache():
+    print("Initial GPU Usage")
+    gpu_usage()                             
+
+    torch.cuda.empty_cache()
+
+    cuda.select_device(0)
+    cuda.close()
+    cuda.select_device(0)
+
+    print("GPU Usage after emptying the cache")
+    gpu_usage()
 
 
 def create_logger_file_and_console(path_file):
@@ -459,7 +475,7 @@ def working(args: argparse.Namespace, name_dataset: str):
     
     # Create model
     fandc_logger.log(logging.INFO, f'Create model for {name_dataset} dataset')
-    topic_model = create_model_bertopic_booking(args.n_topics) if name_dataset == 'booking' else create_model_bertopic_tripadvisor(args.n_topics)
+    topic_model = create_model_bertopic_booking(args.n_topics)
     
     # Fitting model
     fandc_logger.log(logging.INFO, f'Training model for {name_dataset} dataset')
@@ -607,3 +623,5 @@ if __name__ == "__main__":
         working(args, 'tripadvisor')
     else:
         working(args, args.name_dataset)
+
+    free_gpu_cache()