add cache

Author: sfluegel05

chebifier/ensemble/base_ensemble.py

@@ -7,7 +7,7 @@
 from chebai.result.analyse_sem import PredictionSmoother, get_chebi_graph
 
 from chebifier.prediction_models.base_predictor import BasePredictor
-
+from functools import lru_cache
 
 class BaseEnsemble:
 

chebifier/prediction_models/base_predictor.py

@@ -1,6 +1,8 @@
 import json
 from abc import ABC
 
+from functools import lru_cache
+
 
 class BasePredictor(ABC):
     def __init__(
@@ -22,7 +24,16 @@ def __init__(
         self._description = kwargs.get("description", None)
 
     def predict_smiles_list(self, smiles_list: list[str]) -> dict:
-        raise NotImplementedError
+        # list is not hashable, so we convert it to a tuple (useful for caching)
+        return self.predict_smiles_tuple(tuple(smiles_list))
+
+    @lru_cache(maxsize=100)
+    def predict_smiles_tuple(self, smiles_tuple: tuple[str]) -> dict:
+        raise NotImplementedError()
+
+    def predict_smiles(self, smiles: str) -> dict:
+        # by default, use list-based prediction
+        return self.predict_smiles_tuple((smiles,))[0]
 
     @property
     def info_text(self):

chebifier/prediction_models/c3p_predictor.py

@@ -1,3 +1,4 @@
+from functools import lru_cache
 from typing import Optional, List
 from pathlib import Path
 
@@ -17,8 +18,9 @@ def __init__(self, model_name: str, program_directory: Optional[Path]=None, chem
         self.chemical_classes = chemical_classes
         self.chebi_graph = kwargs.get("chebi_graph", None)
 
-    def predict_smiles_list(self, smiles_list: list[str]) -> list:
-        result_list = c3p_classifier.classify(smiles_list, self.program_directory, self.chemical_classes, strict=False)
+    @lru_cache(maxsize=100)
+    def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
+        result_list = c3p_classifier.classify(list(smiles_list), self.program_directory, self.chemical_classes, strict=False)
         result_reformatted = [dict() for _ in range(len(smiles_list))]
         for result in result_list:
             chebi_id = result.class_id.split(":")[1]

chebifier/prediction_models/chebi_lookup.py

@@ -1,3 +1,6 @@
+from functools import lru_cache
+from typing import Optional
+
 from chebifier.prediction_models import BasePredictor
 import os
 import networkx as nx
@@ -51,32 +54,36 @@ def build_smiles_lookup(self):
                     print(f"Failed to parse SMILES {smiles} for ChEBI ID {chebi_id}: {e}")
         return smiles_lookup
 
+    @lru_cache(maxsize=100)
+    def predict_smiles(self, smiles: str) -> Optional[dict]:
+        if not smiles:
+            return None
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+        canonical_smiles = Chem.MolToSmiles(mol)
+        if canonical_smiles in self.lookup_table:
+            parent_candidates = self.lookup_table[canonical_smiles]
+            preds_i = dict()
+            if len(parent_candidates) > 1:
+                print(
+                    f"Multiple matches found in ChEBI for SMILES {smiles}: {', '.join(str(chebi_id) for chebi_id, _ in parent_candidates)}")
+                for k in list(set(pp for _, p in parent_candidates for pp in p)):
+                    preds_i[str(k)] = 1
+            elif len(parent_candidates) == 1:
+                chebi_id, parents = parent_candidates[0]
+                for k in parents:
+                    preds_i[str(k)] = 1
+            else:
+                preds_i = None
+            return preds_i
+        else:
+            return None
 
-    def predict_smiles_list(self, smiles_list: list[str]) -> list:
+    def predict_smiles_tuple(self, smiles_list: list[str]) -> list:
         predictions = []
         for smiles in smiles_list:
-            if not smiles:
-                predictions.append(None)
-                continue
-            mol = Chem.MolFromSmiles(smiles)
-            if mol is None:
-                predictions.append(None)
-                continue
-            canonical_smiles = Chem.MolToSmiles(mol)
-            if canonical_smiles in self.lookup_table:
-                parent_candidates = self.lookup_table[canonical_smiles]
-                preds_i = dict()
-                if len(parent_candidates) > 1:
-                    print(f"Multiple matches found in ChEBI for SMILES {smiles}: {', '.join(str(chebi_id) for chebi_id, _ in parent_candidates)}")
-                    for k in list(set(pp for _, p in parent_candidates for pp in p)):
-                        preds_i[str(k)] = 1
-                elif len(parent_candidates) == 1:
-                    chebi_id, parents = parent_candidates[0]
-                    for k in parents:
-                        preds_i[str(k)] = 1
-                else:
-                    preds_i = None
-                predictions.append(preds_i)
+            predictions.append(self.predict_smiles(smiles))
 
         return predictions
 

chebifier/prediction_models/chemlog_predictor.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import tqdm
 from chemlog.alg_classification.charge_classifier import get_charge_category
 from chemlog.alg_classification.peptide_size_classifier import get_n_amino_acid_residues
@@ -10,6 +12,7 @@
 )
 from chemlog.cli import CLASSIFIERS, _smiles_to_mol, strategy_call
 from chemlog_extra.alg_classification.by_element_classification import XMolecularEntityClassifier, OrganoXCompoundClassifier
+from functools import lru_cache
 
 from .base_predictor import BasePredictor
 
@@ -48,7 +51,7 @@ def __init__(self, model_name: str, **kwargs):
         self.chebi_graph = kwargs.get("chebi_graph", None)
         self.classifier = self.CHEMLOG_CLASSIFIER()
 
-    def predict_smiles_list(self, smiles_list: list[str]) -> list:
+    def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
         mol_list = [_smiles_to_mol(smiles) for smiles in smiles_list]
         res = self.classifier.classify(mol_list)
         if self.chebi_graph is not None:
@@ -88,30 +91,32 @@ def __init__(self, model_name: str, **kwargs):
         # fmt: on
         print(f"Initialised ChemLog model {self.model_name}")
 
-    def predict_smiles_list(self, smiles_list: list[str]) -> list:
+    @lru_cache(maxsize=100)
+    def predict_smiles(self, smiles: str) -> Optional[dict]:
+        mol = _smiles_to_mol(smiles)
+        if mol is None:
+            return None
+        pos_labels = [label for label in self.peptide_labels if label in strategy_call(
+            self.strategy, self.classifier_instances, mol
+        )["chebi_classes"]]
+        if self.chebi_graph:
+            indirect_pos_labels = [str(pr) for label in pos_labels for pr in
+                                   self.chebi_graph.predecessors(int(label))]
+            pos_labels = list(set(pos_labels + indirect_pos_labels))
+        return {
+                label: (
+                    1
+                    if label
+                       in pos_labels
+                    else 0
+                )
+                for label in self.peptide_labels + pos_labels
+            }
+
+    def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
         results = []
         for i, smiles in tqdm.tqdm(enumerate(smiles_list)):
-            mol = _smiles_to_mol(smiles)
-            if mol is None:
-                results.append(None)
-            else:
-                pos_labels = [label for label in self.peptide_labels if label in strategy_call(
-                                self.strategy, self.classifier_instances, mol
-                            )["chebi_classes"]]
-                if self.chebi_graph:
-                    indirect_pos_labels = [str(pr) for label in pos_labels for pr in self.chebi_graph.predecessors(int(label))]
-                    pos_labels = list(set(pos_labels + indirect_pos_labels))
-                results.append(
-                    {
-                        label: (
-                            1
-                            if label
-                            in pos_labels
-                            else 0
-                        )
-                        for label in self.peptide_labels + pos_labels
-                    }
-                )
+            results.append(self.predict_smiles(smiles))
 
         for classifier in self.classifier_instances.values():
             classifier.on_finish()

chebifier/prediction_models/nn_predictor.py

@@ -1,3 +1,5 @@
+from functools import lru_cache
+
 import numpy as np
 import torch
 import tqdm
@@ -50,7 +52,8 @@ def read_smiles(self, smiles):
         d = reader.to_data(dict(features=smiles, labels=None))
         return d
 
-    def predict_smiles_list(self, smiles_list) -> list:
+    @lru_cache(maxsize=100)
+    def predict_smiles_tuple(self, smiles_list: tuple[str]) -> list:
         """Returns a list with the length of smiles_list, each element is either None (=failure) or a dictionary
         Of classes and predicted values."""
         token_dicts = []