split chemlog extra predictors into two for generalisability

Author: sfluegel05

chebifier/cli.py

@@ -3,7 +3,7 @@
 import click
 import yaml
 
-from .model_registry import ENSEMBLES
+from chebifier.model_registry import ENSEMBLES
 
 
 @click.group()

chebifier/model_registry.py

@@ -7,8 +7,9 @@
     ChemlogPeptidesPredictor,
     ElectraPredictor,
     ResGatedPredictor,
-    ChEBILookupPredictor, ChemlogByElementPredictor
+    ChEBILookupPredictor
 )
+from chebifier.prediction_models.chemlog_predictor import ChemlogXMolecularEntityPredictor, ChemlogOrganoXCompoundPredictor
 
 ENSEMBLES = {
     "mv": BaseEnsemble,
@@ -22,7 +23,8 @@
     "resgated": ResGatedPredictor,
     "chemlog_peptides": ChemlogPeptidesPredictor,
     "chebi_lookup": ChEBILookupPredictor,
-    "chemlog_element": ChemlogByElementPredictor
+    "chemlog_element": ChemlogXMolecularEntityPredictor,
+    "chemlog_organox": ChemlogOrganoXCompoundPredictor,
 }
 
 

chebifier/prediction_models/__init__.py

@@ -1,7 +1,7 @@
 from .base_predictor import BasePredictor
-from .chemlog_predictor import ChemlogPeptidesPredictor, ChemlogByElementPredictor
+from .chemlog_predictor import ChemlogPeptidesPredictor, ChemlogExtraPredictor
 from .electra_predictor import ElectraPredictor
 from .gnn_predictor import ResGatedPredictor
 from .chebi_lookup import ChEBILookupPredictor
 __all__ = ["BasePredictor", "ChemlogPeptidesPredictor", "ElectraPredictor", "ResGatedPredictor", "ChEBILookupPredictor",
-           "ChemlogByElementPredictor"]
+           "ChemlogExtraPredictor"]

chebifier/prediction_models/chemlog_predictor.py

@@ -39,21 +39,25 @@
     "Y": "L-tyrosine",
 }
 
-class ChemlogByElementPredictor(BasePredictor):
+class ChemlogExtraPredictor(BasePredictor):
+
+    CHEMLOG_CLASSIFIER = None
 
     def __init__(self, model_name: str, **kwargs):
         super().__init__(model_name, **kwargs)
-        self.x_molecular = XMolecularEntityClassifier()
-        self.organo_x = OrganoXCompoundClassifier()
+        self.classifier = self.CHEMLOG_CLASSIFIER()
 
     def predict_smiles_list(self, smiles_list: list[str]) -> list:
         mol_list = [_smiles_to_mol(smiles) for smiles in smiles_list]
-        return [
-            {str(cls): 1 for cls in self.x_molecular.classify(mol)[0] + self.organo_x.classify(mol)[0]}
-            if mol
-            else None
-            for mol in mol_list
-        ]
+        return self.classifier.classify(mol_list)
+
+class ChemlogXMolecularEntityPredictor(ChemlogExtraPredictor):
+
+    CHEMLOG_CLASSIFIER = XMolecularEntityClassifier
+
+class ChemlogOrganoXCompoundPredictor(ChemlogExtraPredictor):
+
+    CHEMLOG_CLASSIFIER = OrganoXCompoundClassifier
 
 class ChemlogPeptidesPredictor(BasePredictor):
     def __init__(self, model_name: str, **kwargs):