Upload 5 files

config.py

@@ -0,0 +1,29 @@
+import torch
+
+# Set device cuda for GPU if it is available, otherwise run on the CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# loss
+t_p = 0.25  # optimal: 1/8 ~ 1/32
+zeta = 3  # optimal: 2 ~ 5
+# m = 0.2
+
+# Training hyperparameters
+min_epochs = 3
+max_epochs = 30
+learning_rate = 5e-5
+unfreeze_ratio = 1
+mlm_weight = 0.5  # optimal: 0.5~0.75
+
+# Dataset
+batch_size = 100
+split_ratio = 0.2
+
+# Logger
+log_every_n_steps = 50
+ckcpt_every_n_steps = 5000
+
+# Compute related
+accelerator = "gpu"
+devices = 1  # number of gpus
+precision = "16-mixed"

dataset.py

@@ -0,0 +1,351 @@
+import torch
+from transformers import AutoTokenizer
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.utils.rnn import pad_sequence
+import lightning.pytorch as pl
+import config
+import pandas as pd
+import copy
+from ast import literal_eval
+from sklearn.model_selection import train_test_split
+import random
+
+
+def get_code_by_entity(entity, dictionary):
+    """
+    Query the dictionary by entity and return its code.
+    Return the key with the longest value list if multiple keys found.
+    """
+    keys = []
+    length = []
+    for key, values in dictionary.items():
+        if entity in values:
+            keys.append(key)
+            length.append(len(values))
+    d = dict(zip(keys, length))
+    if len(d) > 0:
+        return max(d, key=d.get)
+    else:
+        return None
+
+
+def num_ancestors(df, code):
+    result = len(df.loc[df["concept"] == code, "ancestors"].values[0])
+    return result
+
+
+def get_score(df, code1, code2):
+    result = df[
+        ((df["Code1"] == code1) & (df["Code2"] == code2))
+        | ((df["Code1"] == code2) & (df["Code2"] == code1))
+    ]
+
+    if result.empty:
+        return None
+
+    return result.iloc[0]["score"]
+
+
+def mask(tokenizer, dictionary, unique_d, text, entities, anchor=True):
+    """
+    Randomly select one entity from the entities, mask the first existence in the text and create duplicates with synonyms. The rest are treated as context.
+
+    Returns a dictionary {input_ids, attention_mask, mlm_labels, masked_indices, tags}.
+    """
+    if anchor is True:
+        entity = random.choice(entities)
+        code = get_code_by_entity(entity, dictionary)
+        try:
+            synonyms = dictionary[code]
+        except:
+            return None
+        text_token = tokenizer.tokenize(text)
+        ent_token = tokenizer.tokenize(entity.lower())
+        num_ent_token = len(ent_token)
+
+        input_ids = [copy.deepcopy(text_token) for _ in range(len(synonyms))]
+        mlm_labels = [copy.deepcopy(text_token) for _ in range(len(synonyms))]
+        masked_indices = []
+
+        for i, t in enumerate(mlm_labels):
+            start_indices = [
+                index for index, value in enumerate(t) if value == ent_token[0]
+            ]
+            masked_index = []
+            for start in start_indices:
+                if (
+                    tokenizer.convert_tokens_to_string(t[start : start + num_ent_token])
+                    == entity.lower()
+                ) and len(masked_index) == 0:
+                    syn = tokenizer.tokenize(synonyms[i])
+                    mlm_labels[i][start : start + num_ent_token] = syn
+                    input_ids[i][start : start + num_ent_token] = ["[MASK]"] * len(syn)
+                    masked_index.extend(list(range(start, start + len(syn))))
+            masked_indices.append(masked_index)
+
+        if any(not sublist for sublist in masked_indices):
+            empty_mask_idx = [
+                k for k, sublist in enumerate(masked_indices) if not sublist
+            ]
+            input_ids = [x for i, x in enumerate(input_ids) if i not in empty_mask_idx]
+            mlm_labels = [
+                x for i, x in enumerate(mlm_labels) if i not in empty_mask_idx
+            ]
+            masked_indices = [
+                sublist for k, sublist in enumerate(masked_indices) if sublist
+            ]
+
+        if len(input_ids) <= 1:
+            return None
+
+        input_ids_lst = []
+        attention_mask_lst = []
+        mlm_labels_lst = []
+
+        for j, token in enumerate(input_ids):
+            input_id = torch.tensor(tokenizer.convert_tokens_to_ids(token))
+            input_ids_lst.append(input_id)
+            attention_mask_lst.append(torch.ones_like(input_id))
+            mlm_label = torch.tensor(tokenizer.convert_tokens_to_ids(mlm_labels[j]))
+            for l in range(len(mlm_label)):
+                if l not in masked_indices[j]:
+                    mlm_label[l] = -100
+            mlm_labels_lst.append(mlm_label)
+
+        tags = [1] * len(input_ids_lst)
+        tags[0] = 0
+        codes = [code] * len(input_ids_lst)
+        if code not in unique_d:
+            return None
+
+        out = {
+            "input_ids": input_ids_lst,
+            "attention_mask": attention_mask_lst,
+            "mlm_labels": mlm_labels_lst,
+            "masked_indices": masked_indices,
+            "tags": tags,
+            "codes": codes,
+        }
+
+    if anchor is False:
+        entity = random.choice(entities)
+        code = get_code_by_entity(entity, dictionary)
+        input_ids = tokenizer.tokenize(text)
+        mlm_labels = copy.deepcopy(input_ids)
+        ent_token = tokenizer.tokenize(entity.lower())
+        num_ent_token = len(ent_token)
+        masked_indices = []
+
+        start_indices = []
+        for i, t in enumerate(mlm_labels):
+            if t == ent_token[0]:
+                start_indices.append(i)
+
+        for start in start_indices:
+            if (
+                tokenizer.convert_tokens_to_string(
+                    input_ids[start : start + num_ent_token]
+                )
+                == entity.lower()
+            ) and len(masked_indices) == 0:
+                input_ids[start : start + num_ent_token] = ["[MASK]"] * num_ent_token
+                masked_indices.extend(list(range(start, start + num_ent_token)))
+
+        if len(masked_indices) == 0:
+            return None
+
+        input_ids_lst = []
+        attention_mask_lst = []
+        mlm_labels_lst = []
+
+        input_id = torch.tensor(tokenizer.convert_tokens_to_ids(input_ids))
+        input_ids_lst.append(input_id)
+        attention_mask_lst.append(torch.ones_like(input_id))
+        mlm_labels = tokenizer.convert_tokens_to_ids(mlm_labels)
+        for l in range(len(mlm_labels)):
+            if l not in masked_indices:
+                mlm_labels[l] = -100
+        mlm_labels_lst.append(torch.tensor(mlm_labels))
+
+        tags = [2] * len(input_ids_lst)
+        code = get_code_by_entity(entity, dictionary)
+
+        if code not in unique_d:
+            return None
+
+        codes = [code] * len(input_ids_lst)
+
+        out = {
+            "input_ids": input_ids_lst,
+            "attention_mask": attention_mask_lst,
+            "mlm_labels": mlm_labels_lst,
+            "masked_indices": masked_indices,
+            "tags": tags,
+            "codes": codes,
+        }
+
+    return out
+
+
+class CLDataset(Dataset):
+    def __init__(
+        self,
+        data: pd.DataFrame,
+    ):
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        data_row = self.data.iloc[index]
+        sentence = data_row.sentences
+        concepts = data_row.concepts
+        return [sentence, concepts]
+
+
+def collate_func(batch, tokenizer, dictionary, all_d, pairs):
+    input_ids_lst = []
+    attention_mask_lst = []
+    mlm_labels_lst = []
+    masked_indices_lst = []
+    tags_lst = []
+    codes_lst = []
+    scores_lst = []
+
+    unique_d = pairs["Code1"].unique()
+
+    anchor = batch[0]
+    anchor_masked = mask(tokenizer, dictionary, unique_d, anchor[0], anchor[1])
+    while anchor_masked is None:
+        batch = batch[1:]
+        anchor = batch[0]
+        anchor_masked = mask(tokenizer, dictionary, unique_d, anchor[0], anchor[1])
+
+    for i in range(len(anchor_masked["input_ids"])):
+        input_ids_lst.append(anchor_masked["input_ids"][i])
+        attention_mask_lst.append(anchor_masked["attention_mask"][i])
+        mlm_labels_lst.append(anchor_masked["mlm_labels"][i])
+    masked_indices_lst.extend(anchor_masked["masked_indices"])
+    tags_lst.extend(anchor_masked["tags"])
+    codes_lst.extend(anchor_masked["codes"])
+    ap_code = anchor_masked["codes"][0]
+    ap_score = num_ancestors(all_d, ap_code)
+    scores_lst.extend([ap_score] * len(tags_lst))
+
+    negatives = batch[1:]
+    for neg in negatives:
+        neg_masked = mask(tokenizer, dictionary, unique_d, neg[0], neg[1], False)
+        if neg_masked is None:
+            continue
+
+        for j in range(len(neg_masked["input_ids"])):
+            input_ids_lst.append(neg_masked["input_ids"][j])
+            attention_mask_lst.append(neg_masked["attention_mask"][j])
+        mlm_labels_lst.extend(neg_masked["mlm_labels"])
+        masked_indices_lst.append(neg_masked["masked_indices"])
+        tags_lst.extend(neg_masked["tags"])
+        codes_lst.extend(neg_masked["codes"])
+        n_code = neg_masked["codes"][0]
+        if n_code == ap_code:
+            an_score = num_ancestors(all_d, n_code)
+        else:
+            an_score = get_score(pairs, ap_code, n_code)
+        scores_lst.append(an_score)
+
+    padded_input_ids = pad_sequence(input_ids_lst, padding_value=0)
+    padded_input_ids = torch.transpose(padded_input_ids, 0, 1)
+
+    padded_attention_mask = pad_sequence(attention_mask_lst, padding_value=0)
+    padded_attention_mask = torch.transpose(padded_attention_mask, 0, 1)
+
+    padded_mlm_labels = pad_sequence(mlm_labels_lst, padding_value=-100)
+    padded_mlm_labels = torch.transpose(padded_mlm_labels, 0, 1)
+
+    return {
+        "input_ids": padded_input_ids,
+        "attention_mask": padded_attention_mask,
+        "mlm_labels": padded_mlm_labels,
+        "masked_indices": masked_indices_lst,
+        "tags": tags_lst,
+        "codes": codes_lst,
+        "scores": scores_lst,
+    }
+
+
+def create_dataloader(dataset, tokenizer, dictionary, all_d, pairs, shuffle):
+    return DataLoader(
+        dataset,
+        batch_size=config.batch_size,
+        shuffle=shuffle,
+        num_workers=1,
+        collate_fn=lambda batch: collate_func(
+            batch, tokenizer, dictionary, all_d, pairs
+        ),
+    )
+
+
+class CLDataModule(pl.LightningDataModule):
+    def __init__(self, train_df, val_df, tokenizer, dictionary, all_d, pairs):
+        super().__init__()
+        self.train_df = train_df
+        self.val_df = val_df
+        self.tokenizer = tokenizer
+        self.dictionary = dictionary
+        self.all_d = all_d
+        self.pairs = pairs
+
+    def setup(self, stage=None):
+        self.train_dataset = CLDataset(self.train_df)
+        self.val_dataset = CLDataset(self.val_df)
+
+    def train_dataloader(self):
+        return create_dataloader(
+            self.train_dataset,
+            self.tokenizer,
+            self.dictionary,
+            self.all_d,
+            self.pairs,
+            shuffle=True,
+        )
+
+    def val_dataloader(self):
+        return create_dataloader(
+            self.val_dataset,
+            self.tokenizer,
+            self.dictionary,
+            self.all_d,
+            self.pairs,
+            shuffle=False,
+        )
+
+
+if __name__ == "__main__":
+    query_df = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/query_df.csv"
+    )
+    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]
+    )
+    train_df, val_df = train_test_split(query_df, test_size=config.split_ratio)
+
+    all_d = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/query_all_d.csv"
+    )
+    all_d.drop(columns=["finding_sites", "morphology"], inplace=True)
+    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"]))
+
+    pairs = pd.read_csv("/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/pairs.csv")
+
+    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+    d = CLDataModule(train_df, val_df, tokenizer, dictionary, all_d, pairs)
+    d.setup()
+    train = d.train_dataloader()
+    for batch in train:
+        b = batch
+        break

