<jupyter_start><jupyter_text>Processing Data, Defining Data Classes and Collator and other miscellaneous stuff<jupyter_code># Class to store data samples, text_a, text_b, score
@dataclass
class Example:
    text_a: str
    text_b: str
    score: float


# lowercase and add space around words, remove unnecessary spaces
def pre_process(sentence, cased=False):
    sentence = sentence.replace(":", " : ").replace("/", " / ").replace("[", " [ ").replace("]", " ] ").replace("(", " ( ").replace(")", " ) ").replace("\"", " \" ").replace("-", " - ").replace("?", " ").lstrip().rstrip()
    if cased:
      return re.sub(' +',' ', sentence)
    return re.sub(' +',' ', sentence).lower()


# returns test and train arrays as Example Objects
# test train split is stratified and 80-20 split
def get_data(cased=False):
    train_data = "/content/drive/My Drive/clinical-sts/train.tsv"
    df = pd.read_csv(train_data, sep="\t", names=["sentence_1", "sentence_2", "similarity_score"], encoding="utf-8")
    df["sentence_1"] = df["sentence_1"].apply(lambda sentence: pre_process(sentence, cased))
    df["sentence_2"] = df["sentence_2"].apply(lambda sentence: pre_process(sentence, cased))
    df["input_sample"] = df["sentence_1"] + "<SEP>" + df["sentence_2"]

    ## stratified binned sampling
    min_val = np.amin(df["similarity_score"])
    max_val = np.amax(df["similarity_score"])
    bins     = np.linspace(start=min_val, stop=max_val, num=10)
    y_binned = np.digitize(df["similarity_score"], bins, right=True)

    X_train, X_test, y_train, y_test = train_test_split(
        df["input_sample"], 
        df["similarity_score"], 
        stratify=y_binned,
        test_size=0.2,
        random_state=23
    )

    train_a_b = [Example(text_a=sample.split("<SEP>")[0], 
                    text_b=sample.split("<SEP>")[1], 
                    score=similarity_score) for sample, similarity_score in zip(X_train, y_train)]
    train_b_a = [Example(text_a=sample.split("<SEP>")[1], 
                    text_b=sample.split("<SEP>")[0], 
                    score=similarity_score) for sample, similarity_score in zip(X_train, y_train)]
    train = train_a_b + train_b_a

    test_a_b = [Example(text_a=sample.split("<SEP>")[0], 
                    text_b=sample.split("<SEP>")[1], 
                    score=similarity_score) for sample, similarity_score in zip(X_test, y_test)]
    test_b_a = [Example(text_a=sample.split("<SEP>")[1], 
                    text_b=sample.split("<SEP>")[0], 
                    score=similarity_score) for sample, similarity_score in zip(X_test, y_test)]
    test = test_a_b + test_b_a

    return train, test




# DYNAMIC PADDING AND UNIFORM LENGTH BATCHING - reduces wasted computation and makes it faster to run
# CODE BORROWED FROM https://towardsdatascience.com/divide-hugging-face-transformers-training-time-by-2-or-more-21bf7129db9q-21bf7129db9e


# We'll be creating a custome dataset using this first
@dataclass
class Features:
    og_sample: Example
    input_ids: List[int]
    attention_mask: List[int]
    score: float


class TextDataset(Dataset):
    def __init__(self, tokenizer, 
                 pad_to_max_length, 
                 max_len,
                 examples: List[Example]):
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.examples: List[Example] = examples
        self.current = 0
        self.pad_to_max_length = pad_to_max_length

    # tokenize the sentences and return a Features object for each sentence 
    def encode(self, ex: Example) -> Features:
        encode_dict = self.tokenizer.encode_plus(text=ex.text_a,
                                                 text_pair=ex.text_b,
                                                 add_special_tokens=True,
                                                 max_length=self.max_len,
                                                 pad_to_max_length=self.pad_to_max_length,
                                                 return_token_type_ids=False,
                                                 return_attention_mask=True,
                                                 return_overflowing_tokens=False,
                                                 return_special_tokens_mask=False,
                                                 truncation=True,
                                                 )
        return Features(og_sample=ex,
                        input_ids=encode_dict["input_ids"],
                        attention_mask=encode_dict["attention_mask"],
                        score=ex.score)

    def __getitem__(self, idx) -> Features:
        return self.encode(ex=self.examples[idx])

    def __len__(self):
        return len(self.examples)


