Making PFN conditioned on local availability

automl · Nov 3, 2023 · 7eebf09 · 7eebf09
1 parent 17ef717
commit 7eebf09
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/neps/optimizers/bayesian_optimization/acquisition_samplers/freeze_thaw_sampler.py b/src/neps/optimizers/bayesian_optimization/acquisition_samplers/freeze_thaw_sampler.py
@@ -13,7 +13,7 @@
 
 
 class FreezeThawSampler(AcquisitionSampler):
-    SAMPLES_TO_DRAW = 10  # number of random samples to draw at lowest fidelity
+    SAMPLES_TO_DRAW = 3  # number of random samples to draw at lowest fidelity
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)

diff --git a/src/neps/optimizers/bayesian_optimization/models/__init__.py b/src/neps/optimizers/bayesian_optimization/models/__init__.py
@@ -1,11 +1,16 @@
 from .deepGP import DeepGP
 from .gp import ComprehensiveGP
 from .gp_hierarchy import ComprehensiveGPHierarchy
-from .pfn import PFN_SURROGATE
+
 
 SurrogateModelMapping = {
     "deep_gp": DeepGP,
     "gp": ComprehensiveGP,
     "gp_hierarchy": ComprehensiveGPHierarchy,
-    "pfn": PFN_SURROGATE
 }
+
+try:
+    from .pfn import PFN_SURROGATE  # only if available locally
+    SurrogateModelMapping.update({"pfn": PFN_SURROGATE})
+except:
+    pass