Create train.py

Run this file to train models.

train.py

@@ -0,0 +1,257 @@
+from lightning.pytorch import seed_everything
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.callbacks.early_stopping import EarlyStopping
+import lightning.pytorch as pl
+from pytorch_lightning.loggers import TensorBoardLogger
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from transformers import AutoTokenizer
+from ast import literal_eval
+
+# imports from our own modules
+import config
+from model import (
+    BERTContrastiveLearning_simcse,
+    BERTContrastiveLearning_simcse_w,
+    BERTContrastiveLearning_samp,
+    BERTContrastiveLearning_samp_w,
+)
+from dataset import (
+    ContrastiveLearningDataModule_simcse,
+    ContrastiveLearningDataModule_simcse_w,
+    ContrastiveLearningDataModule_samp,
+    ContrastiveLearningDataModule_samp_w,
+)
+
+if __name__ == "__main__":
+    seed_everything(0, workers=True)
+
+    # Initialize tensorboard logger
+    logger = TensorBoardLogger("logs", name="MIMIC-tr")
+
+    query_df = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/mimic_data/processed_train/processed.csv"
+    )
+    # query_df = query_df.head(1000)
+    query_df["concepts"] = query_df["concepts"].apply(literal_eval)
+    query_df["codes"] = query_df["codes"].apply(literal_eval)
+    query_df["codes"] = query_df["codes"].apply(
+        lambda x: [val for val in x if val is not None]
+    )  # remove None in lists
+    query_df = query_df.drop(columns=["one_hot"])
+    train_df, val_df = train_test_split(query_df, test_size=config.split_ratio)
+
+    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+    sim_df = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/pairwise_scores.csv"
+    )
+
+    all_d = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/all_d_full.csv"
+    )
+    all_d["synonyms"] = all_d["synonyms"].apply(literal_eval)
+    all_d["ancestors"] = all_d["ancestors"].apply(literal_eval)
+    dictionary = dict(zip(all_d["concept"], all_d["synonyms"]))
+
+    # SimCSE
+    data_module1 = ContrastiveLearningDataModule_simcse(
+        train_df,
+        val_df,
+        tokenizer,
+    )
+    data_module1.setup()
+
+    print("Number of training data:", len(data_module1.train_dataset))
+    print("Number of validation data:", len(data_module1.val_dataset))
+
+    model1 = BERTContrastiveLearning_simcse(
+        n_batches=len(data_module1.train_dataset) / config.batch_size,
+        n_epochs=config.max_epochs,
+        lr=config.learning_rate,
+        unfreeze=config.unfreeze_ratio,
+    )
+
+    checkpoint1 = ModelCheckpoint(
+        dirpath="/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/ckpt/simcse/v1",
+        filename="{epoch}-{step}",
+        # save_weights_only=True,
+        save_last=True,
+        every_n_train_steps=config.log_every_n_steps,
+        monitor=None,
+        save_top_k=-1,
+    )
+
+    trainer1 = pl.Trainer(
+        accelerator=config.accelerator,
+        devices=config.devices,
+        strategy="ddp",
+        logger=logger,
+        max_epochs=config.max_epochs,
+        min_epochs=config.min_epochs,
+        precision=config.precision,
+        callbacks=[
+            EarlyStopping(
+                monitor="validation_loss", min_delta=1e-3, patience=3, mode="min"
+            ),
+            checkpoint1,
+        ],
+        profiler="simple",
+        log_every_n_steps=config.log_every_n_steps,
+    )
+
+    trainer1.fit(model1, data_module1)
+
+    # SimCSE_w
+    data_module2 = ContrastiveLearningDataModule_simcse_w(
+        train_df,
+        val_df,
+        query_df,
+        tokenizer,
+        sim_df,
+        all_d,
+    )
+    data_module2.setup()
+
+    print("Number of training data:", len(data_module2.train_dataset))
+    print("Number of validation data:", len(data_module2.val_dataset))
+
+    model2 = BERTContrastiveLearning_simcse_w(
+        n_batches=len(data_module2.train_dataset) / config.batch_size,
+        n_epochs=config.max_epochs,
+        lr=config.learning_rate,
+        unfreeze=config.unfreeze_ratio,
+    )
+
+    checkpoint2 = ModelCheckpoint(
+        dirpath="/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/ckpt/simcse_w/v1",
+        filename="{epoch}-{step}",
+        # save_weights_only=True,
+        save_last=True,
+        every_n_train_steps=config.log_every_n_steps,
+        monitor=None,
+        save_top_k=-1,
+    )
+
+    trainer2 = pl.Trainer(
+        accelerator=config.accelerator,
+        devices=config.devices,
+        strategy="ddp",
+        logger=logger,
+        max_epochs=config.max_epochs,
+        min_epochs=config.min_epochs,
+        precision=config.precision,
+        callbacks=[
+            EarlyStopping(
+                monitor="validation_loss", min_delta=1e-3, patience=3, mode="min"
+            ),
+            checkpoint2,
+        ],
+        profiler="simple",
+        log_every_n_steps=config.log_every_n_steps,
+    )
+
+    trainer2.fit(model2, data_module2)
+
+    # Samp
+    data_module3 = ContrastiveLearningDataModule_samp(
+        train_df,
+        val_df,
+        query_df,
+        tokenizer,
+        dictionary,
+        sim_df,
+    )
+    data_module3.setup()
+
+    print("Number of training data:", len(data_module3.train_dataset))
+    print("Number of validation data:", len(data_module3.val_dataset))
+
+    model3 = BERTContrastiveLearning_samp(
+        n_batches=len(data_module3.train_dataset) / config.batch_size,
+        n_epochs=config.max_epochs,
+        lr=config.learning_rate,
+        unfreeze=config.unfreeze_ratio,
+    )
+
+    checkpoint3 = ModelCheckpoint(
+        dirpath="/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/ckpt/samp/v1",
+        filename="{epoch}-{step}",
+        # save_weights_only=True,
+        save_last=True,
+        every_n_train_steps=config.log_every_n_steps,
+        monitor=None,
+        save_top_k=-1,
+    )
+
+    trainer3 = pl.Trainer(
+        accelerator=config.accelerator,
+        devices=config.devices,
+        strategy="ddp",
+        logger=logger,
+        max_epochs=config.max_epochs,
+        min_epochs=config.min_epochs,
+        precision=config.precision,
+        callbacks=[
+            EarlyStopping(
+                monitor="validation_loss", min_delta=1e-3, patience=3, mode="min"
+            ),
+            checkpoint3,
+        ],
+        profiler="simple",
+        log_every_n_steps=config.log_every_n_steps,
+    )
+
+    trainer3.fit(model3, data_module3)
+
+    # Samp_w
+    data_module4 = ContrastiveLearningDataModule_samp_w(
+        train_df,
+        val_df,
+        query_df,
+        tokenizer,
+        dictionary,
+        sim_df,
+        all_d,
+    )
+    data_module4.setup()
+
+    print("Number of training data:", len(data_module4.train_dataset))
+    print("Number of validation data:", len(data_module4.val_dataset))
+
+    model4 = BERTContrastiveLearning_samp_w(
+        n_batches=len(data_module4.train_dataset) / config.batch_size,
+        n_epochs=config.max_epochs,
+        lr=config.learning_rate,
+        unfreeze=config.unfreeze_ratio,
+    )
+
+    checkpoint4 = ModelCheckpoint(
+        dirpath="/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/ckpt/samp_w/v1",
+        filename="{epoch}-{step}",
+        # save_weights_only=True,
+        save_last=True,
+        every_n_train_steps=config.log_every_n_steps,
+        monitor=None,
+        save_top_k=-1,
+    )
+
+    trainer4 = pl.Trainer(
+        accelerator=config.accelerator,
+        devices=config.devices,
+        strategy="ddp",
+        logger=logger,
+        max_epochs=config.max_epochs,
+        min_epochs=config.min_epochs,
+        precision=config.precision,
+        callbacks=[
+            EarlyStopping(
+                monitor="validation_loss", min_delta=1e-3, patience=3, mode="min"
+            ),
+            checkpoint4,
+        ],
+        profiler="simple",
+        log_every_n_steps=config.log_every_n_steps,
+    )
+
+    trainer4.fit(model4, data_module4)