Make output dir optional

Molmed · Nov 6, 2024 · ac7e713 · ac7e713
1 parent b1eaae5
commit ac7e713
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/conformist/base_cop.py b/src/conformist/base_cop.py
@@ -113,12 +113,10 @@ def do_validation_trial(self,
 
     def predict(self,
                 pds,
-                output_dir,
+                output_dir=None,
                 validate=False,
                 upset_plot_color="black"):
 
-        self.create_output_dir(output_dir)
-
         self.smx = pds.smx
         self.val_smx = pds.smx
         self.val_labels = pds.labels_idx
@@ -131,6 +129,11 @@ def predict(self,
         prediction_sets_text = self.prediction_sets_to_text(prediction_sets)
         formatted_predictions = pds.prediction_sets_df(prediction_sets_text)
 
+        if not output_dir:
+            return formatted_predictions
+
+        self.create_output_dir(output_dir)
+
         # -- WRITE PREDICTIONS TO CSV
         formatted_predictions.to_csv(f'{self.output_dir}/prediction_sets.csv')