Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/datasets/moleculenet/qm9/qm9.csv filter=lfs diff=lfs merge=lfs -text
+data/datasets/moleculenet/qm9/train.csv filter=lfs diff=lfs merge=lfs -text
+images/3dgridvqgan_architecture.png filter=lfs diff=lfs merge=lfs -text

data/3d_grids_sample/00b781f5a45a2cd1cdff1a582f5650f9.smi

@@ -0,0 +1 @@
+O=CNC1=NOC=N1

data/3d_grids_sample/00b781f5a45a2cd1cdff1a582f5650f9_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a17ceb4553e1861d34ebaf0d0b2e5f33975f41632eef13f662ed6098e06eb35
+size 338816

data/3d_grids_sample/00b781f5a45a2cd1cdff1a582f5650f9_0.xyz

@@ -0,0 +1,11 @@
+O -0.3433117873 0.0542505386 0.0030523182
+C -0.0862462134 1.2279140618 0.0081946405
+N 1.1846004685 1.7722542978 -0.0009450675
+C 2.386606203 1.0948306964 -0.0156497688
+N 2.5353255646 -0.2087260828 -0.0269119898
+O 3.9407365203 -0.3429161461 -0.0391639287
+C 4.4419993701 0.8855820068 -0.0338427395
+N 3.5551776308 1.8291029217 -0.0192066804
+H -0.8486325877 2.0303875837 0.0201485951
+H 1.2766287385 2.7775349452 0.0053395306
+H 5.5174865125 0.985073027 -0.0416438097

data/3d_grids_sample/0a18e0f64cbaf1508b32834ece70933c.smi

@@ -0,0 +1 @@
+C#CCC1CC1(C)CO

data/3d_grids_sample/0a18e0f64cbaf1508b32834ece70933c_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c211d6096850f4fd8bb7e862cb0e3cf09492ed5fb165348829ac37ed9d3c344a
+size 2612864

data/3d_grids_sample/0a18e0f64cbaf1508b32834ece70933c_0.xyz

@@ -0,0 +1,21 @@
+C 0.0095240887 1.5961806115 -0.0858857917
+C -0.0675156706 0.086549314 0.0511974473
+C -1.3201762415 -0.5125406846 -0.5508961783
+O -1.3261797579 -0.2420461247 -1.9477734763
+C 1.1982985053 -0.7358014855 -0.0482283446
+C 0.4787861145 -0.6243589949 1.2724834132
+C 1.0934534769 0.1242004421 2.4490763186
+C 1.9615545513 -0.7265702717 3.2588924687
+C 2.6652427105 -1.4441035487 3.9181710713
+H 0.9721361795 1.9954968448 0.2425890448
+H -0.1285829928 1.8750984225 -1.1338899379
+H -0.7769315564 2.0893384777 0.4988274146
+H -1.3405285423 -1.5971735621 -0.3600593089
+H -2.2081101606 -0.0754460014 -0.0643042439
+H -2.18303183 -0.5001196615 -2.2994922363
+H 1.1667370542 -1.6792797399 -0.5822025643
+H 2.1360390085 -0.1991832668 -0.1544507441
+H -0.0724271026 -1.5082395772 1.583953566
+H 1.6685961108 0.9854383459 2.091700259
+H 0.2979810602 0.5352503178 3.0854850798
+H 3.2930319243 -2.0701853274 4.5019859333

data/3d_grids_sample/0aa6c786f42a0d43d8d1dfc7e9ae4939.smi

@@ -0,0 +1 @@
+CCC1C=C2CCC2O1

data/3d_grids_sample/0aa6c786f42a0d43d8d1dfc7e9ae4939_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5adbb7b36ff91928b2ecec9e548962d45c5241d3949d35e5274511b0af7aa574
+size 2471168

data/3d_grids_sample/0aa6c786f42a0d43d8d1dfc7e9ae4939_0.xyz

@@ -0,0 +1,21 @@
+C -0.0055109461 1.5153095531 -0.1082270888
+C -0.0178026585 -0.0131266806 -0.0290282794
+C -1.426157367 -0.5876103682 0.1341883634
+O -2.2512193772 -0.169709702 -0.9848466475
+C -3.0265062489 -1.2762275569 -1.3904613892
+C -2.8104411225 -2.0637374738 -2.7387249862
+C -2.7598726311 -3.3919927241 -1.8878534159
+C -2.4407002514 -2.4767203713 -0.7256518392
+C -1.467181607 -2.1129314787 0.098813579
+H 1.012088138 1.8942209817 -0.2429905156
+H -0.615905307 1.8589202784 -0.9470659107
+H -0.4113398154 1.9616853934 0.8065539523
+H 0.4283203565 -0.4428051608 -0.9340972166
+H 0.5897595853 -0.3532294279 0.8196759242
+H -1.8681962495 -0.175448141 1.0583687871
+H -4.0959419742 -1.0853286499 -1.2201143184
+H -3.5930759641 -1.9863512547 -3.4970500576
+H -1.8354825793 -1.8302996855 -3.1700309799
+H -3.7465941351 -3.8605591083 -1.8233285422
+H -2.0210691777 -4.1500431884 -2.1526779895
+H -0.7049926879 -2.7257531143 0.5639707007