def pad_seq(seq: List[int], max_batch_len: int, pad_value: int) -> List[int]:
    return seq + (max_batch_len - len(seq)) * [pad_value]


# Smart Collator is used to create batches which are dynamically padded with uniform length 
@dataclass
class SmartCollator:  # (DataCollator):
    pad_token_id: int

    def collate_batch(self, batch: List[Features]) -> Dict[str, torch.Tensor]:
        batch_og_sample = list()
        batch_inputs = list()
        batch_attention_masks = list()
        labels = list()
        max_size = max([len(ex.input_ids) for ex in batch])
        for item in batch:
            batch_inputs += [pad_seq(item.input_ids, max_size, self.pad_token_id)]
            batch_attention_masks += [pad_seq(item.attention_mask, max_size, 0)]
            labels.append(item.score)
            batch_og_sample.append(item)

        return {"input_ids": torch.tensor(batch_inputs, dtype=torch.long),
                "attention_mask": torch.tensor(batch_attention_masks, dtype=torch.long),
                "score": torch.tensor(labels, dtype=torch.float),
                "og_sample": batch_og_sample
                }
                
def collate_wrapper(data):
    collator = SmartCollator(pad_token_id=tokenizer.pad_token_id)
    return collator.collate_batch(data)


# USE THIS FUNCTION TO LOAD TEST AND TRAIN DATA AND ITERATE THROUGH THEM
def load_data(tokenizer, batch_size, cased=False):
    # Get train and test Data Examples
    train, test = get_data(cased)


    # Now tokenize the words and convert them to token IDs
    max_sequence_len = 128
    train_set = TextDataset(tokenizer=tokenizer,
                            max_len=max_sequence_len,
                            examples=train,
                            pad_to_max_length=False)

    test_set = TextDataset(tokenizer=tokenizer,
                            max_len=max_sequence_len,
                            examples=test,
                            pad_to_max_length=False)

    train_dataloader = DataLoader(train_set, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=collate_wrapper,
              pin_memory=False, drop_last=False, timeout=0,
              worker_init_fn=None)

    test_dataloader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=collate_wrapper,
              pin_memory=False, drop_last=False, timeout=0,
              worker_init_fn=None)
    
    return train_dataloader, test_dataloader<jupyter_output><empty_output><jupyter_text>**Define the Training Loop for fine-tuning the Model.**  We also have some miscellaneous functions to evaluate our model on the dev-set.        The model is as shown below. We have different learning rates for the bert and LR layer. Note that we take the hidden layer output from BERT and not the CLS embedding. The CLS embedding does not generate any meaningful sentence embedding and BERT was specifically trained for the NSP task. As such, using the CLS embedding directly leads to worse results. We found that using the penultimate hidden layers gave us best results for this task.<jupyter_code># Let's define the training loop and model

def get_bert_output(my_bert, input_ids, attention_mask):    
    outputs = my_bert(input_ids, attention_mask=attention_mask)
    hidden_states = outputs[2]
    sent_embedding = hidden_states[11][:, 0:1, :].squeeze(1).cuda()
    return sent_embedding


class linearRegression(nn.Module):
    def __init__(self):
        super(linearRegression, self).__init__()
        self.linear = nn.Linear(768, 1)
    def forward(self, x):
        out = self.linear(x)
        return out

    
