Implement Tox21 dataset as used in the competition

Author: MGlauer

chebai/cli.py

@@ -15,11 +15,12 @@ def cli():
 @click.argument("experiment")
 @click.argument("batch_size", type=click.INT)
 @click.option("-g", "--group", default="default")
+@click.option("--version", default=None)
 @click.argument("args", nargs=-1)
-def train(experiment, batch_size, group, args):
+def train(experiment, batch_size, group, version, args):
     """Run experiment identified by EXPERIMENT in batches of size BATCH_SIZE."""
     try:
-        ex = experiments.EXPERIMENTS[experiment](batch_size, group)
+        ex = experiments.EXPERIMENTS[experiment](batch_size, group, version=version)
     except KeyError:
         raise Exception(
             "Experiment ID not found. The following are available:"

chebai/experiments.py

@@ -25,8 +25,9 @@ def __init_subclass__(cls, **kwargs):
         if cls.identifier() is not None:
             EXPERIMENTS[cls.identifier()] = cls
 
-    def __init__(self, batch_size, *args, **kwargs):
+    def __init__(self, batch_size, *args, version=None, **kwargs):
         self.dataset = self.build_dataset(batch_size)
+        self.version=version
 
     @classmethod
     def identifier(cls) -> str:
@@ -46,6 +47,7 @@ def train(self, batch_size, *args):
             self.MODEL.NAME,
             loss=self.LOSS,
             model_kwargs=self.model_kwargs(*args),
+            version=self.version
         )
 
     def test(self, ckpt_path, *args):
@@ -264,28 +266,32 @@ def model_kwargs(self, *args) -> Dict:
         return d
 
 
-class ElectraOnTox21Bloat(ElectraOnTox21):
+class ElectraOnTox21MoleculeNet(_ElectraExperiment):
     MODEL = electra.Electra
     LOSS = torch.nn.BCEWithLogitsLoss
 
     @classmethod
     def identifier(cls) -> str:
-        return "Electra+Tox21Bloat"
+        return "Electra+Tox21MN"
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
-        return datasets.Tox21BloatChem(batch_size)
+        return datasets.Tox21MolNetChem(batch_size)
 
+    def model_kwargs(self, *args) -> Dict:
+        d = super().model_kwargs(*args)
+        d["config"]["hidden_dropout_prob"] = 0.4
+        d["config"]["word_dropout"] = 0.2
+        d["optimizer_kwargs"]["weight_decay"] = 1e-4
+        return d
 
-class ElectraOnTox21Ext(_ElectraExperiment):
-    MODEL = electra.Electra
-    LOSS = torch.nn.BCEWithLogitsLoss
 
+class ElectraOnTox21Challenge(_ElectraExperiment):
     @classmethod
     def identifier(cls) -> str:
-        return "Electra+Tox21Ext"
+        return "Electra+Tox21Chal"
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
-        return datasets.Tox21ExtendedChem(batch_size)
+        return datasets.Tox21ChallengeChem(batch_size)
 
 
 class ElectraBPEOnJCIExt(_ElectraExperiment):

chebai/models/base.py

@@ -146,6 +146,7 @@ def run(
         model_kwargs: dict = None,
         loss=torch.nn.BCELoss,
         weighted=False,
+        version=None
     ):
         if model_args is None:
             model_args = []