data/3d_grids_sample/0ae6f92c43122633a151eff5089e8da6.smi

@@ -0,0 +1 @@
+CN(C=O)C(N)=O

data/3d_grids_sample/0ae6f92c43122633a151eff5089e8da6_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37286a82f528eedeca6eb4b02d5a487fb80526fac79b8d4c16d374ee389aa2f4
+size 933248

data/3d_grids_sample/0ae6f92c43122633a151eff5089e8da6_0.xyz

@@ -0,0 +1,13 @@
+C -0.0561551802 1.4880114224 0.0115672899
+N -0.012314902 0.0308990979 -0.0074242366
+C 0.0394071908 -0.6483703608 1.2066473859
+O 0.0604983378 -0.0950939226 2.2806600492
+C -0.042694924 -0.7260819108 -1.1975216915
+N -0.0688282818 0.0226384605 -2.3461653613
+O -0.0805938084 -1.9390867198 -1.2010568454
+H 0.8560644678 1.9243363476 -0.4123282977
+H -0.9270762996 1.8594106628 -0.5357636429
+H -0.1279995908 1.7881285544 1.0569523777
+H 0.0576097412 -1.7352809462 1.059547977
+H 0.2506516948 0.9756066207 -2.3628772038
+H 0.0545094644 -0.513342856 -3.1899679505

data/3d_grids_sample/0b4020095a14325f0f174bc8a43f625d.smi

@@ -0,0 +1 @@
+CCCC(CC)C(C)=O

data/3d_grids_sample/0b4020095a14325f0f174bc8a43f625d_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:daf1f661f359abd9637a71410cfbd09e692daa396b32e047c6606e03874c317a
+size 2903168

data/3d_grids_sample/0b4020095a14325f0f174bc8a43f625d_0.xyz

@@ -0,0 +1,25 @@
+C 0.3762415949 1.4903515185 0.051015763
+C 0.16294672 -0.0262712548 0.0322514822
+C 0.675840727 -0.6751333193 -1.2593347251
+C 0.5134796126 -2.2163401834 -1.316721532
+C -0.9478193302 -2.6564989579 -1.4944772157
+C -1.1489111104 -4.1684742305 -1.3528914865
+C 1.3845485491 -2.7207559592 -2.4735115772
+C 2.8132544677 -3.1065708435 -2.1332551236
+O 0.960203994 -2.7885031108 -3.6052090651
+H 1.4383754499 1.7429425026 -0.0425294685
+H -0.1517461232 1.9744952727 -0.7778578838
+H 0.0095987923 1.9321754361 0.9827049486
+H 0.6755907641 -0.4799924509 0.890858941
+H -0.9033038021 -0.2426745921 0.1645631834
+H 0.1643795151 -0.2410453696 -2.127887566
+H 1.7371611953 -0.4154116304 -1.3751070827
+H 0.9136495565 -2.6311564942 -0.3805859669
+H -1.5741060384 -2.1375358991 -0.7612854955
+H -1.2777225554 -2.3336553985 -2.4874152623
+H -0.5803069276 -4.7154607671 -2.1113564713
+H -0.8279135541 -4.5236840194 -0.3667438252
+H -2.2026624643 -4.437369643 -1.4740466355
+H 2.8024629719 -4.041189801 -1.5589057987
+H 3.3927775837 -3.2531585575 -3.0456419515
+H 3.2903552816 -2.3529525778 -1.4982282451

data/3d_grids_sample/0b77c16e04a8ac8f84ae9ccaf9b1aaa0.smi

@@ -0,0 +1 @@
+C1=NC2=CC=NN2C=N1

data/3d_grids_sample/0b77c16e04a8ac8f84ae9ccaf9b1aaa0_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08a45fe3a80e5c8cc22b3a0f59d531353b6d401f88661f593abdb69dc44b45c4
+size 380288

data/3d_grids_sample/0b77c16e04a8ac8f84ae9ccaf9b1aaa0_0.xyz

