Move callback to callback folder

examples/6_advanced_features/1_wandb_logging.py

@@ -7,79 +7,15 @@
 """
 from __future__ import annotations
 
-from ConfigSpace import Configuration, ConfigurationSpace
-
 import numpy as np
-from smac import HyperparameterOptimizationFacade, Scenario
+from ConfigSpace import Configuration, ConfigurationSpace
 from sklearn import datasets
-from sklearn.svm import SVC
 from sklearn.model_selection import cross_val_score
+from sklearn.svm import SVC
 
-from smac import Callback
 import smac
-from wandb import Table
-
-from smac.runhistory import TrialInfo, TrialValue
-from dataclasses import asdict
-
-
-class WandBCallback(Callback):
-    def __init__(
-            self,
-            project: str,
-            entity: str,
-            id: str | None = None,
-            outdir: str | None = None,
-            mode: str | None = None,
-            resume: str = "allow",
-            job_type: str | None = None,
-            group: str | None = None,
-            config: dict | str | None = None,
-            save_code: bool = True,
-            **kwargs
-        ) -> None:
-        import wandb
-        self.run = wandb.init(
-            id=id,
-            resume=resume,
-            mode=mode,
-            project=project,
-            job_type=job_type,
-            entity=entity,
-            group=group,
-            dir=outdir,
-            config=config,
-            save_code=save_code,
-            **kwargs
-        )
-        super().__init__()
-
-    def on_tell_end(self, smbo: smac.main.smbo.SMBO, info: TrialInfo, value: TrialValue) -> bool | None:
-        info_dict = asdict(info)
-        info_dict["config"] = info_dict["config"].get_dictionary()
-        value_dict = asdict(value)
-        log_dict = info_dict | value_dict
-        log_dict["step"] = smbo.runhistory.finished
-        self.run.log(data=log_dict)
-        return super().on_tell_end(smbo, info, value)
-
-
-    def on_end(self, smbo: smac.main.smbo.SMBO) -> None:
-        intensifier_data = smbo.intensifier.get_data()
-        trajectory = intensifier_data["trajectory"]
-        import pandas as pd
-        df = pd.DataFrame(data=trajectory)
-        print(df)
-        # trajectory = Table(dataframe=df, allow_mixed_types=True)
-        df["costs"] = df["costs"].apply(lambda x: x[0])  # TODO properly log multi costs
-        for index, row in df.iterrows():
-            print(dict(row))
-            self.run.log(dict(row))
-        self.run.finish()
-        return super().on_end(smbo)
-
-
-
+from smac import HyperparameterOptimizationFacade, Scenario
+from smac.callback import WandBCallback
 
 iris = datasets.load_iris()
 
@@ -103,4 +39,4 @@ def train(config: Configuration, seed: int = 0) -> float:
 
 # Use SMAC to find the best configuration/hyperparameters
 smac = HyperparameterOptimizationFacade(scenario, train, callbacks=[wandb_callback], overwrite=True)
-incumbent = smac.optimize()
+incumbent = smac.optimize()

smac/callback/__init__.py

@@ -1,7 +1,9 @@
 from smac.callback.callback import Callback
 from smac.callback.metadata_callback import MetadataCallback
+from smac.callback.wandb_logging import WandBCallback
 
 __all__ = [
     "Callback",
     "MetadataCallback",
+    "WandBCallback",
 ]

smac/callback/wandb_logging.py

@@ -0,0 +1,62 @@
+from dataclasses import asdict
+
+import smac
+from smac import Callback
+from smac.runhistory import TrialInfo, TrialValue
+
+
+class WandBCallback(Callback):
+    def __init__(
+        self,
+        project: str,
+        entity: str,
+        id: str | None = None,
+        outdir: str | None = None,
+        mode: str | None = None,
+        resume: str = "allow",
+        job_type: str | None = None,
+        group: str | None = None,
+        config: dict | str | None = None,
+        save_code: bool = True,
+        **kwargs,
+    ) -> None:
+        import wandb
+
+        self.run = wandb.init(
+            id=id,
+            resume=resume,
+            mode=mode,
+            project=project,
+            job_type=job_type,
+            entity=entity,
+            group=group,
+            dir=outdir,
+            config=config,
+            save_code=save_code,
+            **kwargs,
+        )
+        super().__init__()
+
+    def on_tell_end(self, smbo: smac.main.smbo.SMBO, info: TrialInfo, value: TrialValue) -> bool | None:
+        info_dict = asdict(info)
+        info_dict["config"] = info_dict["config"].get_dictionary()
+        value_dict = asdict(value)
+        log_dict = info_dict | value_dict
+        log_dict["step"] = smbo.runhistory.finished
+        self.run.log(data=log_dict)
+        return super().on_tell_end(smbo, info, value)
+
+    def on_end(self, smbo: smac.main.smbo.SMBO) -> None:
+        intensifier_data = smbo.intensifier.get_data()
+        trajectory = intensifier_data["trajectory"]
+        import pandas as pd
+
+        df = pd.DataFrame(data=trajectory)
+        print(df)
+        # trajectory = Table(dataframe=df, allow_mixed_types=True)
+        df["costs"] = df["costs"].apply(lambda x: x[0])  # TODO properly log multi costs
+        for index, row in df.iterrows():
+            print(dict(row))
+            self.run.log(dict(row))
+        self.run.finish()
+        return super().on_end(smbo)