def run_new_method(my_bert, optimizer, regression_head, regression_optimizer, train_dataloader, test_dataloader, epochs=10, freeze_layers=False):
    old_test_loss = float('inf')

    # freeze_layers = "0,1,2,3,4"
    if freeze_layers:
        freeze_layers = "0,1,2,3,4,5"
        layer_indexes = [int(x) for x in freeze_layers.split(",")]
        for layer_idx in layer_indexes:
            for param in list(my_bert.encoder.layer[layer_idx].parameters()):
                param.requires_grad = False
            print ("Froze Layer: ", layer_idx)


    for epoch_num in range(epochs):
        total_loss = 0
        batch_count = 0
        # put model in train mode
        my_bert.train()
        regression_head.train()
        for step_num, batch_data in enumerate(train_dataloader):
            input_ids = batch_data["input_ids"].to(device)
            attention_mask = batch_data["attention_mask"].to(device)
            score = batch_data["score"].to(device)
            score = score.unsqueeze(1)

            my_bert_optimizer.zero_grad()
            regression_optimizer.zero_grad()
            bert_embedding = get_bert_output(my_bert, input_ids, attention_mask)
            predicted_score = regression_head(bert_embedding)

            loss_func = nn.MSELoss()
            batch_loss = loss_func(predicted_score, score)

            
            batch_loss.backward()
            my_bert_optimizer.step()
            regression_optimizer.step()
            total_loss += batch_loss.item()
            batch_count += 1
            train_loss = total_loss/batch_count
        print("Epoch: {} Train Loss:{}".format(epoch_num, train_loss))

        # put model in test mode
        my_bert.eval()
        regression_head.eval()
        test_loss = 0
        test_batch_count = 0
        with torch.no_grad():
            for step_num, batch_data in enumerate(test_dataloader):
                input_ids = batch_data["input_ids"].to(device)
                attention_mask = batch_data["attention_mask"].to(device)
                score = batch_data["score"].to(device)
                score = score.unsqueeze(1)

                bert_embedding = get_bert_output(my_bert, input_ids, attention_mask)
                predicted_score = regression_head(bert_embedding)

                loss_func = nn.MSELoss()
                batch_loss = loss_func(predicted_score, score)
                test_loss += batch_loss.item()
                test_batch_count += 1
        curr_test_loss = test_loss/test_batch_count
        print("Epoch: {} Test Loss:{}\n".format(epoch_num, curr_test_loss))
        if curr_test_loss < 0.61 or train_loss < 0.25:
            print("yay, exit")
            break
        # curr_test_loss = test_loss/test_batch_count
        # if curr_test_loss-old_test_loss >= 0.03 or curr_test_loss<0.61:
        #     print("new test loss is greater; breaking")
        #     break
        # old_test_loss = curr_test_loss
    return my_bert, regression_head




'''
    Method to evaluate and calculate pcc on dev/test dataset, and show terrible predictions
'''
def evaluate_model(model, regression_head, test_dataloader, show_bad_predictions=True, prediction_difference=2.0):

    actual = list()
    predicted = list()
    og_data = list()

    model.eval()
    with torch.no_grad():
        test_batch_count = 0
        test_loss =0
        for step_num, batch_data in enumerate(test_dataloader):
            input_ids = batch_data["input_ids"].to(device)
            attention_mask = batch_data["attention_mask"].to(device)
            score = batch_data["score"].to(device)
            score = score.unsqueeze(1)

            bert_embedding = get_bert_output(model, input_ids, attention_mask)
            predicted_score = regression_head(bert_embedding)

            actual.extend(score.tolist())
            predicted.extend(predicted_score.tolist())
            original_samples = batch_data["og_sample"]
            og_data.extend(original_samples)

    # show bad predictions in a table
    table = Texttable()
    table.add_row(["Actual", "Predicted", "Difference", "Text Sample"])

    for act, pre, og_data in zip(actual, predicted, og_data):
        if abs(pre[0]-act[0]) > prediction_difference:
            og = og_data.og_sample.text_a + "    |||    " + og_data.og_sample.text_b
            print("{:.2f}    {:.2f}          {:.2f}     {}".format(act[0], pre[0], abs(pre[0]-act[0]), og))

    print('\n\n\n')
    actual = [item[0] for item in actual]
    predicted = [item[0] for item in predicted]

    correlation, p_value = pearsonr(actual, predicted)
    print(correlation)

    d = {"a": actual, "p": predicted}
    dx = pd.DataFrame(d)
    dx.plot.hist(bins=20, alpha=0.25)
    return correlation


def get_optimizer_params(model):
    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'gamma', 'beta']
    opt_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
        'weight_decay_rate': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
        'weight_decay_rate': 0.0}
    ]
    return opt_parameters<jupyter_output><empty_output><jupyter_text>**Select the base model to fine tune, and pass it to the train loop**We then evaluate the PCC on the dev set.  Note: Only bad examples and scores are printed.<jupyter_code>MODEL = "bert-base-uncased"
my_bert = BertModel.from_pretrained(MODEL, output_hidden_states=True)
my_bert.cuda()

regression_head = linearRegression().cuda()
linear_regression_optimizer = torch.optim.Adam(regression_head.parameters(), lr=1e-4)