@@ -0,0 +1,13 @@
+C 0.0592192692 1.368762351 0.0093929006
+N 1.3122139041 1.8377169369 -0.0001856573
+N 2.0617385494 0.7170865182 -0.013873644
+C 3.4112525731 0.634008202 -0.0275595021
+N 4.0198493068 -0.5162508406 -0.0400141125
+C 3.2187933425 -1.6270159107 -0.0384447767
+N 1.909989902 -1.6610752725 -0.0258242213
+C 1.292800839 -0.4606631607 -0.0130004803
+C -0.0242733248 -0.0346384436 0.0023033991
+H -0.7575875186 2.0764585031 0.0213803521
+H 3.9530882705 1.5750751636 -0.0275449692
+H 3.7430423978 -2.5779490055 -0.0489488368
+H -0.9034361909 -0.6555191213 0.0074057982

data/3d_grids_sample/0b88d98ac218831353fb8c61aea0cfe8.smi

@@ -0,0 +1 @@
+C1=C(C2CC2)NN=N1

data/3d_grids_sample/0b88d98ac218831353fb8c61aea0cfe8_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a8ab9d5c9c61e0547e3a35f52a60354f12a348d681f1d7dc90dc42fed3bb577
+size 1382528

data/3d_grids_sample/0b88d98ac218831353fb8c61aea0cfe8_0.xyz

@@ -0,0 +1,15 @@
+C 0.002730801 1.5364883046 -0.0652481333
+C 1.3129515531 0.8070922114 -0.076629951
+C 0.0187818375 0.0226915472 0.0386968119
+C -0.3380018988 -0.6381328815 1.3049549011
+C -0.1608829834 -0.3235247554 2.6381345568
+N -0.6971640624 -1.3104697961 3.4082804995
+N -1.2012930219 -2.2290873625 2.6387797732
+N -0.9871902476 -1.8273143514 1.3659063658
+H -0.2724161486 2.0867325654 0.827333906
+H -0.3566469563 1.9707288421 -0.9912789551
+H 1.8591050209 0.7372804887 -1.0105007132
+H 1.9370129991 0.8562070416 0.8084014071
+H -0.2904338773 -0.518969552 -0.8499640332
+H 0.3118506955 0.5391482945 3.0790387057
+H -1.3069889708 -2.4074874964 0.6066944785

data/3d_grids_sample/0b94b07e1ec5e58964bfc7e670a359fb.smi

@@ -0,0 +1 @@
+CNC1=NC(O)=C(C)N1

data/3d_grids_sample/0b94b07e1ec5e58964bfc7e670a359fb_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04ca6c042fdc6c92bad68d750b4770bab1cbba2f1c6a016350429da115f7145f
+size 2628416

data/3d_grids_sample/0b94b07e1ec5e58964bfc7e670a359fb_0.xyz

@@ -0,0 +1,18 @@
+C 0.2021885028 1.2833147668 0.1828347559
+N -0.2183533194 -0.0600235049 -0.1763289174
+C 0.4792782987 -0.7562350227 -1.1452089495
+N 0.5549235456 -2.0713645176 -1.2135880027
+C 1.2651968247 -2.3183417769 -2.3633449011
+O 1.5374666899 -3.5896693927 -2.7487695103
+C 1.6425751596 -1.1756799637 -3.0064836641
+C 2.4261233754 -0.9183617354 -4.2431958005
+N 1.1239546418 -0.1576445941 -2.1801339353
+H 1.2655154169 1.3576381934 0.4589582859
+H 0.0149717259 1.9764828738 -0.6460648109
+H -0.4014962835 1.6279653806 1.0247916362
+H -0.4239665093 -0.6625739292 0.6073847095
+H 1.1295866096 -4.1597040968 -2.0868842735
+H 1.8581192728 -0.3597192603 -5.0000767893
+H 3.3521095167 -0.3582763401 -4.0525753219
+H 2.7082322022 -1.875782489 -4.6877635002
+H 1.1534957895 0.8267204188 -2.3790310507

data/datasets/moleculenet/qm9/qm9.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b256baa2f324e03d73c6160f9b573d752a8e16512e6939f0f3f04bb3d7dd8c60
+size 39132032

data/datasets/moleculenet/qm9/train.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc96b329d6e0efe1d59359ddb96881209743518ac2910ac3b4a184f70b866f2d
+size 31696921

finetune/args.py

