Simplify getting started

This PR updates the BaseEnsemble constructor to allow the following:

1. Passing a string or path to the configuration
2. Not passing a configuration at all, which will automatically load the default configuration. This is now the default, since most users won't want to have to configure it (it should have reasonable defaults)

Author: cthoyt

README.md

@@ -75,15 +75,11 @@ python -m chebifier predict --help
 You can also use the package programmatically:
 
 ```python
-from chebifier.ensemble.base_ensemble import BaseEnsemble
-import yaml
+from chebifier import BaseEnsemble
 
-# Load configuration from YAML file
-with open('configs/example_config.yml', 'r') as f:
-    config = yaml.safe_load(f)
-
-# Instantiate ensemble model
-ensemble = BaseEnsemble(config)
+# Instantiate ensemble model. If desired, can pass
+# a path to a configuration, like 'configs/example_config.yml'
+ensemble = BaseEnsemble()
 
 # Make predictions
 smiles_list = ["CC(=O)OC1=CC=CC=C1C(=O)O", "C1=CC=C(C=C1)C(=O)O"]

chebifier/__init__.py

@@ -2,5 +2,10 @@
 # even if multiple subpackages are imported later.
 
 from ._custom_cache import PerSmilesPerModelLRUCache
+from chebifier.ensemble.base_ensemble import BaseEnsemble
+
+__all__ = [
+    "BaseEnsemble",
+]
 
 modelwise_smiles_lru_cache = PerSmilesPerModelLRUCache(max_size=100)

chebifier/cli.py

@@ -4,6 +4,7 @@
 import yaml
 
 from chebifier.model_registry import ENSEMBLES
+from chebifier.utils import get_default_configs
 
 
 @click.group()
@@ -75,12 +76,7 @@ def predict(
     # Load configuration from YAML file
     if not ensemble_config:
         print("Using default ensemble configuration")
-        with (
-            importlib.resources.files("chebifier")
-            .joinpath("ensemble.yml")
-            .open("r") as f
-        ):
-            config = yaml.safe_load(f)
+        config = get_default_configs()
     else:
         print(f"Loading ensemble configuration from {ensemble_config}")
         with open(ensemble_config, "r") as f:

chebifier/ensemble/base_ensemble.py

@@ -1,23 +1,33 @@
 import os
 import time
+from pathlib import Path
+from typing import Union
 
 import torch
 import tqdm
+import yaml
 
 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
+from chebifier.utils import get_disjoint_files, load_chebi_graph, get_default_configs
 
 
 class BaseEnsemble:
     def __init__(
         self,
-        model_configs: dict,
+        model_configs: Union[str, Path, dict, None] = None,
         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
 

chebifier/utils.py

@@ -1,8 +1,11 @@
+import importlib.resources
 import os
 
 import networkx as nx
 import requests
 import fastobo
+import yaml
+
 from chebifier.hugging_face import download_model_files
 import pickle
 
@@ -129,3 +132,12 @@ def get_disjoint_files():
     # pickle.dump(chebi_graph, open("chebi_graph.pkl", "wb"))
     chebi_graph = load_chebi_graph()
     print(chebi_graph)
+
+
+def get_default_configs():
+    with (
+        importlib.resources.files("chebifier")
+            .joinpath("ensemble.yml")
+            .open("r") as f
+    ):
+        return yaml.safe_load(f)