move config processing logic to base ensemble

Author: sfluegel05

chebifier/cli.py

@@ -4,8 +4,6 @@
 import yaml
 
 from chebifier.model_registry import ENSEMBLES
-from chebifier.utils import get_default_configs
-
 
 @click.group()
 def cli():
@@ -73,38 +71,10 @@ def predict(
     resolve_inconsistencies=True,
 ):
     """Predict ChEBI classes for SMILES strings using an ensemble model."""
-    # Load configuration from YAML file
-    if not ensemble_config:
-        print("Using default ensemble configuration")
-        config = get_default_configs()
-    else:
-        print(f"Loading ensemble configuration from {ensemble_config}")
-        with open(ensemble_config, "r") as f:
-            config = yaml.safe_load(f)
-
-    with (
-        importlib.resources.files("chebifier")
-        .joinpath("model_registry.yml")
-        .open("r") as f
-    ):
-        model_registry = yaml.safe_load(f)
-
-    new_config = {}
-    for model_name, entry in config.items():
-        if "load_model" in entry:
-            if entry["load_model"] not in model_registry:
-                raise ValueError(
-                    f"Model {entry['load_model']} not found in model registry. "
-                    f"Available models are: {','.join(model_registry.keys())}."
-                )
-            new_config[model_name] = {**model_registry[entry["load_model"]], **entry}
-        else:
-            new_config[model_name] = entry
-    config = new_config
 
     # Instantiate ensemble model
     ensemble = ENSEMBLES[ensemble_type](
-        config,
+        ensemble_config,
         chebi_version=chebi_version,
         resolve_inconsistencies=resolve_inconsistencies,
     )

chebifier/ensemble/base_ensemble.py

@@ -6,12 +6,13 @@
 import torch
 import tqdm
 import yaml
+import importlib
 
 from chebifier.check_env import check_package_installed
 from chebifier.hugging_face import download_model_files
 from chebifier.inconsistency_resolution import PredictionSmoother
 from chebifier.prediction_models.base_predictor import BasePredictor
-from chebifier.utils import get_disjoint_files, load_chebi_graph, get_default_configs
+from chebifier.utils import get_disjoint_files, load_chebi_graph, get_default_configs, process_config
 
 
 class BaseEnsemble:
@@ -21,22 +22,34 @@ def __init__(
         chebi_version: int = 241,
         resolve_inconsistencies: bool = True,
     ):
-        if model_configs is None:
-            model_configs = get_default_configs()
-        elif isinstance(model_configs, (str, Path)):
-            # Load configuration from YAML file
-            with open(model_configs) as file:
-                model_configs = yaml.safe_load(file)
-
         # Deferred Import: To avoid circular import error
         from chebifier.model_registry import MODEL_TYPES
 
+        # Load configuration from YAML file
+        if not model_configs:
+            config = get_default_configs()
+        elif isinstance(model_configs, dict):
+            config = model_configs
+        else:
+            print(f"Loading ensemble configuration from {model_configs}")
+            with open(model_configs, "r") as f:
+                config = yaml.safe_load(f)
+
+        with (
+            importlib.resources.files("chebifier")
+                    .joinpath("model_registry.yml")
+                    .open("r") as f
+        ):
+            model_registry = yaml.safe_load(f)
+
+        processed_configs = process_config(config, model_registry)
+
         self.chebi_graph = load_chebi_graph()
         self.disjoint_files = get_disjoint_files()
 
         self.models = []
         self.positive_prediction_threshold = 0.5
-        for model_name, model_config in model_configs.items():
+        for model_name, model_config in processed_configs.items():
             model_cls = MODEL_TYPES[model_config["type"]]
             if "hugging_face" in model_config:
                 hugging_face_kwargs = download_model_files(model_config["hugging_face"])

chebifier/utils.py

@@ -126,18 +126,27 @@ def get_disjoint_files():
     return disjoint_files
 
 
-if __name__ == "__main__":
-    # chebi_graph = build_chebi_graph(chebi_version=241)
-    # save the graph to a file
-    # pickle.dump(chebi_graph, open("chebi_graph.pkl", "wb"))
-    chebi_graph = load_chebi_graph()
-    print(chebi_graph)
-
-
 def get_default_configs():
+    default_config_name = "ensemble.yml"
+    print(f"Using default ensemble configuration from {default_config_name}")
     with (
         importlib.resources.files("chebifier")
-            .joinpath("ensemble.yml")
+            .joinpath(default_config_name)
             .open("r") as f
     ):
         return yaml.safe_load(f)
+
+
+def process_config(config, model_registry):
+    new_config = {}
+    for model_name, entry in config.items():
+        if "load_model" in entry:
+            if entry["load_model"] not in model_registry:
+                raise ValueError(
+                    f"Model {entry['load_model']} not found in model registry. "
+                    f"Available models are: {','.join(model_registry.keys())}."
+                )
+            new_config[model_name] = {**model_registry[entry["load_model"]], **entry}
+        else:
+            new_config[model_name] = entry
+    return new_config