@@ -0,0 +1,40 @@
+import argparse
+
+
+def get_parser(parser=None):
+    if parser is None:
+        parser = argparse.ArgumentParser()
+
+    parser.add_argument("--data_root", type=str, required=False, default="")
+    parser.add_argument("--grid_path", type=str, required=False, default="")
+
+    parser.add_argument(
+        "--lr_start", type=float, default=3 * 1e-4, help="Initial lr value"
+    )
+    parser.add_argument(
+        "--max_epochs", type=int, required=False, default=1, help="max number of epochs"
+    )
+    
+    parser.add_argument("--num_workers", type=int, default=0, required=False)
+    parser.add_argument("--dropout", type=float, default=0.1, required=False)
+    parser.add_argument("--n_batch", type=int, default=512, help="Batch size")
+    parser.add_argument("--dataset_name", type=str, required=False, default="sol")
+    parser.add_argument("--measure_name", type=str, required=False, default="measure")
+    parser.add_argument("--checkpoints_folder", type=str, required=True)
+    parser.add_argument("--model_path", type=str, default="./smi_ted/")
+    parser.add_argument("--ckpt_filename", type=str, default="smi_ted_Light_40.pt")
+    parser.add_argument("--restart_filename", type=str, default="")
+    parser.add_argument('--n_output', type=int, default=1)
+    parser.add_argument("--save_every_epoch", type=int, default=0)
+    parser.add_argument("--save_ckpt", type=int, default=1)
+    parser.add_argument("--start_seed", type=int, default=0)
+    parser.add_argument("--target_metric", type=str, default="rmse")
+    parser.add_argument("--loss_fn", type=str, default="mae")
+
+    return parser
+
+
+def parse_args():
+    parser = get_parser()
+    args = parser.parse_args()
+    return args

finetune/dataset/__init__.py

@@ -0,0 +1,7 @@
+# from dataset.breast_uka import BreastUKA
+# from dataset.mrnet import MRNetDataset
+# from dataset.brats import BRATSDataset
+# from dataset.adni import ADNIDataset
+# from dataset.duke import DUKEDataset
+# from dataset.lidc import LIDCDataset
+from dataset.default import GridDataset

finetune/dataset/default.py

@@ -0,0 +1,35 @@
+import os
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch.multiprocessing as mp
+
+class GridDataset(Dataset):
+    def __init__(self, dataset, target: str, root_dir: str, internal_resolution: int):
+        super().__init__()
+        self.dataset = dataset
+        self.target = target
+        self.root_dir = root_dir
+        self.internal_resolution = internal_resolution
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int):
+        target = self.dataset.iloc[idx][self.target]
+        filename = self.dataset.iloc[idx]['3d_grid']
+        try:
+            numpy_file = np.load(os.path.join(self.root_dir, filename))
+            torch_np = torch.from_numpy(numpy_file)
+            torch_np = torch_np.unsqueeze(0).unsqueeze(0).float()  # Convert to float and move to appropriate device
+            interpolated_data = F.interpolate(input=torch_np, size=(self.internal_resolution, self.internal_resolution, self.internal_resolution), mode='trilinear')
+
+            # Apply tanh and log operations
+            # interpolated_data_tanh = torch.tanh(interpolated_data)
+            interpolated_data_log = torch.log(interpolated_data + 1).squeeze(0)  # Adding 1 to avoid log(0)
+
+            return interpolated_data_log, target
+        except Exception as e:
+            print(f"Error loading file '{filename}': {e}")
+            return None

finetune/finetune_regression.py

@@ -0,0 +1,92 @@
+# Deep learning
+import torch
+import torch.nn as nn
+from torch import optim
+from trainers import TrainerRegressor
+from vq_gan_3d.model.vqgan_DDP import load_VQGAN
+from utils import init_weights, RMSELoss
+
+# Parallel
+from torch.distributed import init_process_group, destroy_process_group
+
+# Data
+import pandas as pd
+import numpy as np
+
+# Standard library
+import math
+import args
+import os
+
+
+def ddp_setup():
+    init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+def main(config):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    ddp_setup()
+
+    # load dataset
+    df_train = pd.read_csv(f"{config.data_root}/train.csv")
+    df_valid = pd.read_csv(f"{config.data_root}/valid.csv")
+    df_test  = pd.read_csv(f"{config.data_root}/test.csv")
+
+    # load model
+    model = load_VQGAN(folder=config.model_path, filename=config.ckpt_filename)
+    model.net.apply(init_weights)
+    print(model.net)
+
+    # disable gradients to frozen parts
+    for param in model.decoder.parameters():  # decoder
+        param.requires_grad = False
+    for param in model.post_vq_conv.parameters():  # after codebook
+        param.requires_grad = False
+    for param in model.codebook.parameters():  # codebook
+        param.requires_grad = False
+    for param in model.image_discriminator.parameters():  # GAN discriminator
+        param.requires_grad = False
+
+    if config.loss_fn == 'rmse':
+        loss_function = RMSELoss()
+    elif config.loss_fn == 'mae':
+        loss_function = nn.L1Loss()
+
+    # init trainer
+    trainer = TrainerRegressor(
+        raw_data=(df_train, df_valid, df_test),
+        grids_path=config.grid_path,
+        dataset_name=config.dataset_name,
+        target=config.measure_name,
+        batch_size=config.n_batch,
+        hparams=config,
+        internal_resolution=model.config['model']['internal_resolution'],
+        target_metric=config.target_metric,
+        seed=config.start_seed,
+        num_workers=config.num_workers,
+        checkpoints_folder=config.checkpoints_folder,
+        restart_filename=config.restart_filename,
+        device=device,
+        save_every_epoch=bool(config.save_every_epoch),
+        save_ckpt=bool(config.save_ckpt)
+    )
+    trainer.compile(
+        model=model,
+        optimizer=optim.AdamW(
+            list(model.encoder.parameters())
+            +list(model.pre_vq_conv.parameters())
+            +list(model.net.parameters()), 
+            lr=config.lr_start, betas=(0.9, 0.999)
+        ),
+        loss_fn=loss_function
+    )
+    trainer.fit(max_epochs=config.max_epochs)
+    trainer.evaluate()
+    destroy_process_group()
+
+
+if __name__ == '__main__':
+    parser = args.get_parser()
+    config = parser.parse_args()
+    main(config)

