Skip to content

Commit 4cc3e9e

Browse files
committed
use chebifier graph instead of chebai graph
1 parent 955ab83 commit 4cc3e9e

File tree

1 file changed

+2
-10
lines changed

1 file changed

+2
-10
lines changed

chebifier/prediction_models/chebi_lookup.py

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import networkx as nx
77
from rdkit import Chem
88
import json
9-
9+
from chebifier.utils import load_chebi_graph
1010

1111
class ChEBILookupPredictor(BasePredictor):
1212

@@ -23,15 +23,7 @@ def __init__(
2323
or "ChEBI Lookup: If the SMILES is equivalent to a ChEBI entry, retrieve the classification of that entry."
2424
)
2525
self.chebi_version = chebi_version
26-
self.chebi_graph = kwargs.get("chebi_graph", None)
27-
if self.chebi_graph is None:
28-
from chebai.preprocessing.datasets.chebi import ChEBIOver50
29-
30-
self.chebi_dataset = ChEBIOver50(chebi_version=self.chebi_version)
31-
self.chebi_dataset._download_required_data()
32-
self.chebi_graph = self.chebi_dataset._extract_class_hierarchy(
33-
os.path.join(self.chebi_dataset.raw_dir, "chebi.obo")
34-
)
26+
self.chebi_graph = kwargs.get("chebi_graph", load_chebi_graph())
3527
self.lookup_table = self.get_smiles_lookup()
3628

3729
def get_smiles_lookup(self):

0 commit comments

Comments
 (0)