@@ -177,7 +178,9 @@ def run(
         else:
             trainer_kwargs = dict(gpus=0)
 
-        tb_logger = pl_loggers.TensorBoardLogger("logs/", name=name)
+        tb_logger = pl_loggers.TensorBoardLogger("logs/", name=name, version=version)
+        if os.path.isdir(tb_logger.log_dir):
+            raise IOError("Fixed logging directory does already exist:", tb_logger.log_dir)
         best_checkpoint_callback = ModelCheckpoint(
             dirpath=os.path.join(tb_logger.log_dir, "best_checkpoints"),
             filename="{epoch}-{val_F1Score_micro:.4f}--{val_loss:.4f}",

chebai/models/electra.py

@@ -151,10 +151,9 @@ def __init__(self, **kwargs):
         self.config = ElectraConfig(**kwargs["config"], output_attentions=True)
         self.word_dropout = nn.Dropout(kwargs["config"].get("word_dropout", 0))
         if pretrained_checkpoint:
-            elpre = ElectraPre.load_from_checkpoint(pretrained_checkpoint)
-            with TemporaryDirectory() as td:
-                elpre.as_pretrained.save_pretrained(td)
-                self.electra = ElectraModel.from_pretrained(td, config=self.config)
+            with open(pretrained_checkpoint, "rb") as fin:
+                model_dict = torch.load(fin,map_location=self.device)
+                self.electra = ElectraModel.from_pretrained(None, state_dict=model_dict['state_dict'], config=self.config)
         else:
             self.electra = ElectraModel(config=self.config)
 

chebai/preprocessing/datasets/tox21.py

@@ -13,12 +13,12 @@
 from chebai.preprocessing import reader as dr
 import pysmiles
 import numpy as np
-import rdkit
+from rdkit import Chem
 import zipfile
 import shutil
 
 
-class Tox21Base(XYBaseDataModule):
+class Tox21MolNet(XYBaseDataModule):
     HEADERS = [
         "NR-AR",
         "NR-AR-LBD",
@@ -36,7 +36,122 @@ class Tox21Base(XYBaseDataModule):
 
     @property
     def _name(self):
-        return "tox21"
+        return "Tox21mn"
+
+    @property
+    def label_number(self):
+        return 12
+
+    @property
+    def raw_file_names(self):
+        return ["tox21.csv"]
+
+    @property
+    def processed_file_names(self):
+        return ["test.pt", "train.pt", "validation.pt"]
+
+    def download(self):
+        with NamedTemporaryFile("rb") as gout:
+            request.urlretrieve(
+                "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz",
+                gout.name,
+            )
+            with gzip.open(gout.name) as gfile:
+                with open(os.path.join(self.raw_dir, "tox21.csv"), "wt") as fout:
+                    fout.write(gfile.read().decode())
+
+    def setup_processed(self):
+        print("Create splits")
+        data = self._load_data_from_file(os.path.join(self.raw_dir, f"tox21.csv"))
+        groups = np.array([d["group"] for d in data])
+        if not all(g is None for g in groups):
+            split_size = int(len(set(groups)) * self.train_split)
+            os.makedirs(self.processed_dir, exist_ok=True)
+            splitter = GroupShuffleSplit(train_size=split_size, n_splits=1)
+
+            train_split_index, temp_split_index = next(
+                splitter.split(data, groups=groups)
+            )
+
+            split_groups = groups[temp_split_index]
+
+            splitter = GroupShuffleSplit(
+                train_size=int(len(set(split_groups)) * self.train_split), n_splits=1
+            )
+            test_split_index, validation_split_index = next(
+                splitter.split(temp_split_index, groups=split_groups)
+            )
+            train_split = [data[i] for i in train_split_index]
+            test_split = [
+                d
+                for d in (data[temp_split_index[i]] for i in test_split_index)
+                if d["original"]
+            ]
+            validation_split = [
+                d
+                for d in (data[temp_split_index[i]] for i in validation_split_index)
+                if d["original"]
+            ]
+        else:
+            train_split, test_split = train_test_split(
+                data, train_size=self.train_split, shuffle=True
+            )
+            test_split, validation_split = train_test_split(
+                test_split, train_size=0.5, shuffle=True
+            )
+        for k, split in [
+            ("test", test_split),
+            ("train", train_split),
+            ("validation", validation_split),
+        ]:
+            print("transform", k)
+            torch.save(
+                split,
+                os.path.join(self.processed_dir, f"{k}.pt"),
+            )
+
+    def setup(self, **kwargs):
+        if any(
+            not os.path.isfile(os.path.join(self.raw_dir, f))
+            for f in self.raw_file_names
+        ):
+            self.download()
+        if any(
+            not os.path.isfile(os.path.join(self.processed_dir, f))
+            for f in self.processed_file_names
+        ):
+            self.setup_processed()
+
+    def _load_dict(self, input_file_path):
+        with open(input_file_path, "r") as input_file:
+            reader = csv.DictReader(input_file)
+            for row in reader:
+                smiles = row["smiles"]
+                labels = [
+                    bool(int(l)) if l else None for l in (row[k] for k in self.HEADERS)
+                ]
+                yield dict(features=smiles, labels=labels, ident=row["mol_id"])
+
+
+class Tox21Challenge(XYBaseDataModule):
+    HEADERS = [
+        "NR-AR",
+        "NR-AR-LBD",
+        "NR-AhR",
+        "NR-Aromatase",
+        "NR-ER",
+        "NR-ER-LBD",
+        "NR-PPAR-gamma",
+        "SR-ARE",
+        "SR-ATAD5",
+        "SR-HSE",
+        "SR-MMP",
+        "SR-p53",
+    ]
+
+    @property
+    def _name(self):
+        return "tox21chal"
 
     @property
     def label_number(self):
@@ -81,23 +196,24 @@ def _retrieve_file(self, url, target_file, compression=None):
                         shutil.move(os.path.join(td.name, f), target_path)
 
     def _load_data_from_file(self, path):
-        sdf = rdkit.Chem.SDMolSupplier(path)
+        sdf = Chem.SDMolSupplier(path)
         data = []
         for mol in sdf:
             if mol is not None:
                 d = dict(
                     labels=[int(mol.GetProp(h)) if h in mol.GetPropNames() else None for h in self.HEADERS],
                     ident=[mol.GetProp(k) for k in ("DSSTox_CID", "Compound ID") if k in mol.GetPropNames() ][0],
-                    features=rdkit.Chem.MolToSmiles(mol))
+                    features=Chem.MolToSmiles(mol))
                 data.append(self.reader.to_data(d))
         return data
 
     def setup_processed(self):
         for k in ("train", "validation"):
-            torch.save(self._load_data_from_file(os.path.join(self.raw_dir, f"{k}.sdf")), os.path.join(self.processed_dir, f"{k}.pt"))
+            d = self._load_data_from_file(os.path.join(self.raw_dir, f"{k}.sdf"))
+            torch.save(d, os.path.join(self.processed_dir, f"{k}.pt"))
 
         with open(os.path.join(self.raw_dir, f"test.smiles")) as fin:
-            headers = next(fin)
+            next(fin)
             test_smiles = dict(reversed(row.strip().split("\t")) for row in fin)
         with open(os.path.join(self.raw_dir, f"test_results.txt")) as fin:
             headers = next(fin).strip().split("\t")
@@ -128,35 +244,9 @@ def _load_dict(self, input_file_path):
                 yield dict(features=smiles, labels=labels, ident=row["mol_id"])
 
 
-class Tox21Chem(Tox21Base):
+class Tox21ChallengeChem(Tox21Challenge):
     READER = dr.ChemDataReader
 
 
-class Tox21Graph(Tox21Base):
-    READER = dr.GraphReader
-
-
-
-class Tox21ExtendedChem(MergedDataset):
-    MERGED = [Tox21Chem, Hazardous, JCIExtendedTokenData]
-
-    @property
-    def limits(self):
-        return [None, 5000, 5000]
-
-    def _process_data(self, subset_id, data):
-        res = dict(
-            features=data["features"], labels=data["labels"], ident=data["ident"]
-        )
-        # Feature: non-toxic
-        if subset_id == 0:
-            res["labels"] = [not any(res["labels"])]
-        elif subset_id == 1:
-            res["labels"] = [False]
-        elif subset_id == 2:
-            res["labels"] = [True]
-        return res
-
-    @property
-    def label_number(self):
-        return 1
+class Tox21MolNetChem(Tox21MolNet):
+    READER = dr.ChemDataReader