Upload 3 files

dataset.py

@@ -1,21 +1,66 @@
-import pandas as pd
 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 sys
 
 sys.path.append("/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag")
 from data_proc.data_gen import (
     positive_generator,
+    positive_generator_alter,
     negative_generator,
+    negative_generator_alter,
+    negative_generator_random,
+    negative_generator_v2,
     get_mentioned_code,
 )
 
 
-##### General
-class ContrastiveLearningDataset(Dataset):
+def tokenize(text, tokenizer, tag):
+    inputs = tokenizer(
+        text,
+        return_token_type_ids=False,
+        return_tensors="pt",
+    )
+
+    inputs["input_ids"] = inputs["input_ids"][0]
+    inputs["attention_mask"] = inputs["attention_mask"][0]
+    inputs["mlm_ids"] = copy.deepcopy(inputs["input_ids"])
+    inputs["mlm_labels"] = copy.deepcopy(inputs["input_ids"])
+
+    tokens_to_ignore = torch.tensor([101, 102, 0])  # [CLS], [SEP], [PAD]
+    valid_tokens = inputs["input_ids"][
+        ~torch.isin(inputs["input_ids"], tokens_to_ignore)
+    ]
+    num_of_token_to_mask = int(len(valid_tokens) * config.mask_pct)
+    token_to_mask = valid_tokens[
+        torch.randperm(valid_tokens.size(0))[:num_of_token_to_mask]
+    ]
+    inputs["mlm_ids"] = [
+        103 if x in token_to_mask else x for x in inputs["mlm_ids"]
+    ]  # [MASK]
+    inputs["mlm_labels"] = [
+        y if y in token_to_mask else -100 for y in inputs["mlm_labels"]
+    ]
+
+    inputs["mlm_ids"] = torch.tensor(inputs["mlm_ids"])
+    inputs["mlm_labels"] = torch.tensor(inputs["mlm_labels"])
+    if tag == "A":
+        inputs["tag"] = 0
+    elif tag == "P":
+        inputs["tag"] = 1
+    elif tag == "N":
+        inputs["tag"] = 2
+    return inputs
+
+
+class CLDataset(Dataset):
     def __init__(
         self,
         data: pd.DataFrame,
@@ -31,728 +76,146 @@ class ContrastiveLearningDataset(Dataset):
         return sentence
 
 
-def max_pairwise_sim(sentence1, sentence2, current_df, query_df, sim_df, all_d):
-    """Returns the maximum ontology similarity score between concept pairs mentioned in sentence1 and sentence2.
-
-    Args:
-        sentence1: anchor sentence
-        sentence2: negative sentence
-        current_df: the dataset where anchor sentence stays
-        query_df: the union of training and validation sets
-        dictionary: cardiac-related {concepts: synonyms}
-        sim_df: the dataset of pairwise ontology similarity score
-        all_d: the dataset of [concepts, synonyms, list of ancestor concepts]
-    """
-    # retrieve concepts from the two sentences
-    anchor_codes = get_mentioned_code(sentence1, current_df)
-    other_codes = get_mentioned_code(sentence2, query_df)
-
-    # create snomed-ct code pairs and calculate the score using sim_df
-    code_pairs = list(zip(anchor_codes, other_codes))
-    sim_scores = []
-    for pair in code_pairs:
-        code1 = pair[0]
-        code2 = pair[1]
-        if code1 == code2:
-            result = len(all_d.loc[all_d["concept"] == code1, "ancestors"].values[0])
-            sim_scores.append(result)
-        else:
-            try:
-                result = sim_df.loc[
-                    (sim_df["Code1"] == code1) & (sim_df["Code2"] == code2), "score"
-                ].values[0]
-                sim_scores.append(result)
-            except:
-                result = sim_df.loc[
-                    (sim_df["Code1"] == code2) & (sim_df["Code2"] == code1), "score"
-                ].values[0]
-            sim_scores.append(result)
-    if len(sim_scores) > 0:
-        return max(sim_scores)
-    else:
-        return 0
-
-
-##### SimCSE
-def collate_simcse(batch, tokenizer):
-    """
-    Use the first sample in the batch as the anchor,
-    use the duplicate of anchor as the positive,
-    use the rest of the batch as negatives.
-    """
-    anchor = batch[0]  # use the first sample in the batch as anchor
-    positive = anchor[:]  # create a duplicate of anchor as positive
-    negatives = batch[1:]  # everything else as negatives
-    df = pd.DataFrame(columns=["label", "input_ids", "attention_mask"])
-
-    anchor_token = tokenizer.encode_plus(
-        anchor,
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        return_tensors="pt",
-    )
-    anchor_row = pd.DataFrame(
-        {
-            "label": 0,
-            "input_ids": anchor_token["input_ids"].tolist(),
-            "attention_mask": anchor_token["attention_mask"].tolist(),
-        }
-    )
-    df = pd.concat([df, anchor_row])
-
-    pos_token = tokenizer.encode_plus(
-        positive,
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        return_tensors="pt",
-    )
-    pos_row = pd.DataFrame(
-        {
-            "label": 1,
-            "input_ids": pos_token["input_ids"].tolist(),
-            "attention_mask": pos_token["attention_mask"].tolist(),
-        }
-    )
-    df = pd.concat([df, pos_row])
-
-    for neg in negatives:
-        neg_token = tokenizer.encode_plus(
-            neg,
-            return_token_type_ids=False,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        neg_row = pd.DataFrame(
-            {
-                "label": 2,
-                "input_ids": neg_token["input_ids"].tolist(),
-                "attention_mask": neg_token["attention_mask"].tolist(),
-            }
-        )
-        df = pd.concat([df, neg_row])
-
-    label = torch.tensor(df["label"].tolist())
-
-    input_ids_tsr = list(map(lambda x: torch.tensor(x), df["input_ids"]))
-    padded_input_ids = pad_sequence(input_ids_tsr, padding_value=tokenizer.pad_token_id)
-    padded_input_ids = torch.transpose(padded_input_ids, 0, 1)
-
-    attention_mask_tsr = list(map(lambda x: torch.tensor(x), df["attention_mask"]))
-    padded_attention_mask = pad_sequence(attention_mask_tsr, padding_value=0)
-    padded_attention_mask = torch.transpose(padded_attention_mask, 0, 1)
-
-    return {
-        "label": label,
-        "input_ids": padded_input_ids,
-        "attention_mask": padded_attention_mask,
-    }
-
-
-def create_dataloader_simcse(
-    dataset,
-    tokenizer,
-    shuffle,
-):
-    return DataLoader(
-        dataset,
-        batch_size=config.batch_size_simcse,
-        shuffle=shuffle,
-        num_workers=config.num_workers,
-        collate_fn=lambda batch: collate_simcse(
-            batch,
-            tokenizer,
-        ),
-    )
-
-
-class ContrastiveLearningDataModule_simcse(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_df,
-        val_df,
-        tokenizer,
-    ):
-        super().__init__()
-        self.train_df = train_df
-        self.val_df = val_df
-        self.tokenizer = tokenizer
-
-    def setup(self, stage=None):
-        self.train_dataset = ContrastiveLearningDataset(self.train_df)
-        self.val_dataset = ContrastiveLearningDataset(self.val_df)
-
-    def train_dataloader(self):
-        return create_dataloader_simcse(
-            self.train_dataset,
-            self.tokenizer,
-            shuffle=True,
-        )
-
-    def val_dataloader(self):
-        return create_dataloader_simcse(
-            self.val_dataset,
-            self.tokenizer,
-            shuffle=False,
-        )
-
+def collate_func(batch, tokenizer, current_df, query_df, dictionary, all_d):
 
-##### SimCSE_w
-def collate_simcse_w(
-    batch,
-    current_df,
-    query_df,
-    tokenizer,
-    sim_df,
-    all_d,
-):
-    """
-    Anchor: 0
-    Positive: 1
-    Negative: 2
-    """
     anchor = batch[0]
-    positive = anchor[:]
-    negatives = batch[1:]
-    df = pd.DataFrame(columns=["label", "input_ids", "attention_mask", "score"])
-
-    anchor_token = tokenizer.encode_plus(
-        anchor,
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        return_tensors="pt",
-    )
-
-    anchor_row = pd.DataFrame(
-        {
-            "label": 0,
-            "input_ids": anchor_token["input_ids"].tolist(),
-            "attention_mask": anchor_token["attention_mask"].tolist(),
-            "score": 1,
-        }
-    )
-    df = pd.concat([df, anchor_row])
-
-    pos_token = tokenizer.encode_plus(
-        positive,
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        return_tensors="pt",
-    )
-    pos_row = pd.DataFrame(
-        {
-            "label": 1,
-            "input_ids": pos_token["input_ids"].tolist(),
-            "attention_mask": pos_token["attention_mask"].tolist(),
-            "score": 1,
-        }
-    )
-    df = pd.concat([df, pos_row])
-
-    for neg in negatives:
-        neg_token = tokenizer.encode_plus(
-            neg,
-            return_token_type_ids=False,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        score = max_pairwise_sim(anchor, neg, current_df, query_df, sim_df, all_d)
-        offset = 8
-        score = score + offset
-        neg_row = pd.DataFrame(
-            {
-                "label": 2,
-                "input_ids": neg_token["input_ids"].tolist(),
-                "attention_mask": neg_token["attention_mask"].tolist(),
-                "score": score,
-            }
-        )
-        df = pd.concat([df, neg_row])
-
-    label = torch.tensor(df["label"].tolist())
-
-    input_ids_tsr = list(map(lambda x: torch.tensor(x), df["input_ids"]))
-    padded_input_ids = pad_sequence(input_ids_tsr, padding_value=tokenizer.pad_token_id)
-    padded_input_ids = torch.transpose(padded_input_ids, 0, 1)
-
-    attention_mask_tsr = list(map(lambda x: torch.tensor(x), df["attention_mask"]))
-    padded_attention_mask = pad_sequence(attention_mask_tsr, padding_value=0)
-    padded_attention_mask = torch.transpose(padded_attention_mask, 0, 1)
-
-    score = torch.tensor(df["score"].tolist())
-
-    return {
-        "label": label,
-        "input_ids": padded_input_ids,
-        "attention_mask": padded_attention_mask,
-        "score": score,
-    }
-
-
-def create_dataloader_simcse_w(
-    dataset,
-    current_df,
-    query_df,
-    tokenizer,
-    sim_df,
-    all_d,
-    shuffle,
-):
-    return DataLoader(
-        dataset,
-        batch_size=config.batch_size_simcse,
-        shuffle=shuffle,
-        num_workers=config.num_workers,
-        collate_fn=lambda batch: collate_simcse_w(
-            batch,
-            current_df,
-            query_df,
-            tokenizer,
-            sim_df,
-            all_d,
-        ),
-    )
-
-
-class ContrastiveLearningDataModule_simcse_w(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_df,
-        val_df,
-        query_df,
-        tokenizer,
-        sim_df,
-        all_d,
-    ):
-        super().__init__()
-        self.train_df = train_df
-        self.val_df = val_df
-        self.query_df = query_df
-        self.tokenizer = tokenizer
-        self.sim_df = sim_df
-        self.all_d = all_d
-
-    def setup(self, stage=None):
-        self.train_dataset = ContrastiveLearningDataset(self.train_df)
-        self.val_dataset = ContrastiveLearningDataset(self.val_df)
-
-    def train_dataloader(self):
-        return create_dataloader_simcse_w(
-            self.train_dataset,
-            self.train_df,
-            self.query_df,
-            self.tokenizer,
-            self.sim_df,
-            self.all_d,
-            shuffle=True,
-        )
-
-    def val_dataloader(self):
-        return create_dataloader_simcse_w(
-            self.val_dataset,
-            self.val_df,
-            self.query_df,
-            self.tokenizer,
-            self.sim_df,
-            self.all_d,
-            shuffle=False,
-        )
-
-
-##### Samp
-def collate_samp(
-    sentence,
-    current_df,
-    query_df,
-    tokenizer,
-    dictionary,
-    sim_df,
-):
-
-    anchor = sentence[0]
-    positives = positive_generator(
-        anchor, current_df, query_df, dictionary, num_pos=config.num_pos
-    )
-    negatives = negative_generator(
+    positives = positive_generator_alter(
         anchor,
         current_df,
-        query_df,
-        dictionary,
-        sim_df,
-        num_neg=config.num_neg,
-    )
-    df = pd.DataFrame(columns=["label", "input_ids", "attention_mask"])
-    anchor_token = tokenizer.encode_plus(
-        anchor,
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        return_tensors="pt",
-    )
-
-    anchor_row = pd.DataFrame(
-        {
-            "label": 0,
-            "input_ids": anchor_token["input_ids"].tolist(),
-            "attention_mask": anchor_token["attention_mask"].tolist(),
-        }
-    )
-    df = pd.concat([df, anchor_row])
-
-    for pos in positives:
-        token = tokenizer.encode_plus(
-            pos,
-            return_token_type_ids=False,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        row = pd.DataFrame(
-            {
-                "label": 1,
-                "input_ids": token["input_ids"].tolist(),
-                "attention_mask": token["attention_mask"].tolist(),
-            }
-        )
-        df = pd.concat([df, row])
-
-    for neg in negatives:
-        token = tokenizer.encode_plus(
-            neg,
-            return_token_type_ids=False,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        row = pd.DataFrame(
-            {
-                "label": 2,
-                "input_ids": token["input_ids"].tolist(),
-                "attention_mask": token["attention_mask"].tolist(),
-            }
-        )
-        df = pd.concat([df, row])
-
-    label = torch.tensor(df["label"].tolist())
-
-    input_ids_tsr = list(map(lambda x: torch.tensor(x), df["input_ids"]))
-    padded_input_ids = pad_sequence(input_ids_tsr, padding_value=tokenizer.pad_token_id)
-    padded_input_ids = torch.transpose(padded_input_ids, 0, 1)
-
-    attention_mask_tsr = list(map(lambda x: torch.tensor(x), df["attention_mask"]))
-    padded_attention_mask = pad_sequence(attention_mask_tsr, padding_value=0)
-    padded_attention_mask = torch.transpose(padded_attention_mask, 0, 1)
-
-    return {
-        "label": label,
-        "input_ids": padded_input_ids,
-        "attention_mask": padded_attention_mask,
-    }
-
-
-def create_dataloader_samp(
-    dataset,
-    current_df,
-    query_df,
-    tokenizer,
-    dictionary,
-    sim_df,
-    shuffle,
-):
-    return DataLoader(
-        dataset,
-        batch_size=config.batch_size,
-        shuffle=shuffle,
-        num_workers=config.num_workers,
-        collate_fn=lambda batch: collate_samp(
-            batch,
-            current_df,
-            query_df,
-            tokenizer,
-            dictionary,
-            sim_df,
-        ),
-    )
-
-
-class ContrastiveLearningDataModule_samp(pl.LightningDataModule):
-    def __init__(
-        self,
-        train_df,
-        val_df,
-        query_df,
-        tokenizer,
         dictionary,
-        sim_df,
-    ):
-        super().__init__()
-        self.train_df = train_df
-        self.val_df = val_df
-        self.query_df = query_df
-        self.tokenizer = tokenizer
-        self.dictionary = dictionary
-        self.sim_df = sim_df
-
-    def setup(self, stage=None):
-        self.train_dataset = ContrastiveLearningDataset(self.train_df)
-        self.val_dataset = ContrastiveLearningDataset(self.val_df)
-
-    def train_dataloader(self):
-        return create_dataloader_samp(
-            self.train_dataset,
-            self.train_df,
-            self.query_df,
-            self.tokenizer,
-            self.dictionary,
-            self.sim_df,
-            shuffle=True,
-        )
-
-    def val_dataloader(self):
-        return create_dataloader_samp(
-            self.val_dataset,
-            self.val_df,
-            self.query_df,
-            self.tokenizer,
-            self.dictionary,
-            self.sim_df,
-            shuffle=False,
-        )
-
-
-##### Samp_w
-def collate_samp_w(
-    sentence,
-    current_df,
-    query_df,
-    tokenizer,
-    dictionary,
-    sim_df,
-    all_d,
-):
-    """
-    Anchor: 0
-    Positive: 1
-    Negative: 2
-    """
-    anchor = sentence[0]
-    positives = positive_generator(
-        anchor, current_df, query_df, dictionary, num_pos=config.num_pos
+        num_pos=config.num_pos,
     )
-    negatives = negative_generator(
+    negatives = negative_generator_v2(
         anchor,
         current_df,
         query_df,
-        dictionary,
-        sim_df,
+        all_d,
         num_neg=config.num_neg,
     )
-    df = pd.DataFrame(columns=["label", "input_ids", "attention_mask", "score"])
-    anchor_token = tokenizer.encode_plus(
-        anchor,
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        return_tensors="pt",
-    )
 
-    anchor_row = pd.DataFrame(
-        {
-            "label": 0,
-            "input_ids": anchor_token["input_ids"].tolist(),
-            "attention_mask": anchor_token["attention_mask"].tolist(),
-            "score": 1,
-        }
-    )
-    df = pd.concat([df, anchor_row])
+    inputs = []
+
+    anchor_dict = tokenize(anchor, tokenizer, "A")
+    inputs.append(anchor_dict)
 
     for pos in positives:
-        token = tokenizer.encode_plus(
-            pos,
-            return_token_type_ids=False,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        row = pd.DataFrame(
-            {
-                "label": 1,
-                "input_ids": token["input_ids"].tolist(),
-                "attention_mask": token["attention_mask"].tolist(),
-                "score": 1,
-            }
-        )
-        df = pd.concat([df, row])
+        pos_dict = tokenize(pos, tokenizer, "P")
+        inputs.append(pos_dict)
 
     for neg in negatives:
-        token = tokenizer.encode_plus(
-            neg,
-            return_token_type_ids=False,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        score = max_pairwise_sim(anchor, neg, current_df, query_df, sim_df, all_d)
-        offset = 8  # all negative scores start with 8 to distinguish from the positives
-        score = score + offset
-        row = pd.DataFrame(
-            {
-                "label": 2,
-                "input_ids": token["input_ids"].tolist(),
-                "attention_mask": token["attention_mask"].tolist(),
-                "score": score,
-            }
-        )
-        df = pd.concat([df, row])
+        neg_dict = tokenize(neg, tokenizer, "N")
+        inputs.append(neg_dict)
 
-    label = torch.tensor(df["label"].tolist())
+    tags = torch.tensor([d["tag"] for d in inputs])
 
-    input_ids_tsr = list(map(lambda x: torch.tensor(x), df["input_ids"]))
-    padded_input_ids = pad_sequence(input_ids_tsr, padding_value=tokenizer.pad_token_id)
+    input_ids_tsr = [d["input_ids"] for d in inputs]
+    padded_input_ids = pad_sequence(input_ids_tsr, padding_value=0)
     padded_input_ids = torch.transpose(padded_input_ids, 0, 1)
 
-    attention_mask_tsr = list(map(lambda x: torch.tensor(x), df["attention_mask"]))
+    attention_mask_tsr = [d["attention_mask"] for d in inputs]
     padded_attention_mask = pad_sequence(attention_mask_tsr, padding_value=0)
     padded_attention_mask = torch.transpose(padded_attention_mask, 0, 1)
 
-    score = torch.tensor(df["score"].tolist())
+    mlm_ids_tsr = [d["mlm_ids"] for d in inputs]
+    padded_mlm_ids = pad_sequence(mlm_ids_tsr, padding_value=0)
+    padded_mlm_ids = torch.transpose(padded_mlm_ids, 0, 1)
+
+    mlm_labels_tsr = [d["mlm_labels"] for d in inputs]
+    padded_mlm_labels = pad_sequence(mlm_labels_tsr, padding_value=-100)
+    padded_mlm_labels = torch.transpose(padded_mlm_labels, 0, 1)
 
     return {
-        "label": label,
+        "tags": tags,
         "input_ids": padded_input_ids,
         "attention_mask": padded_attention_mask,
-        "score": score,
+        "mlm_ids": padded_mlm_ids,
+        "mlm_labels": padded_mlm_labels,
     }
 
 
-def create_dataloader_samp_w(
-    dataset,
-    current_df,
-    query_df,
-    tokenizer,
-    dictionary,
-    sim_df,
-    all_d,
-    shuffle,
+def create_dataloader(
+    dataset, tokenizer, shuffle, current_df, query_df, dictionary, all_d
 ):
     return DataLoader(
         dataset,
         batch_size=config.batch_size,
         shuffle=shuffle,
-        num_workers=config.num_workers,
-        collate_fn=lambda batch: collate_samp_w(
-            batch,
-            current_df,
-            query_df,
-            tokenizer,
-            dictionary,
-            sim_df,
-            all_d,
+        num_workers=1,
+        collate_fn=lambda batch: collate_func(
+            batch, tokenizer, current_df, query_df, dictionary, all_d
         ),
     )
 
 
-class ContrastiveLearningDataModule_samp_w(pl.LightningDataModule):
+class CLDataModule(pl.LightningDataModule):
     def __init__(
         self,
         train_df,
         val_df,
-        query_df,
         tokenizer,
+        query_df,
         dictionary,
-        sim_df,
         all_d,
     ):
         super().__init__()
         self.train_df = train_df
         self.val_df = val_df
-        self.query_df = query_df
         self.tokenizer = tokenizer
+        self.query_df = query_df
         self.dictionary = dictionary
-        self.sim_df = sim_df
         self.all_d = all_d
 
     def setup(self, stage=None):
-        self.train_dataset = ContrastiveLearningDataset(self.train_df)
-        self.val_dataset = ContrastiveLearningDataset(self.val_df)
+        self.train_dataset = CLDataset(self.train_df)
+        self.val_dataset = CLDataset(self.val_df)
 
     def train_dataloader(self):
-        return create_dataloader_samp_w(
+        return create_dataloader(
             self.train_dataset,
-            self.train_df,
-            self.query_df,
             self.tokenizer,
-            self.dictionary,
-            self.sim_df,
-            self.all_d,
             shuffle=True,
+            current_df=self.train_df,
+            query_df=self.query_df,
+            dictionary=self.dictionary,
+            all_d=self.all_d,
         )
 
     def val_dataloader(self):
-        return create_dataloader_samp_w(
+        return create_dataloader(
             self.val_dataset,
-            self.val_df,
-            self.query_df,
             self.tokenizer,
-            self.dictionary,
-            self.sim_df,
-            self.all_d,
             shuffle=False,
+            current_df=self.val_df,
+            query_df=self.query_df,
+            dictionary=self.dictionary,
+            all_d=self.all_d,
         )
 
 
-#### Test
-from transformers import AutoTokenizer
-from ast import literal_eval
-from sklearn.model_selection import train_test_split
+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)
+    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 
-query_df = pd.read_csv(
-    "/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag/data_proc/mimic_data/processed_train/processed.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]
-)  # 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/query_all_d.csv"
+    )
+    all_d["synonyms"] = all_d["synonyms"].apply(literal_eval)
+    all_d["ancestors"] = all_d["ancestors"].apply(literal_eval)
+    all_d["finding_sites"] = all_d["finding_sites"].apply(literal_eval)
+    all_d["morphology"] = all_d["morphology"].apply(literal_eval)
+    dictionary = dict(zip(all_d["concept"], all_d["synonyms"]))
 
-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"]))
-
-d1 = ContrastiveLearningDataModule_simcse(train_df, val_df, tokenizer)
-d1.setup()
-train_d1 = d1.train_dataloader()
-for batch in train_d1:
-    b1 = batch
-    break
-
-d2 = ContrastiveLearningDataModule_simcse_w(
-    train_df, val_df, query_df, tokenizer, sim_df, all_d
-)
-d2.setup()
-train_d2 = d2.train_dataloader()
-for batch in train_d2:
-    b2 = batch
-    break
-
-d3 = ContrastiveLearningDataModule_samp(
-    train_df, val_df, query_df, tokenizer, dictionary, sim_df
-)
-d3.setup()
-train_d3 = d3.train_dataloader()
-for batch in train_d3:
-    b3 = batch
-    break
-
-d4 = ContrastiveLearningDataModule_samp_w(
-    train_df, val_df, query_df, tokenizer, dictionary, sim_df, all_d
-)
-d4.setup()
-train_d4 = d4.train_dataloader()
-for batch in train_d4:
-    b4 = batch
-    break
+    d = CLDataModule(train_df, val_df, tokenizer, query_df, dictionary, all_d)
+    d.setup()
+    train = d.train_dataloader()
+    for batch in train:
+        b = batch
+        break

loss.py

@@ -4,124 +4,17 @@ import torch.nn.functional as F
 import config
 
 
-class ContrastiveLoss_simcse(nn.Module):
-    """SimCSE loss"""
-
-    def __init__(self):
-        super(ContrastiveLoss_simcse, self).__init__()
-        self.temperature = config.temperature
-
-    def forward(self, feature_vectors, labels):
-        normalized_features = F.normalize(
-            feature_vectors, p=2, dim=0
-        )  # normalize along columns
-
-        # Identify indices for each label
-        anchor_indices = (labels == 0).nonzero().squeeze(dim=1)
-        positive_indices = (labels == 1).nonzero().squeeze(dim=1)
-        negative_indices = (labels == 2).nonzero().squeeze(dim=1)
-
-        # Extract tensors based on labels
-        anchor = normalized_features[anchor_indices]
-        positives = normalized_features[positive_indices]
-        negatives = normalized_features[negative_indices]
-        pos_and_neg = torch.cat([positives, negatives])
-
-        denominator = torch.sum(
-            torch.exp(
-                torch.div(
-                    torch.matmul(anchor, torch.transpose(pos_and_neg, 0, 1)),
-                    self.temperature,
-                )
-            )
-        )
-
-        numerator = torch.exp(
-            torch.div(
-                torch.matmul(anchor, torch.transpose(positives, 0, 1)),
-                self.temperature,
-            )
-        )
-
-        loss = -torch.log(
-            torch.div(
-                numerator,
-                denominator,
-            )
-        )
-
-        return loss
-
-
-class ContrastiveLoss_simcse_w(nn.Module):
-    """SimCSE loss with weighting."""
-
-    def __init__(self):
-        super(ContrastiveLoss_simcse_w, self).__init__()
-        self.temperature = config.temperature
-
-    def forward(self, feature_vectors, labels, scores):
-        normalized_features = F.normalize(
-            feature_vectors, p=2, dim=0
-        )  # normalize along columns
-
-        # Identify indices for each label
-        anchor_indices = (labels == 0).nonzero().squeeze(dim=1)
-        positive_indices = (labels == 1).nonzero().squeeze(dim=1)
-        negative_indices = (labels == 2).nonzero().squeeze(dim=1)
-
-        pos_scores = scores[positive_indices].float()
-        normalized_neg_scores = F.normalize(
-            scores[negative_indices].float(), p=2, dim=0
-        )  # l2-norm
-        normalized_neg_scores += 1
-        scores = torch.cat([pos_scores, normalized_neg_scores])
-
-        # Extract tensors based on labels
-        anchor = normalized_features[anchor_indices]
-        positives = normalized_features[positive_indices]
-        negatives = normalized_features[negative_indices]
-        pos_and_neg = torch.cat([positives, negatives])
-
-        denominator = torch.sum(
-            torch.exp(
-                scores
-                * torch.div(
-                    torch.matmul(anchor, torch.transpose(pos_and_neg, 0, 1)),
-                    self.temperature,
-                )
-            )
-        )
-
-        numerator = torch.exp(
-            torch.div(
-                torch.matmul(anchor, torch.transpose(positives, 0, 1)),
-                self.temperature,
-            )
-        )
-
-        loss = -torch.log(
-            torch.div(
-                numerator,
-                denominator,
-            )
-        )
-
-        return loss
-
-
-class ContrastiveLoss_samp(nn.Module):
+class CL_loss(nn.Module):
     """Supervised contrastive loss without weighting."""
 
     def __init__(self):
-        super(ContrastiveLoss_samp, self).__init__()
+        super(CL_loss, self).__init__()
         self.temperature = config.temperature
 
     def forward(self, feature_vectors, labels):
-        # Normalize feature vectors
         normalized_features = F.normalize(
-            feature_vectors, p=2, dim=0
-        )  # normalize along columns
+            feature_vectors, p=2, dim=1
+        )  # normalize by row, each row euc is approximately 1
 
         # Identify indices for each label
         anchor_indices = (labels == 0).nonzero().squeeze(dim=1)
@@ -139,82 +32,35 @@ class ContrastiveLoss_samp(nn.Module):
         denominator = torch.sum(
             torch.exp(
                 torch.div(
-                    torch.matmul(anchor, torch.transpose(pos_and_neg, 0, 1)),
+                    F.cosine_similarity(anchor, pos_and_neg, dim=1),
                     self.temperature,
                 )
             )
         )
 
-        sum_log_ent = torch.sum(
-            torch.log(
-                torch.div(
-                    torch.exp(
-                        torch.div(
-                            torch.matmul(anchor, torch.transpose(positives, 0, 1)),
-                            self.temperature,
-                        )
-                    ),
-                    denominator,
-                )
-            )
-        )
-
-        scale = -1 / pos_cardinal
-
-        return scale * sum_log_ent
-
-
-class ContrastiveLoss_samp_w(nn.Module):
-    """Supervised contrastive loss with weighting."""
-
-    def __init__(self):
-        super(ContrastiveLoss_samp_w, self).__init__()
-        self.temperature = config.temperature
-
-    def forward(self, feature_vectors, labels, scores):
-        # Normalize feature vectors
-        normalized_features = F.normalize(
-            feature_vectors, p=2, dim=0
-        )  # normalize along columns
-
-        # Identify indices for each label
-        anchor_indices = (labels == 0).nonzero().squeeze(dim=1)
-        positive_indices = (labels == 1).nonzero().squeeze(dim=1)
-        negative_indices = (labels == 2).nonzero().squeeze(dim=1)
-
-        # Normalize score vector
-        num_skip = len(positive_indices) + 1
-        pos_scores = scores[: (num_skip - 1)].float()  # exclude anchor
-        normalized_neg_scores = F.normalize(
-            scores[num_skip:].float(), p=2, dim=0
-        )  # l2-norm
-        normalized_neg_scores += 1
-        scores = torch.cat([pos_scores, normalized_neg_scores])
-
-        # Extract tensors based on labels
-        anchor = normalized_features[anchor_indices]
-        positives = normalized_features[positive_indices]
-        negatives = normalized_features[negative_indices]
-        pos_and_neg = torch.cat([positives, negatives])
-
-        pos_cardinal = positives.shape[0]
-
-        denominator = torch.sum(
-            torch.exp(
-                scores
-                * torch.div(
-                    torch.matmul(anchor, torch.transpose(pos_and_neg, 0, 1)),
-                    self.temperature,
-                )
-            )
-        )
+        # if not torch.isfinite(denominator):
+        #     print("Denominator is Inf!")
+
+        # if not torch.isfinite(
+        #     torch.exp(
+        #         torch.div(F.cosine_similarity(anchor, pos_and_neg, dim=1)),
+        #         self.temperature,
+        #     )
+        # ).all():
+        #     print("Exp is Inf!")
+        #     print(
+        #         torch.exp(
+        #             torch.div(F.cosine_similarity(anchor, pos_and_neg, dim=1)),
+        #             self.temperature,
+        #         )
+        #     )
 
         sum_log_ent = torch.sum(
             torch.log(
                 torch.div(
                     torch.exp(
                         torch.div(
-                            torch.matmul(anchor, torch.transpose(positives, 0, 1)),
+                            F.cosine_similarity(anchor, positives, dim=1),
                             self.temperature,
                         )
                     ),
@@ -224,5 +70,6 @@ class ContrastiveLoss_samp_w(nn.Module):
         )
 
         scale = -1 / pos_cardinal
+        out = scale * sum_log_ent
 
-        return scale * sum_log_ent
+        return out

model.py

@@ -2,31 +2,35 @@ 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 (
-    ContrastiveLoss_simcse,
-    ContrastiveLoss_simcse_w,
-    ContrastiveLoss_samp,
-    ContrastiveLoss_samp_w,
-)
+from loss import CL_loss
+import pandas as pd
 
 
-class BERTContrastiveLearning_simcse(pl.LightningModule):
-    def __init__(self, n_batches=None, n_epochs=None, lr=None, **kwargs):
+class CL_model(pl.LightningModule):
+    def __init__(
+        self, n_batches=None, n_epochs=None, lr=None, mlm_weight=None, **kwargs
+    ):
         super().__init__()
-        ### Parameters
+
+        ## Params
         self.n_batches = n_batches
         self.n_epochs = n_epochs
         self.lr = lr
+        self.mlm_weight = mlm_weight
+        # self.first_neg_idx = 0
+        self.config = AutoConfig.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 
-        ### Architecture
+        ## Encoder
         self.bert = AutoModel.from_pretrained(
             "emilyalsentzer/Bio_ClinicalBERT", return_dict=True
         )
-        # Unfreeze encoder
+        # 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
@@ -43,378 +47,138 @@ class BERTContrastiveLearning_simcse(pl.LightningModule):
             )
         for param in list(self.bert.parameters())[: -self.num_unfreeze_layer]:
             param.requires_grad = False
-        # Random dropouts
-        self.dropout1 = nn.Dropout(p=0.1)
-        self.dropout2 = nn.Dropout(p=0.1)
-        # Linear projector
+
+        self.lm_head = BertLMPredictionHead(self.config)
         self.projector = nn.Linear(self.bert.config.hidden_size, 128)
         print("Model Initialized!")
 
-        ### Loss
-        self.criterion = ContrastiveLoss_simcse()
-
-        ### Logs
-        self.train_loss, self.val_loss, self.test_loss = [], [], []
-        self.training_step_outputs = []
-        self.validation_step_outputs = []
-
-    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]
-
-    def forward(self, input_ids, attention_mask):
-        emb = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls = emb.pooler_output
-        out = self.projector(cls)
-        anchor_out = self.dropout1(out[0:1])
-        rest_out = self.dropout2(out[1:])
-        output = torch.cat([anchor_out, rest_out])
-        return cls, output
+        ## Losses
+        self.cl_loss = CL_loss()
+        self.mlm_loss = nn.CrossEntropyLoss()
+
+        ## Logs
+        self.train_loss, self.val_loss = [], []
+        self.train_cl_loss, self.val_cl_loss = [], []
+        self.train_mlm_loss, self.val_mlm_loss = [], []
+        self.training_step_outputs, self.validation_step_outputs = [], []
+
+    def forward(self, input_ids, attention_mask, mlm_ids, eval=False):
+        # Contrastive
+        unmasked = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        cls = unmasked.pooler_output
+        if eval is True:
+            return cls
+        output = self.projector(cls)
+
+        # MLM
+        masked = self.bert(input_ids=mlm_ids, attention_mask=attention_mask)
+        pred = self.lm_head(masked.last_hidden_state)
+        pred = pred.view(-1, self.config.vocab_size)
+        return cls, output, pred
 
     def training_step(self, batch, batch_idx):
-        label = batch["label"]
+        tags = batch["tags"]
         input_ids = batch["input_ids"]
         attention_mask = batch["attention_mask"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label)
-        logs = {"loss": loss}
+        mlm_ids = batch["mlm_ids"]
+        mlm_labels = batch["mlm_labels"].reshape(-1)
+        cls, output, pred = self(input_ids, attention_mask, mlm_ids)
+        loss_cl = self.cl_loss(output, tags)
+        loss_mlm = self.mlm_loss(pred, mlm_labels)
+        loss = 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)
         return loss
 
     def on_train_epoch_end(self):
-        loss = (
+        avg_loss = (
             torch.stack([x["loss"] for x in self.training_step_outputs])
             .mean()
             .detach()
             .cpu()
             .numpy()
         )
-        self.train_loss.append(loss)
-        print("train_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.training_step_outputs.clear()
-
-    def validation_step(self, batch, batch_idx):
-        label = batch["label"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label)
-        logs = {"loss": loss}
-        self.validation_step_outputs.append(logs)
-        self.log("validation_loss", loss, prog_bar=True, logger=True, sync_dist=True)
-        return loss
-
-    def on_validation_epoch_end(self):
-        loss = (
-            torch.stack([x["loss"] for x in self.validation_step_outputs])
+        self.train_loss.append(avg_loss)
+        avg_cl_loss = (
+            torch.stack([x["loss_cl"] for x in self.training_step_outputs])
             .mean()
             .detach()
             .cpu()
             .numpy()
         )
-        self.val_loss.append(loss)
-        print("val_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.validation_step_outputs.clear()
-
-
-class BERTContrastiveLearning_simcse_w(pl.LightningModule):
-    def __init__(self, n_batches=None, n_epochs=None, lr=None, **kwargs):
-        super().__init__()
-        ### Parameters
-        self.n_batches = n_batches
-        self.n_epochs = n_epochs
-        self.lr = lr
-
-        ### Architecture
-        self.bert = AutoModel.from_pretrained(
-            "emilyalsentzer/Bio_ClinicalBERT", return_dict=True
-        )
-        # Unfreeze encoder
-        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
-        # Random dropouts
-        self.dropout1 = nn.Dropout(p=0.1)
-        self.dropout2 = nn.Dropout(p=0.1)
-        # Linear projector
-        self.projector = nn.Linear(self.bert.config.hidden_size, 128)
-        print("Model Initialized!")
-
-        ### Loss
-        self.criterion = ContrastiveLoss_simcse_w()
-
-        ### Logs
-        self.train_loss, self.val_loss, self.test_loss = [], [], []
-        self.training_step_outputs = []
-        self.validation_step_outputs = []
-
-    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]
-
-    def forward(self, input_ids, attention_mask):
-        emb = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls = emb.pooler_output
-        out = self.projector(cls)
-        anchor_out = self.dropout1(out[0:1])
-        rest_out = self.dropout2(out[1:])
-        output = torch.cat([anchor_out, rest_out])
-        return cls, output
-
-    def training_step(self, batch, batch_idx):
-        label = batch["label"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        score = batch["score"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label, score)
-        logs = {"loss": loss}
-        self.training_step_outputs.append(logs)
-        self.log("train_loss", loss, prog_bar=True, logger=True, sync_dist=True)
-        return loss
-
-    def on_train_epoch_end(self):
-        loss = (
-            torch.stack([x["loss"] for x in self.training_step_outputs])
+        self.train_cl_loss.append(avg_cl_loss)
+        avg_mlm_loss = (
+            torch.stack([x["loss_mlm"] for x in self.training_step_outputs])
             .mean()
             .detach()
             .cpu()
             .numpy()
         )
-        self.train_loss.append(loss)
-        print("train_epoch:", self.current_epoch, "avg_loss:", loss)
+        self.train_mlm_loss.append(avg_mlm_loss)
+        print(
+            "train_epoch:",
+            self.current_epoch,
+            "avg_loss:",
+            avg_loss,
+            "avg_cl_loss:",
+            avg_cl_loss,
+            "avg_mlm_loss:",
+            avg_mlm_loss,
+        )
         self.training_step_outputs.clear()
 
     def validation_step(self, batch, batch_idx):
-        label = batch["label"]
+        tags = batch["tags"]
         input_ids = batch["input_ids"]
         attention_mask = batch["attention_mask"]
-        score = batch["score"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label, score)
-        logs = {"loss": loss}
+        mlm_ids = batch["mlm_ids"]
+        mlm_labels = batch["mlm_labels"].reshape(-1)
+        cls, output, pred = self(input_ids, attention_mask, mlm_ids)
+        loss_cl = self.cl_loss(output, tags)
+        loss_mlm = self.mlm_loss(pred, mlm_labels)
+        loss = 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("validation_loss", loss, prog_bar=True, logger=True, sync_dist=True)
         return loss
 
     def on_validation_epoch_end(self):
-        loss = (
+        avg_loss = (
             torch.stack([x["loss"] for x in self.validation_step_outputs])
             .mean()
             .detach()
             .cpu()
             .numpy()
         )
-        self.val_loss.append(loss)
-        print("val_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.validation_step_outputs.clear()
-
-
-class BERTContrastiveLearning_samp(pl.LightningModule):
-
-    def __init__(self, n_batches=None, n_epochs=None, lr=None, **kwargs):
-        super().__init__()
-        ### Parameters
-        self.n_batches = n_batches
-        self.n_epochs = n_epochs
-        self.lr = lr
-
-        ### Architecture
-        self.bert = AutoModel.from_pretrained(
-            "emilyalsentzer/Bio_ClinicalBERT", return_dict=True
-        )
-        # Unfreeze encoder
-        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
-        # Linear projector
-        self.projector = nn.Linear(self.bert.config.hidden_size, 128)
-        print("Model Initialized!")
-
-        ### Loss
-        self.criterion = ContrastiveLoss_samp()
-
-        ### Logs
-        self.train_loss, self.val_loss, self.test_loss = [], [], []
-        self.training_step_outputs = []
-        self.validation_step_outputs = []
-
-    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]
-
-    def forward(self, input_ids, attention_mask):
-        emb = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls = emb.pooler_output
-        out = self.projector(cls)
-        return cls, out
-
-    def training_step(self, batch, batch_idx):
-        label = batch["label"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label)
-        logs = {"loss": loss}
-        self.training_step_outputs.append(logs)
-        self.log("train_loss", loss, prog_bar=True, logger=True, sync_dist=True)
-        return loss
-
-    def on_train_epoch_end(self):
-        loss = (
-            torch.stack([x["loss"] for x in self.training_step_outputs])
+        self.val_loss.append(avg_loss)
+        avg_cl_loss = (
+            torch.stack([x["loss_cl"] for x in self.validation_step_outputs])
             .mean()
             .detach()
             .cpu()
             .numpy()
         )
-        self.train_loss.append(loss)
-        print("train_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.training_step_outputs.clear()
-
-    def validation_step(self, batch, batch_idx):
-        label = batch["label"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label)
-        logs = {"loss": loss}
-        self.validation_step_outputs.append(logs)
-        self.log("validation_loss", loss, prog_bar=True, logger=True, sync_dist=True)
-        return loss
-
-    def on_validation_epoch_end(self):
-        loss = (
-            torch.stack([x["loss"] for x in self.validation_step_outputs])
+        self.val_cl_loss.append(avg_cl_loss)
+        avg_mlm_loss = (
+            torch.stack([x["loss_mlm"] for x in self.validation_step_outputs])
             .mean()
             .detach()
             .cpu()
             .numpy()
         )
-        self.val_loss.append(loss)
-        print("val_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.validation_step_outputs.clear()
-
-
-class BERTContrastiveLearning_samp_w(pl.LightningModule):
-
-    def __init__(self, n_batches=None, n_epochs=None, lr=None, **kwargs):
-        super().__init__()
-        ### Parameters
-        self.n_batches = n_batches
-        self.n_epochs = n_epochs
-        self.lr = lr
-
-        ### Architecture
-        self.bert = AutoModel.from_pretrained(
-            "emilyalsentzer/Bio_ClinicalBERT", return_dict=True
+        self.val_mlm_loss.append(avg_mlm_loss)
+        print(
+            "val_epoch:",
+            self.current_epoch,
+            "avg_loss:",
+            avg_loss,
+            "avg_cl_loss:",
+            avg_cl_loss,
+            "avg_mlm_loss:",
+            avg_mlm_loss,
         )
-        # Unfreeze encoder
-        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
-        # Linear projector
-        self.projector = nn.Linear(self.bert.config.hidden_size, 128)
-        print("Model Initialized!")
-
-        ### Loss
-        self.criterion = ContrastiveLoss_samp_w()
-
-        ### Logs
-        self.train_loss, self.val_loss, self.test_loss = [], [], []
-        self.training_step_outputs = []
-        self.validation_step_outputs = []
+        self.validation_step_outputs.clear()
 
     def configure_optimizers(self):
         # Optimizer
@@ -424,69 +188,9 @@ class BERTContrastiveLearning_samp_w(pl.LightningModule):
         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]
-
-    def forward(self, input_ids, attention_mask):
-        emb = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls = emb.pooler_output
-        out = self.projector(cls)
-        return cls, out
-
-    def training_step(self, batch, batch_idx):
-        label = batch["label"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        score = batch["score"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
+        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
         )
-        loss = self.criterion(out, label, score)
-        logs = {"loss": loss}
-        self.training_step_outputs.append(logs)
-        self.log("train_loss", loss, prog_bar=True, logger=True, sync_dist=True)
-        return loss
-
-    def on_train_epoch_end(self):
-        loss = (
-            torch.stack([x["loss"] for x in self.training_step_outputs])
-            .mean()
-            .detach()
-            .cpu()
-            .numpy()
-        )
-        self.train_loss.append(loss)
-        print("train_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.training_step_outputs.clear()
-
-    def validation_step(self, batch, batch_idx):
-        label = batch["label"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        score = batch["score"]
-        cls, out = self(
-            input_ids,
-            attention_mask,
-        )
-        loss = self.criterion(out, label, score)
-        logs = {"loss": loss}
-        self.validation_step_outputs.append(logs)
-        self.log("validation_loss", loss, prog_bar=True, logger=True, sync_dist=True)
-        return loss
-
-    def on_validation_epoch_end(self):
-        loss = (
-            torch.stack([x["loss"] for x in self.validation_step_outputs])
-            .mean()
-            .detach()
-            .cpu()
-            .numpy()
-        )
-        self.val_loss.append(loss)
-        print("val_epoch:", self.current_epoch, "avg_loss:", loss)
-        self.validation_step_outputs.clear()
+        return [optimizer], [scheduler]