tokenizer = BertTokenizer.from_pretrained(MODEL, output_hidden_states=True)
train_dataloader, test_dataloader = load_data(tokenizer=tokenizer, batch_size=8)
bert_config = BertConfig.from_pretrained(MODEL, output_hidden_states=True)
my_bert_optimizer = AdamW(get_optimizer_params(my_bert), lr=1e-5)
updated_model, regression_head = run_new_method(my_bert=my_bert, optimizer=my_bert_optimizer, # optimizer
                                                regression_head=regression_head, 
                                                regression_optimizer=linear_regression_optimizer,
                                                train_dataloader=train_dataloader,
                                                test_dataloader=test_dataloader,
                                                epochs=3)

pcc = evaluate_model(updated_model, regression_head, test_dataloader)
model_loc = "/content/drive/My Drive/clinical-sts/models/bert-{}-{:.2f}.pth".format(MODEL, pcc)
regression_loc = "/content/drive/My Drive/clinical-sts/models/regression-{}-{:.2f}.pth".format(MODEL, pcc)
torch.save(updated_model, model_loc)
torch.save(regression_head, regression_loc)
print("Saved at:\n'{}'\n'{}'".format(model_loc, regression_loc))


MODEL = "/content/drive/My Drive/clinical-sts/bio-bert/"
my_bert = BertModel.from_pretrained(MODEL, output_hidden_states=True)
my_bert.cuda()

regression_head = linearRegression().cuda()
linear_regression_optimizer = torch.optim.Adam(regression_head.parameters(), lr=1e-4)


tokenizer = AutoTokenizer.from_pretrained(MODEL, output_hidden_states=True)
train_dataloader, test_dataloader = load_data(tokenizer=tokenizer, batch_size=8)
bert_config = BertConfig.from_pretrained(MODEL, output_hidden_states=True)
my_bert_optimizer = AdamW(get_optimizer_params(my_bert), lr=1e-5)
updated_model, regression_head = run_new_method(my_bert=my_bert, optimizer=my_bert_optimizer, # optimizer
                                                regression_head=regression_head, 
                                                regression_optimizer=linear_regression_optimizer,
                                                train_dataloader=train_dataloader,
                                                test_dataloader=test_dataloader,
                                                epochs=3)

pcc = evaluate_model(updated_model, regression_head, test_dataloader)
model_name = "bio-bert"
model_loc = "/content/drive/My Drive/clinical-sts/models/bert-{}-{:.2f}.pth".format(model_name, pcc)
regression_loc = "/content/drive/My Drive/clinical-sts/models/regression-{}-{:.2f}.pth".format(model_name, pcc)
torch.save(updated_model, model_loc)
torch.save(regression_head, regression_loc)
print("Saved at:\n'{}'\n'{}'".format(model_loc, regression_loc))
MODEL = "//content/drive/My Drive/clinical-sts/NCBI-BERT/"
my_bert = BertModel.from_pretrained(MODEL, output_hidden_states=True)
my_bert.cuda()

regression_head = linearRegression().cuda()
linear_regression_optimizer = torch.optim.Adam(regression_head.parameters(), lr=1e-4)


tokenizer = BertTokenizer.from_pretrained(MODEL, output_hidden_states=True)
train_dataloader, test_dataloader = load_data(tokenizer=tokenizer, batch_size=8)
bert_config = BertConfig.from_pretrained(MODEL, output_hidden_states=True)
my_bert_optimizer = AdamW(get_optimizer_params(my_bert), lr=1e-5)
updated_model, regression_head = run_new_method(my_bert=my_bert, optimizer=my_bert_optimizer, # optimizer
                                                regression_head=regression_head, 
                                                regression_optimizer=linear_regression_optimizer,
                                                train_dataloader=train_dataloader,
                                                test_dataloader=test_dataloader,
                                                epochs=3)

pcc = evaluate_model(updated_model, regression_head, test_dataloader)
model_name = "blue-bert"
model_loc = "/content/drive/My Drive/clinical-sts/models/bert-{}-{:.3f}.pth".format(model_name, pcc)
regression_loc = "/content/drive/My Drive/clinical-sts/models/regression-{}-{:.3f}.pth".format(model_name, pcc)
torch.save(updated_model, model_loc)
torch.save(regression_head, regression_loc)
print("Saved at:\n'{}'\n'{}'".format(model_loc, regression_loc))<jupyter_output>Epoch: 0 Train Loss:1.1125290060016324
Epoch: 0 Test Loss:0.6453494093863361

Epoch: 1 Train Loss:0.39784646396640566
Epoch: 1 Test Loss:0.6885681696325899

Epoch: 2 Train Loss:0.2653962019943489
Epoch: 2 Test Loss:0.758003762148949