finetune/run_finetune_qm9_alpha.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'alpha' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/alpha' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_cv.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'cv' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/cv' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_g298.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'g298' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/g298' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_gap.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'gap' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/gap' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_h298.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'h298' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/h298' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_homo.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'homo' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/homo' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_lumo.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'lumo' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/lumo' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_mu.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'mu' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/mu' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_r2.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'r2' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/r2' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_u0.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'u0' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/u0' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_u298.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'u298' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/u298' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/run_finetune_qm9_zpve.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc_per_node=1 \
+    finetune_regression.py \
+        --n_batch 8 \
+        --dropout 0.1 \
+        --lr_start 3e-5 \
+        --num_workers 16 \
+        --max_epochs 100 \
+        --model_path '../data/checkpoints/pretrained' \
+        --ckpt_filename 'VQGAN_43.pt' \
+        --data_root '../data/datasets/moleculenet/qm9' \
+        --grid_path '/data_npy/qm9' \
+        --dataset_name qm9 \
+        --measure_name 'zpve' \
+        --checkpoints_folder '../data/checkpoints/finetuned/qm9/zpve' \
+        --loss_fn 'mae' \
+        --target_metric 'mae' \
+        --save_ckpt 1 \
+        --start_seed 0 \
+        --save_every_epoch 0 \
+        --restart_filename '' \

finetune/trainers.py

