| 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 CL_loss |
| import pandas as pd |
|
|
|
|
| class CL_model(pl.LightningModule): |
| def __init__( |
| self, n_batches=None, n_epochs=None, lr=None, mlm_weight=None, **kwargs |
| ): |
| super().__init__() |
|
|
| |
| self.n_batches = n_batches |
| self.n_epochs = n_epochs |
| self.lr = lr |
| self.mlm_weight = mlm_weight |
| |
| self.config = AutoConfig.from_pretrained("emilyalsentzer/Bio_ClinicalBERT") |
|
|
| |
| self.bert = AutoModel.from_pretrained( |
| "emilyalsentzer/Bio_ClinicalBERT", return_dict=True |
| ) |
| |
| self.bert_layer_num = sum(1 for _ in self.bert.named_parameters()) |
| self.num_unfreeze_layer = self.bert_layer_num |
| self.ratio_unfreeze_layer = 0.0 |
| if kwargs: |
| for key, value in kwargs.items(): |
| if key == "unfreeze" and isinstance(value, float): |
| assert ( |
| value >= 0.0 and value <= 1.0 |
| ), "ValueError: value must be a ratio between 0.0 and 1.0" |
| self.ratio_unfreeze_layer = value |
| if self.ratio_unfreeze_layer > 0.0: |
| self.num_unfreeze_layer = int( |
| self.bert_layer_num * self.ratio_unfreeze_layer |
| ) |
| for param in list(self.bert.parameters())[: -self.num_unfreeze_layer]: |
| param.requires_grad = False |
|
|
| self.lm_head = BertLMPredictionHead(self.config) |
| self.projector = nn.Linear(self.bert.config.hidden_size, 128) |
| print("Model Initialized!") |
|
|
| |
| self.cl_loss = CL_loss() |
| self.mlm_loss = nn.CrossEntropyLoss() |
|
|
| |
| 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): |
| |
| 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) |
|
|
| |
| 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): |
| tags = batch["tags"] |
| input_ids = batch["input_ids"] |
| attention_mask = batch["attention_mask"] |
| 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): |
| avg_loss = ( |
| torch.stack([x["loss"] for x in self.training_step_outputs]) |
| .mean() |
| .detach() |
| .cpu() |
| .numpy() |
| ) |
| 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.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_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): |
| tags = batch["tags"] |
| input_ids = batch["input_ids"] |
| attention_mask = batch["attention_mask"] |
| 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): |
| avg_loss = ( |
| torch.stack([x["loss"] for x in self.validation_step_outputs]) |
| .mean() |
| .detach() |
| .cpu() |
| .numpy() |
| ) |
| 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.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_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, |
| ) |
| self.validation_step_outputs.clear() |
|
|
| def configure_optimizers(self): |
| |
| self.trainable_params = [ |
| param for param in self.parameters() if param.requires_grad |
| ] |
| optimizer = AdamW(self.trainable_params, lr=self.lr) |
|
|
| |
| warmup_steps = self.n_batches // 3 |
| total_steps = self.n_batches * self.n_epochs - warmup_steps |
| scheduler = get_linear_schedule_with_warmup( |
| optimizer, warmup_steps, total_steps |
| ) |
| return [optimizer], [scheduler] |
|
|
