from data.load_data import SupportedDatasets
from data.load_multi_data import SupportedMultiDatasets
from train.rdkit_compare import rdkit_conf, rdkit_multi_conf


# rdkit_conf(
#     dataset_name=SupportedDatasets.QM9,
#     dataset_token='phi-psi',
#     force_save=False,
#     use_tqdm=True,
#     use_cuda=True
# )
# rdkit_multi_conf(
#     dataset_name=SupportedMultiDatasets.GEOM_QM9_SMALL,
#     dataset_token='phi-psi',
#     force_save=False,
#     use_tqdm=True,
#     use_cuda=True
# )
rdkit_multi_conf(
    dataset_name=SupportedMultiDatasets.GEOM_QM9,
    dataset_token='phi-psi',
    force_save=False,
    use_tqdm=False,
    use_cuda=True
)
# rdkit_multi_conf(
#     dataset_name=SupportedMultiDatasets.GEOM_DRUGS_SMALL,
#     dataset_token='phi-psi',
#     force_save=False,
#     use_tqdm=True,
#     use_cuda=True
# )
rdkit_multi_conf(
    dataset_name=SupportedMultiDatasets.GEOM_DRUGS,
    dataset_token='phi-psi',
    force_save=False,
    use_tqdm=False,
    use_cuda=True
)
