add single chemlog model

Author: sfluegel05

chebifier/model_registry.py

@@ -11,6 +11,7 @@
 )
 from chebifier.prediction_models.c3p_predictor import C3PPredictor
 from chebifier.prediction_models.chemlog_predictor import (
+    ChemlogAllPredictor,
     ChemlogOrganoXCompoundPredictor,
     ChemlogXMolecularEntityPredictor,
 )
@@ -27,6 +28,7 @@
     "electra": ElectraPredictor,
     "resgated": ResGatedPredictor,
     "gat": GATPredictor,
+    "chemlog": ChemlogAllPredictor,
     "chemlog_peptides": ChemlogPeptidesPredictor,
     "chebi_lookup": ChEBILookupPredictor,
     "chemlog_element": ChemlogXMolecularEntityPredictor,

chebifier/prediction_models/chemlog_predictor.py

@@ -32,6 +32,31 @@
 }
 
 
+class ChemlogAllPredictor(BasePredictor):
+    def __init__(self, model_name: str, **kwargs):
+        super().__init__(model_name, **kwargs)
+        self.chebi_graph = kwargs.get("chebi_graph", None)
+        self.predictors = [
+            ChemlogXMolecularEntityPredictor("chemlog_x_molecular_entity", **kwargs),
+            ChemlogOrganoXCompoundPredictor("chemlog_organo_x_compound", **kwargs),
+            ChemlogPeptidesPredictor("chemlog_peptides", **kwargs),
+        ]
+
+    @modelwise_smiles_lru_cache.batch_decorator
+    def predict_smiles_list(self, smiles_list: list[str]) -> list:
+        results = []
+        for predictor in self.predictors:
+            predictor_results = predictor._predict_smiles_list(smiles_list)
+            for i, res in enumerate(predictor_results):
+                if i >= len(results):
+                    results.append(dict())
+                results[i].update(res)
+        return results
+
+    def explain_smiles(self, smiles):
+        return self.predictors[2].explain_smiles(smiles)
+
+
 class ChemlogExtraPredictor(BasePredictor):
 
     def __init__(self, model_name: str, **kwargs):
@@ -41,6 +66,9 @@ def __init__(self, model_name: str, **kwargs):
 
     @modelwise_smiles_lru_cache.batch_decorator
     def predict_smiles_list(self, smiles_list: list[str]) -> list:
+        return self._predict_smiles_list(smiles_list)
+
+    def _predict_smiles_list(self, smiles_list: list[str]) -> list:
         from chemlog.cli import _smiles_to_mol
 
         mol_list = [_smiles_to_mol(smiles) for smiles in smiles_list]
@@ -124,6 +152,9 @@ def predict_smiles(self, smiles: str) -> Optional[dict]:
 
     @modelwise_smiles_lru_cache.batch_decorator
     def predict_smiles_list(self, smiles_list: list[str]) -> list:
+        return self._predict_smiles_list(smiles_list)
+
+    def _predict_smiles_list(self, smiles_list: list[str]) -> list:
         results = []
         for i, smiles in tqdm.tqdm(enumerate(smiles_list)):
             results.append(self.predict_smiles(smiles))