@@ -0,0 +1,359 @@
+# Deep learning
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
+from dataset.default import GridDataset
+from utils import RMSELoss
+
+# Data
+import pandas as pd
+import numpy as np
+
+# Standard library
+import random
+import args
+import os
+import copy
+import shutil
+from tqdm import tqdm
+
+# Machine Learning
+from sklearn.metrics import mean_absolute_error, r2_score, accuracy_score, roc_auc_score, roc_curve, auc, precision_recall_curve
+from scipy import stats
+from utils import RMSE, sensitivity, specificity
+
+
+class Trainer:
+
+    def __init__(self, raw_data, grids_path, dataset_name, target, batch_size, hparams, internal_resolution,
+                 target_metric='rmse', seed=0, num_workers=0, checkpoints_folder='./checkpoints', restart_filename=None, save_every_epoch=False, save_ckpt=True, device='cpu'):
+        # data
+        self.df_train = raw_data[0]
+        self.df_valid = raw_data[1]
+        self.df_test = raw_data[2]
+        self.grids_path = grids_path
+        self.dataset_name = dataset_name
+        self.target = target
+        self.batch_size = batch_size
+        self.hparams = hparams
+        self.internal_resolution = internal_resolution
+        self.num_workers = num_workers
+        self._prepare_data()
+
+        # config
+        self.target_metric = target_metric
+        self.seed = seed
+        self.checkpoints_folder = checkpoints_folder
+        self.restart_filename = restart_filename
+        self.start_epoch = 1
+        self.save_every_epoch = save_every_epoch
+        self.save_ckpt = save_ckpt
+        self.best_vloss = float('inf')
+        self.last_filename = None
+        self._set_seed(seed)
+
+        # multi-gpu
+        self.local_rank = int(os.environ["LOCAL_RANK"])
+        self.global_rank = int(os.environ["RANK"])
+
+    def _prepare_data(self):
+        train_dataset = GridDataset(
+            dataset=self.df_train,
+            target=self.target,
+            root_dir=self.grids_path, 
+            internal_resolution=self.internal_resolution,
+        )
+        valid_dataset = GridDataset(
+            dataset=self.df_valid,
+            target=self.target,
+            root_dir=self.grids_path, 
+            internal_resolution=self.internal_resolution,
+        )
+        test_dataset = GridDataset(
+            dataset=self.df_test,
+            target=self.target,
+            root_dir=self.grids_path, 
+            internal_resolution=self.internal_resolution,
+        )
+
+        # create dataloader
+        self.train_loader = DataLoader(
+            train_dataset, 
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=DistributedSampler(train_dataset), 
+            shuffle=False, 
+            pin_memory=True
+        )
+        self.valid_loader = DataLoader(
+            valid_dataset, 
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=DistributedSampler(valid_dataset), 
+            shuffle=False, 
+            pin_memory=True
+        )
+        self.test_loader = DataLoader(
+            test_dataset, 
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=DistributedSampler(test_dataset), 
+            shuffle=False, 
+            pin_memory=True
+        )
+
+    def compile(self, model, optimizer, loss_fn):
+        self.model = model.to(self.local_rank)
+        self.optimizer = optimizer
+        self.loss_fn = loss_fn
+        self._print_configuration()
+        
+        if self.restart_filename:
+            self._load_checkpoint(self.restart_filename)
+            print('Checkpoint restored!')
+
+        self.model = DDP(self.model, device_ids=[self.local_rank])
+
+    def fit(self, max_epochs=500):
+        for epoch in range(self.start_epoch, max_epochs+1):
+            print(f'\n=====Epoch [{epoch}/{max_epochs}]=====')
+
+            # training
+            self.model.train()
+            self.train_loader.sampler.set_epoch(epoch)
+            train_loss = self._train_one_epoch()
+
+            # validation
+            self.model.eval()
+            val_preds, val_loss, val_metrics = self._validate_one_epoch(self.valid_loader)
+            tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader)
+
+            if self.global_rank == 0:
+                for m in val_metrics.keys():
+                    print(f"[VALID] Evaluation {m.upper()}: {round(val_metrics[m], 4)}")
+                print('-'*64)
+                for m in tst_metrics.keys():
+                    print(f"[TEST] Evaluation {m.upper()}: {round(tst_metrics[m], 4)}")
+
+            ############################### Save Finetune checkpoint #######################################
+            if ((val_loss < self.best_vloss) or self.save_every_epoch) and self.save_ckpt and self.global_rank == 0:
+                # remove old checkpoint
+                if (self.last_filename != None) and (not self.save_every_epoch):
+                    os.remove(os.path.join(self.checkpoints_folder, self.last_filename))
+
+                # filename
+                model_name = f'{str(self.model.module)}-Finetune'
+                self.last_filename = f"{model_name}_seed{self.seed}_{self.dataset_name}_epoch={epoch}_valloss={round(val_loss, 4)}.pt"
+
+                # update best loss
+                self.best_vloss = val_loss
+
+                # save checkpoint
+                print('Saving checkpoint...')
+                self._save_checkpoint(epoch, self.last_filename)
+
+    def evaluate(self, verbose=True):
+        if verbose:
+            print("\n=====Test Evaluation=====")
+
+        # set model evaluation mode
+        model_inf = copy.deepcopy(self.model)
+        model_inf.eval()
+
+        # evaluate on test set
+        tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader, model_inf)
+
+        if verbose and self.global_rank == 0:
+            # show metrics
+            for m in tst_metrics.keys():
+                print(f"[TEST] Evaluation {m.upper()}: {round(tst_metrics[m], 4)}")
+
+            # save predictions
+            pd.DataFrame(tst_preds).to_csv(
+                os.path.join(
+                    self.checkpoints_folder, 
+                    f'{self.dataset_name}_{self.target if isinstance(self.target, str) else self.target[0]}_predict_test_seed{self.seed}.csv'
+                ), 
+                index=False
+            )
+
+    def _train_one_epoch(self):
+        raise NotImplementedError
+
+    def _validate_one_epoch(self, data_loader, model=None):
+        raise NotImplementedError
+
+    def _print_configuration(self):
+        print('----Finetune information----')
+        print('Dataset:\t', self.dataset_name)
+        print('Target:\t\t', self.target)
+        print('Batch size:\t', self.batch_size)
+        print('LR:\t\t', self._get_lr())
+        print('Device:\t\t', self.local_rank)
+        print('Optimizer:\t', self.optimizer.__class__.__name__)
+        print('Loss function:\t', self.loss_fn.__class__.__name__)
+        print('Seed:\t\t', self.seed)
+        print('Train size:\t', self.df_train.shape[0])
+        print('Valid size:\t', self.df_valid.shape[0])
+        print('Test size:\t', self.df_test.shape[0])
+
+    def _load_checkpoint(self, filename):
+        ckpt_path = os.path.join(self.checkpoints_folder, filename)
+        ckpt_dict = torch.load(ckpt_path, map_location='cpu')
+        self.model.load_state_dict(ckpt_dict['MODEL_STATE'])
+        self.start_epoch = ckpt_dict['EPOCHS_RUN'] + 1
+        self.best_vloss = ckpt_dict['finetune_info']['best_vloss']
+
+    def _save_checkpoint(self, current_epoch, filename):
+        if not os.path.exists(self.checkpoints_folder):
+            os.makedirs(self.checkpoints_folder)
+
+        self.model.module.config['finetune'] = vars(self.hparams)
+        hparams = self.model.module.config
+
+        ckpt_dict = {
+            'MODEL_STATE': self.model.module.state_dict(),
+            'EPOCHS_RUN': current_epoch,
+            'hparams': hparams,
+            'finetune_info': {
+                'dataset': self.dataset_name,
+                'target`': self.target,
+                'batch_size': self.batch_size,
+                'lr': self._get_lr(),
+                'device': self.local_rank,
+                'optim': self.optimizer.__class__.__name__,
+                'loss_fn': self.loss_fn.__class__.__name__,
+                'train_size': self.df_train.shape[0],
+                'valid_size': self.df_valid.shape[0],
+                'test_size': self.df_test.shape[0],
+                'best_vloss': self.best_vloss,
+            },
+            'seed': self.seed,
+        }
+
+        assert list(ckpt_dict.keys()) == ['MODEL_STATE', 'EPOCHS_RUN', 'hparams', 'finetune_info', 'seed']
+
+        torch.save(ckpt_dict, os.path.join(self.checkpoints_folder, filename))
+
+    def _set_seed(self, value):
+        random.seed(value)
+        torch.manual_seed(value)
+        np.random.seed(value)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(value)
+            torch.cuda.manual_seed_all(value)
+            cudnn.deterministic = True
+            cudnn.benchmark = False
+
+    def _get_lr(self):
+        for param_group in self.optimizer.param_groups:
+            return param_group['lr']
+
+
+class TrainerRegressor(Trainer):
+
+    def __init__(self, raw_data, grids_path, dataset_name, target, batch_size, hparams, internal_resolution,
+                 target_metric='rmse', seed=0, num_workers=0, checkpoints_folder='./checkpoints', restart_filename=None, save_every_epoch=False, save_ckpt=True, device='cpu'):
+        super().__init__(raw_data, grids_path, dataset_name, target, batch_size, hparams, internal_resolution,
+                         target_metric, seed, num_workers, checkpoints_folder, restart_filename, save_every_epoch, save_ckpt, device) 
+
+    def _train_one_epoch(self):
+        running_loss = 0.0
+
+        if self.global_rank == 0:
+            pbar = tqdm(total=len(self.train_loader))
+        for idx, data in enumerate(self.train_loader):
+            # Every data instance is an input + label pair
+            grids, targets = data
+            targets = targets.to(self.local_rank)
+            grids = grids.to(self.local_rank)
+
+            # zero the parameter gradients (otherwise they are accumulated)
+            self.optimizer.zero_grad()
+
+            # Make predictions for this batch
+            embeddings = self.model.module.feature_extraction(grids)
+            outputs = self.model.module.net(embeddings).squeeze(1)
+
+            # Compute the loss and its gradients
+            loss = self.loss_fn(outputs, targets)
+            loss.backward()
+
+            # Adjust learning weights
+            self.optimizer.step()
+
+            # print statistics
+            running_loss += loss.item()
+
+            # progress bar
+            if self.global_rank == 0:
+                pbar.update(1)
+                pbar.set_description('[TRAINING]')
+                pbar.set_postfix(loss=running_loss/(idx+1))
+                pbar.refresh()
+        if self.global_rank == 0:
+            pbar.close()
+
+        return running_loss / len(self.train_loader)
+
+    def _validate_one_epoch(self, data_loader, model=None):
+        data_targets = []
+        data_preds = []
+        running_loss = 0.0
+
+        model = self.model if model is None else model
+
+        if self.global_rank == 0:
+            pbar = tqdm(total=len(data_loader))
+        with torch.no_grad():
+            for idx, data in enumerate(data_loader):
+                # Every data instance is an input + label pair
+                grids, targets = data
+                targets = targets.to(self.local_rank)
+                grids = grids.to(self.local_rank)
+
+                # Make predictions for this batch
+                embeddings = model.module.feature_extraction(grids)
+                predictions = model.module.net(embeddings).squeeze(1)
+
+                # Compute the loss
+                loss = self.loss_fn(predictions, targets)
+
+                data_targets.append(targets.view(-1))
+                data_preds.append(predictions.view(-1))
+
+                # print statistics
+                running_loss += loss.item()
+
+                # progress bar
+                if self.global_rank == 0:
+                    pbar.update(1)
+                    pbar.set_description('[EVALUATION]')
+                    pbar.set_postfix(loss=running_loss/(idx+1))
+                    pbar.refresh()
+        if self.global_rank == 0:
+            pbar.close()
+
+        # Put together predictions and labels from batches
+        preds = torch.cat(data_preds, dim=0).cpu().numpy()
+        tgts = torch.cat(data_targets, dim=0).cpu().numpy()
+
+        # Compute metrics
+        mae = mean_absolute_error(tgts, preds)
+        r2 = r2_score(tgts, preds)
+        rmse = RMSE(preds, tgts)
+        spearman = stats.spearmanr(tgts, preds).correlation # scipy 1.12.0
+
+        # Rearange metrics
+        metrics = {
+            'mae': mae,
+            'r2': r2,
+            'rmse': rmse,
+            'spearman': spearman,
+        }
+
+        return preds, running_loss / len(data_loader), metrics