loss.py

@@ -0,0 +1,58 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import config
+
+
+class PCCL(nn.Module):
+    """
+    Pair-wise Cost-sensitive Contrastive Loss.
+
+    feature_matrix: (B,F)
+    label: (B,)
+    """
+
+    def __init__(self):
+        super(PCCL, self).__init__()
+        self.t_p = config.t_p  # positive temperature
+        self.zeta = config.zeta  # temperature ratio
+        self.t_n = config.t_p * self.zeta  # negative temperature
+        # self.m = config.m  # fixed margin
+
+    def forward(self, feature_matrix, label, score):
+        feature_matrix_normalized = F.normalize(feature_matrix, p=2, dim=1)
+        anchor = feature_matrix_normalized[0 : label.index(1)]
+        positives = feature_matrix_normalized[label.index(1) : label.index(2)]
+        pos_cardinal = positives.shape[0]
+        negatives = feature_matrix_normalized[label.index(2) :]
+        min_score = min(score)
+        max_score = max(score)
+        normalized_score = [((x - min_score) / (max_score - min_score)) for x in score]
+        pos_scores = torch.tensor(normalized_score[label.index(1) : label.index(2)])
+        neg_scores = torch.tensor(normalized_score[label.index(2) :])
+
+        # within-class similarity
+        s_i_p = F.cosine_similarity(positives, anchor, dim=1)
+        # between-class similarity
+        s_i_n = F.cosine_similarity(negatives, anchor, dim=1)
+
+        pos_scores = pos_scores.to(s_i_p.device)
+        neg_scores = neg_scores.to(s_i_n.device)
+
+        # pair-wise relaxation factors
+        alpha_i_p = 1 + torch.max(torch.zeros_like(s_i_p), (pos_scores - s_i_p))
+        alpha_i_n = 1 + torch.max(torch.zeros_like(s_i_n), (neg_scores + s_i_n))
+
+        # normalization factor
+        z = torch.sum(torch.exp(torch.div(alpha_i_p, self.t_p) * s_i_p)) + torch.sum(
+            torch.exp(torch.div(alpha_i_n, self.t_n) * s_i_n)
+        )
+
+        # loss
+        loss = torch.sum(
+            torch.log(torch.div(torch.exp(alpha_i_p * torch.div(s_i_p, self.t_p)), z))
+        )
+
+        scale = -1 / pos_cardinal
+
+        return scale * loss