2.00    4.12          2.12     discussed diagnosis and treatment plan; patient expressed understanding of the content.    |||    they verbalized agreement and understanding of the plan and contact information.
1.00    4.27          3.27     the patient is awake, alert, and oriented times three.    |||    patient is appreciative, understands and is in agreement with plan of care.
4.00    1.90          2.10     per protocol continue current dose until next inr in 5 - 7 days.    |||    then resume current dose until next inr check, .
3.00    0.47          2.53     after explaining the procedure to the patient / caregiver, ear lavage of both ears was attempted .    |||    the impression and treatment plan were explained in detail to the patient / family who ex[...]<jupyter_text><jupyter_code>MODEL = "emilyalsentzer/Bio_ClinicalBERT"
my_bert = BertModel.from_pretrained(MODEL, output_hidden_states=True)
my_bert.cuda()

regression_head = linearRegression().cuda()
linear_regression_optimizer = torch.optim.Adam(regression_head.parameters(), lr=1e-4)


tokenizer = BertTokenizer.from_pretrained(MODEL, output_hidden_states=True)
train_dataloader, test_dataloader = load_data(tokenizer=tokenizer, batch_size=8, cased=True)
bert_config = BertConfig.from_pretrained(MODEL, output_hidden_states=True)
my_bert_optimizer = AdamW(get_optimizer_params(my_bert), lr=1e-5)
updated_model, regression_head = run_new_method(my_bert=my_bert, optimizer=my_bert_optimizer, # optimizer
                                                regression_head=regression_head, 
                                                regression_optimizer=linear_regression_optimizer,
                                                train_dataloader=train_dataloader,
                                                test_dataloader=test_dataloader,
                                                epochs=3)

pcc = evaluate_model(updated_model, regression_head, test_dataloader)
model_name = MODEL.replace("/", "-")
model_loc = "/content/drive/My Drive/clinical-sts/models/bert-{}-{:.2f}.pth".format(model_name, pcc)
regression_loc = "/content/drive/My Drive/clinical-sts/models/regression-{}-{:.2f}.pth".format(model_name, pcc)
torch.save(updated_model, model_loc)
torch.save(regression_head, regression_loc)
print("Saved at:\n'{}'\n'{}'".format(model_loc, regression_loc))


MODEL = "sentence-transformers/bert-base-nli-stsb-mean-tokens"
my_bert = BertModel.from_pretrained(MODEL, output_hidden_states=True)
my_bert.cuda()

regression_head = linearRegression().cuda()
linear_regression_optimizer = torch.optim.Adam(regression_head.parameters(), lr=1e-4)


tokenizer = BertTokenizer.from_pretrained(MODEL, output_hidden_states=True)
train_dataloader, test_dataloader = load_data(tokenizer=tokenizer, batch_size=8)
bert_config = BertConfig.from_pretrained(MODEL, output_hidden_states=True)
my_bert_optimizer = AdamW(get_optimizer_params(my_bert), lr=1e-5)
updated_model, regression_head = run_new_method(my_bert=my_bert, optimizer=my_bert_optimizer, # optimizer
                                                regression_head=regression_head, 
                                                regression_optimizer=linear_regression_optimizer,
                                                train_dataloader=train_dataloader,
                                                test_dataloader=test_dataloader,
                                                epochs=3)

pcc = evaluate_model(updated_model, regression_head, test_dataloader)
model_name = MODEL.replace("/", "-")
model_loc = "/content/drive/My Drive/clinical-sts/models/bert-{}-{:.2f}.pth".format(model_name, pcc)
regression_loc = "/content/drive/My Drive/clinical-sts/models/regression-{}-{:.2f}.pth".format(model_name, pcc)
torch.save(updated_model, model_loc)
torch.save(regression_head, regression_loc)
print("Saved at:\n'{}'\n'{}'".format(model_loc, regression_loc))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10320886554808847, 'fraction_numerical': 0.013618763629890071, 'mean_word_length': 2.7531317855353263, 'pattern_counts': {'":': 7, '<': 28, '<?xml version=': 0, '>': 31, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 70, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '96684', 'license': 'MIT', 'n_tokens_mistral': 7121, 'n_tokens_neox': 6764, 'n_words': 1329, 'path': 'clinical_sts_train.ipynb', 'repo_name': 'sreenivasanramesh/clinical-sts', 'size': '22469'}