finetune/utils.py

@@ -0,0 +1,126 @@
+# Deep learning
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+from sklearn.metrics import confusion_matrix
+
+# Data
+import pandas as pd
+import numpy as np
+
+# Standard library
+import os
+
+# Chemistry
+from rdkit import Chem
+from rdkit.Chem import PandasTools
+from rdkit.Chem import Descriptors
+PandasTools.RenderImagesInAllDataFrames(True)
+
+
+def normalize_smiles(smi, canonical=True, isomeric=False):
+    try:
+        normalized = Chem.MolToSmiles(
+        Chem.MolFromSmiles(smi), canonical=canonical, isomericSmiles=isomeric
+        )
+    except:
+        normalized = None
+    return normalized
+
+
+class RMSELoss:
+    def __init__(self):
+        pass
+
+    def __call__(self, yhat, y):
+        return torch.sqrt(torch.mean((yhat-y)**2))
+
+
+def RMSE(predictions, targets):
+    return np.sqrt(((predictions - targets) ** 2).mean())
+
+
+def sensitivity(y_true, y_pred):
+    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
+    return (tp/(tp+fn))
+
+
+def specificity(y_true, y_pred):
+    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
+    return (tn/(tn+fp)) 
+
+
+def init_weights(module):
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+
+def get_optim_groups(module, keep_decoder=False):
+    # setup optimizer
+    # separate out all parameters to those that will and won't experience regularizing weight decay
+    decay = set()
+    no_decay = set()
+    whitelist_weight_modules = (torch.nn.Linear,)
+    blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+    for mn, m in module.named_modules():
+        for pn, p in m.named_parameters():
+            fpn = '%s.%s' % (mn, pn) if mn else pn # full param name
+
+            if not keep_decoder and 'decoder' in fpn: # exclude decoder components
+                continue
+
+            if pn.endswith('bias'):
+                # all biases will not be decayed
+                no_decay.add(fpn)
+            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                # weights of whitelist modules will be weight decayed
+                decay.add(fpn)
+            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                # weights of blacklist modules will NOT be weight decayed
+                no_decay.add(fpn)
+
+    # validate that we considered every parameter
+    param_dict = {pn: p for pn, p in module.named_parameters()}
+    
+    # create the pytorch optimizer object
+    optim_groups = [
+        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.0},
+        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+    ]
+
+    return optim_groups
+
+
+class CustomDataset(Dataset):
+    def __init__(self, dataset, target):
+        self.dataset = dataset
+        self.target = target
+
+    def __len__(self):
+        return len(self.dataset)    
+    
+    def __getitem__(self, idx):
+        smiles = self.dataset['canon_smiles'].iloc[idx]
+        labels = self.dataset[self.target].iloc[idx]
+        return smiles, labels
+
+
+class CustomDatasetMultitask(Dataset):
+    def __init__(self, dataset, targets):
+        self.dataset = dataset
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.dataset)    
+    
+    def __getitem__(self, idx):
+        smiles = self.dataset['canon_smiles'].iloc[idx]
+        labels = self.dataset[self.targets].iloc[idx].to_numpy()
+        mask = [0.0 if np.isnan(x) else 1.0 for x in labels]
+        labels = [0.0 if np.isnan(x) else x for x in labels]
+        return smiles, torch.tensor(labels, dtype=torch.float32), torch.tensor(mask)

inference/run_embeddings_eval_xgboost.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+python -u ./scripts/evaluate_embeddings_xgboost.py --task zpve

inference/run_extract_embeddings.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+python -u ./scripts/extract_embeddings.py \
+    --dataset_path '../data/datasets/moleculenet/qm9.csv' \
+    --save_dataset_path '../data/embeddings/qm9_embeddings.csv' \
+    --ckpt_filename 'VQGAN_43.pt' \
+    --data_dir '../data/sample_data_schema' \
+    --batch_size 2 \
+    --num_workers 0 \