model.py

@@ -0,0 +1,305 @@
+import lightning.pytorch as pl
+from transformers import (
+    AdamW,
+    AutoModel,
+    AutoConfig,
+    get_linear_schedule_with_warmup,
+)
+from transformers.models.bert.modeling_bert import BertLMPredictionHead
+import torch
+from torch import nn
+from loss import PCCL
+import config
+
+
+class CL_model(pl.LightningModule):
+    def __init__(
+        self, n_batches=None, n_epochs=None, lr=None, mlm_weight=None, **kwargs
+    ):
+        super().__init__()
+
+        ## Params
+        self.n_batches = n_batches
+        self.n_epochs = n_epochs
+        self.lr = lr
+        self.mlm_weight = mlm_weight
+        self.config = AutoConfig.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+        ## Encoder
+        self.bert = AutoModel.from_pretrained(
+            "emilyalsentzer/Bio_ClinicalBERT", return_dict=True
+        )
+        # Unfreeze layers
+        self.bert_layer_num = sum(1 for _ in self.bert.named_parameters())
+        self.num_unfreeze_layer = self.bert_layer_num
+        self.ratio_unfreeze_layer = 0.0
+        if kwargs:
+            for key, value in kwargs.items():
+                if key == "unfreeze" and isinstance(value, float):
+                    assert (
+                        value >= 0.0 and value <= 1.0
+                    ), "ValueError: value must be a ratio between 0.0 and 1.0"
+                    self.ratio_unfreeze_layer = value
+        if self.ratio_unfreeze_layer > 0.0:
+            self.num_unfreeze_layer = int(
+                self.bert_layer_num * self.ratio_unfreeze_layer
+            )
+        for param in list(self.bert.parameters())[: -self.num_unfreeze_layer]:
+            param.requires_grad = False
+
+        self.lm_head = BertLMPredictionHead(self.config)
+        # self.projector = nn.Linear(self.bert.config.hidden_size, 128)
+        print("Model Initialized!")
+
+        ## Losses
+        self.cl_loss = PCCL()
+        self.mlm_loss = nn.CrossEntropyLoss()
+
+        ## Logs
+        self.num_batches = 0
+        self.train_loss, self.val_loss = 0, 0
+        self.train_loss_cl, self.val_loss_cl = 0, 0
+        self.train_loss_mlm, self.val_loss_mlm = 0, 0
+        self.training_step_outputs, self.validation_step_outputs = [], []
+
+    def forward(self, input_ids, attention_mask, masked_indices, eval=False):
+        embs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        cls_tokens = embs.pooler_output
+        mask_tokens = []
+        for idx, value in enumerate(masked_indices):
+            masks = embs.last_hidden_state[idx][value]
+            avg_mask = torch.mean(masks, dim=0)
+            mask_tokens.append(avg_mask)
+        mask_tokens = torch.stack(mask_tokens)
+        cls_concat_mask = torch.cat((cls_tokens, mask_tokens), dim=1)
+        if eval is True:
+            return cls_tokens, mask_tokens, cls_concat_mask
+
+        mlm_pred = self.lm_head(embs.last_hidden_state)
+        mlm_pred = mlm_pred.view(-1, self.config.vocab_size)
+        return cls_concat_mask, mlm_pred
+
+    def training_step(self, batch, batch_idx):
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        mlm_labels = batch["mlm_labels"]
+        masked_indices = batch["masked_indices"]
+        tags = batch["tags"]
+        scores = batch["scores"]
+        cls_concat_mask, mlm_pred = self(input_ids, attention_mask, masked_indices)
+        loss_cl = self.cl_loss(cls_concat_mask, tags, scores)
+        loss_mlm = self.mlm_loss(mlm_pred, mlm_labels.reshape(-1))
+        loss = (1 - self.mlm_weight) * loss_cl + self.mlm_weight * loss_mlm
+        logs = {"loss": loss, "loss_cl": loss_cl, "loss_mlm": loss_mlm}
+        self.training_step_outputs.append(logs)
+        self.log("train_loss", loss, prog_bar=True, logger=True, sync_dist=True)
+
+        self.num_batches += 1
+        self.train_loss_cl += loss_cl
+        self.train_loss_mlm += loss_mlm
+        self.train_loss += loss
+
+        if self.num_batches % config.log_every_n_steps == 0:
+            avg_loss_cl = self.train_loss_cl / self.num_batches
+            avg_loss_mlm = self.train_loss_mlm / self.num_batches
+            avg_loss = self.train_loss / self.num_batches
+            self.log(
+                "train_avg_cl_loss",
+                avg_loss_cl,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+            self.log(
+                "train_avg_mlm_loss",
+                avg_loss_mlm,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+            self.log(
+                "train_avg_loss", avg_loss, prog_bar=True, logger=True, sync_dist=True
+            )
+            self.train_loss_cl = 0
+            self.train_loss_mlm = 0
+            self.train_loss = 0
+            self.num_batches = 0
+
+        return loss
+
+    def on_train_epoch_end(self):
+        e_t_avg_loss = (
+            torch.stack([x["loss"] for x in self.training_step_outputs])
+            .mean()
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        self.log(
+            "avg_loss_train_epoch",
+            e_t_avg_loss.item(),
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        e_t_avg_loss_cl = (
+            torch.stack([x["loss_cl"] for x in self.training_step_outputs])
+            .mean()
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        self.log(
+            "avg_loss_cl_train_epoch",
+            e_t_avg_loss_cl.item(),
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        e_t_avg_loss_mlm = (
+            torch.stack([x["loss_mlm"] for x in self.training_step_outputs])
+            .mean()
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        self.log(
+            "avg_loss_mlm_train_epoch",
+            e_t_avg_loss_mlm.item(),
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        print(
+            "train_epoch:",
+            self.current_epoch,
+            "avg_loss:",
+            e_t_avg_loss,
+            "avg_cl_loss:",
+            e_t_avg_loss_cl,
+            "avg_mlm_loss:",
+            e_t_avg_loss_mlm,
+        )
+        self.training_step_outputs.clear()
+
+    def validation_step(self, batch, batch_idx):
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        mlm_labels = batch["mlm_labels"]
+        masked_indices = batch["masked_indices"]
+        tags = batch["tags"]
+        scores = batch["scores"]
+        cls_concat_mask, mlm_pred = self(input_ids, attention_mask, masked_indices)
+        loss_cl = self.cl_loss(cls_concat_mask, tags, scores)
+        loss_mlm = self.mlm_loss(mlm_pred, mlm_labels.reshape(-1))
+        loss = (1 - self.mlm_weight) * loss_cl + self.mlm_weight * loss_mlm
+        logs = {"loss": loss, "loss_cl": loss_cl, "loss_mlm": loss_mlm}
+        self.validation_step_outputs.append(logs)
+        self.log("val_loss", loss, prog_bar=True, logger=True, sync_dist=True)
+
+        self.num_batches += 1
+        self.val_loss_cl += loss_cl
+        self.val_loss_mlm += loss_mlm
+        self.val_loss += loss
+
+        if self.num_batches % config.log_every_n_steps == 0:
+            avg_loss_cl = self.val_loss_cl / self.num_batches
+            avg_loss_mlm = self.val_loss_mlm / self.num_batches
+            avg_loss = self.val_loss / self.num_batches
+            self.log(
+                "val_avg_cl_loss",
+                avg_loss_cl,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+            self.log(
+                "val_avg_mlm_loss",
+                avg_loss_mlm,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+            self.log(
+                "val_avg_loss",
+                avg_loss,
+                prog_bar=True,
+                logger=True,
+                sync_dist=True,
+            )
+            self.val_loss_cl = 0
+            self.val_loss_mlm = 0
+            self.val_loss = 0
+            self.num_batches = 0
+
+        return loss
+
+    def on_validation_epoch_end(self):
+        e_v_avg_loss = (
+            torch.stack([x["loss"] for x in self.validation_step_outputs])
+            .mean()
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        self.log(
+            "avg_loss_val_epoch",
+            e_v_avg_loss.item(),
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        e_v_avg_loss_cl = (
+            torch.stack([x["loss_cl"] for x in self.validation_step_outputs])
+            .mean()
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        self.log(
+            "avg_loss_cl_val_epoch",
+            e_v_avg_loss_cl.item(),
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        e_v_avg_loss_mlm = (
+            torch.stack([x["loss_mlm"] for x in self.validation_step_outputs])
+            .mean()
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        self.log(
+            "avg_loss_mlm_val_epoch",
+            e_v_avg_loss_mlm.item(),
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        print(
+            "val_epoch:",
+            self.current_epoch,
+            "avg_loss:",
+            e_v_avg_loss,
+            "avg_cl_loss:",
+            e_v_avg_loss_cl,
+            "avg_mlm_loss:",
+            e_v_avg_loss_mlm,
+        )
+        self.validation_step_outputs.clear()
+
+    def configure_optimizers(self):
+        # Optimizer
+        self.trainable_params = [
+            param for param in self.parameters() if param.requires_grad
+        ]
+        optimizer = AdamW(self.trainable_params, lr=self.lr)
+
+        # Scheduler
+        warmup_steps = self.n_batches // 3
+        total_steps = self.n_batches * self.n_epochs - warmup_steps
+        scheduler = get_linear_schedule_with_warmup(
+            optimizer, warmup_steps, total_steps
+        )
+        return [optimizer], [scheduler]

train.py

@@ -0,0 +1,81 @@
+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
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from transformers import AutoTokenizer
+from ast import literal_eval
+from pytorch_lightning.loggers import TensorBoardLogger
+
+# imports from our own modules
+import config
+from model import CL_model
+from dataset import CLDataModule
+
+if __name__ == "__main__":
+    seed_everything(0, workers=True)
+    logger = TensorBoardLogger(
+        "/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/prompt/logs", name="CL"
+    )
+
+    query_df = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/query_df.csv"
+    )
+    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]
+    )
+    train_df, val_df = train_test_split(query_df, test_size=config.split_ratio)
+
+    all_d = pd.read_csv(
+        "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/query_all_d.csv"
+    )
+    all_d.drop(columns=["finding_sites", "morphology"], inplace=True)
+    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"]))
+
+    pairs = pd.read_csv("/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/pairs.csv")
+
+    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+    data_module = CLDataModule(train_df, val_df, tokenizer, dictionary, all_d, pairs)
+    data_module.setup()
+
+    model = CL_model(
+        n_batches=len(data_module.train_dataset) / config.batch_size,
+        n_epochs=config.max_epochs,
+        lr=config.learning_rate,
+        mlm_weight=config.mlm_weight,
+        unfreeze=config.unfreeze_ratio,
+    )
+
+    checkpoint = ModelCheckpoint(
+        dirpath="/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/train/ckpt/v2",
+        filename="{epoch}-{step}",
+        save_weights_only=True,
+        save_last=True,
+        every_n_train_steps=config.ckcpt_every_n_steps,
+        monitor=None,
+        save_top_k=-1,
+    )
+
+    trainer = 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="val_loss", min_delta=1e-3, patience=2, mode="min"),
+            checkpoint,
+        ],
+        profiler="simple",
+        log_every_n_steps=config.log_every_n_steps,
+    )
+
+    trainer.fit(model, data_module)