Merge pull request #14 from ChEB-AI/feature/remove_chebai_lookup_dependency

use chebifier graph instead of chebai graph

Author: sfluegel05

chebifier/prediction_models/chebi_lookup.py

@@ -6,6 +6,7 @@
 import networkx as nx
 from rdkit import Chem
 import json
+from chebifier.utils import load_chebi_graph
 
 
 class ChEBILookupPredictor(BasePredictor):
@@ -23,15 +24,7 @@ def __init__(
             or "ChEBI Lookup: If the SMILES is equivalent to a ChEBI entry, retrieve the classification of that entry."
         )
         self.chebi_version = chebi_version
-        self.chebi_graph = kwargs.get("chebi_graph", None)
-        if self.chebi_graph is None:
-            from chebai.preprocessing.datasets.chebi import ChEBIOver50
-
-            self.chebi_dataset = ChEBIOver50(chebi_version=self.chebi_version)
-            self.chebi_dataset._download_required_data()
-            self.chebi_graph = self.chebi_dataset._extract_class_hierarchy(
-                os.path.join(self.chebi_dataset.raw_dir, "chebi.obo")
-            )
+        self.chebi_graph = kwargs.get("chebi_graph", load_chebi_graph())
         self.lookup_table = self.get_smiles_lookup()
 
     def get_smiles_lookup(self):