Create dataset.py

This file holds 4 dataset modules for the 4 models respectively: SimCSE, SimCSE_w, Samp, Samp_w.
Run the test at the end to see what's in each training batch.

dataset.py

@@ -0,0 +1,758 @@
+import pandas as pd
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.utils.rnn import pad_sequence
+import lightning.pytorch as pl
+import config
+import sys
+
+sys.path.append("/home/sunx/data/aiiih/projects/sunx/ccf_fuzzy_diag")
+from data_proc.data_gen import (
+    positive_generator,
+    negative_generator,
+    get_mentioned_code,
+)
+
+
+##### General
+class ContrastiveLearningDataset(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
+        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,
+        )
+
+
+##### 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(
+        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
+    )
+    negatives = negative_generator(
+        anchor,
+        current_df,
+        query_df,
+        dictionary,
+        sim_df,
+        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])
+
+    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])
+
+    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])
+
+    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_samp_w(
+    dataset,
+    current_df,
+    query_df,
+    tokenizer,
+    dictionary,
+    sim_df,
+    all_d,
+    shuffle,
+):
+    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,
+        ),
+    )
+
+
+class ContrastiveLearningDataModule_samp_w(pl.LightningDataModule):
+    def __init__(
+        self,
+        train_df,
+        val_df,
+        query_df,
+        tokenizer,
+        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.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)
+
+    def train_dataloader(self):
+        return create_dataloader_samp_w(
+            self.train_dataset,
+            self.train_df,
+            self.query_df,
+            self.tokenizer,
+            self.dictionary,
+            self.sim_df,
+            self.all_d,
+            shuffle=True,
+        )
+
+    def val_dataloader(self):
+        return create_dataloader_samp_w(
+            self.val_dataset,
+            self.val_df,
+            self.query_df,
+            self.tokenizer,
+            self.dictionary,
+            self.sim_df,
+            self.all_d,
+            shuffle=False,
+        )
+
+
+#### Test
+from transformers import AutoTokenizer
+from ast import literal_eval
+from sklearn.model_selection import train_test_split
+
+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/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