python-chebifier/chebifier/model_registry.py at e10e12981a9ed0ea572a5ff5a9b1449c0e3b5044 · ChEB-AI/python-chebifier · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from chebifier.ensemble.base_ensemble import BaseEnsemble
from chebifier.ensemble.weighted_majority_ensemble import (
    WMVwithF1Ensemble,
    WMVwithPPVNPVEnsemble,
)
from chebifier.prediction_models import (
    ChEBILookupPredictor,
    ChemlogPeptidesPredictor,
    ElectraPredictor,
    GNNPredictor,
)
from chebifier.prediction_models.c3p_predictor import C3PPredictor
from chebifier.prediction_models.chemlog_predictor import (
    ChemlogAllPredictor,
    ChemLogLopsterClingoPredictor,
    ChemlogLopsterPredictor,
    ChemlogOrganoXCompoundPredictor,
    ChemlogXMolecularEntityPredictor,
)

ENSEMBLES = {
    "mv": BaseEnsemble,
    "wmv-ppvnpv": WMVwithPPVNPVEnsemble,
    "wmv-f1": WMVwithF1Ensemble,
}


MODEL_TYPES = {
    "electra": ElectraPredictor,
    "resgated": GNNPredictor,
    "gat": GNNPredictor,
    "chemlog": ChemlogAllPredictor,
    "chemlog_peptides": ChemlogPeptidesPredictor,
    "chebi_lookup": ChEBILookupPredictor,
    "chemlog_element": ChemlogXMolecularEntityPredictor,
    "chemlog_organox": ChemlogOrganoXCompoundPredictor,
    "lopster": ChemlogLopsterPredictor,
    "lopster_clingo": ChemLogLopsterClingoPredictor,
    "c3p": C3PPredictor,
}


common_keys = MODEL_TYPES.keys() & ENSEMBLES.keys()
assert (
    not common_keys
), f"Overlapping keys between MODEL_TYPES and ENSEMBLES: {common_keys}"